import os
import json
import importlib
import sys
sys.path.append("../")
sys.path.append("../../")
from tqdm import tqdm 
from src.utils import get_tool_dict, get_simulated_attacker_tool_response, get_score, Logger
from src.params import parse_arguments
from src.prompts.agent_prompts import PROMPT_DICT
from src.output_parsing import evaluate_output_prompted
from src.models import MODELS
from src.attack_defense_tools import  *


def predict_one_case(system_prompt, user_prompt, params, item, model_class, tool_dict, step_2=False):
    available_tool_names = [item['User Tool']] + item['Attacker Tools']
    available_tool_descriptions = json.dumps([tool_dict[tool] for tool in available_tool_names], indent=True)
    
    input = item['User Instruction']
    
    if step_2:
        agent_scratchpad = f"\n{item['Thought']}\nAction: {item['User Tool']}\nAction Input: {item['Tool Parameters']}\nObservation: {item['Tool Response']}\nThought: {item['Attacker Tool Thought']}\nAction: {item['Attacker Tools'][0]}\nAction Input: {item['Attacker Tool Parameters']}\nObservation: {item['Attacker Tool Response']}\n"
    else:
        agent_scratchpad = f"\n{item['Thought']}\nAction: {item['User Tool']}\nAction Input: {item['Tool Parameters']}\nObservation: {item['Tool Response']}\n"

    user_prompt_filled = user_prompt.format(available_tool_descriptions = available_tool_descriptions, input = input, available_tool_names = available_tool_names, agent_scratchpad = agent_scratchpad)
    model_input = model_class.prepare_input(system_prompt, user_prompt_filled)
    output = model_class.call_model(model_input)
    return output

def main(params):
    logger = Logger(params['log_file'])
    logger.log(params)
    model_class = MODELS[params['model_type']](params)
    system_prompt, user_prompt = PROMPT_DICT[params['prompt_type']]
    output_dir = f"../results/prompted_{params['model_type']}_{params['model_name']}_{params['prompt_type']}"
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    # data_dir = '../data'
    
    tool_dict = get_tool_dict()
    
    # output_files = {}
    for attack in params['attacks']:
        for defense in params['defenses']:
            # TODO change the output file name
            file_name = f"test_cases_atk_{attack}_def_{defense}.json"
            # test_case_file = os.path.join(data_dir, file_name)
            test_case_file = params['test_case_file']
            output_file = os.path.join(output_dir, file_name)
            # output_files[attack] = output_file

            with open(test_case_file, 'r') as f:
                data = json.load(f)

            if params['use_cache']:
                existing_data_step1 = {}
                existing_data_step2 = {}
                if os.path.exists(output_file):
                    with open(output_file, 'r') as f:
                        for line in f.readlines():
                            try:
                                tmp = json.loads(line)
                                #TODO conduct attack and defense
                                if type(tmp['output']) == str and len(tmp['output']):
                                    existing_data_step1[(tmp['User Instruction'], tmp['Tool Response'])] = tmp['output']
                                if type(tmp['output Step 2']) == str and len(tmp['output Step 2']):
                                    existing_data_step2[(tmp['User Instruction'], tmp['Tool Response'], tmp['output'])] = tmp['output Step 2']
                            except:
                                continue
                    print(f"Got {len(existing_data_step1)} existing outputs for the first step and {len(existing_data_step2)} existing outputs for the second step")

            if not params['only_get_score']:
                with open(output_file, 'w') as f:
                    for item in tqdm(data):
                        try:
                            #TODO conduct attack and defense
                            item = apply_attack_and_defense(item, eval(attack), eval(defense))
                            if params['use_cache'] and (item['User Instruction'], item['Tool Response']) in existing_data_step1:
                                # print("Use existing output for step 1")
                                output = existing_data_step1[(item['User Instruction'], item['Tool Response'])]
                            else:
                                output = predict_one_case(system_prompt, user_prompt, params, item, model_class, tool_dict)
                            item['output'] = output

                            parsed_output = evaluate_output_prompted(item['output'], item['Attacker Tools'][0], item['Attacker Tools'], item['User Tool'], item["Attacker Instruction"])
                            item['eval'] = parsed_output['eval']
                            if item['eval'] == 'invalid':
                                item['invalid_reason'] = parsed_output['invalid_reason']

                            # if not params['only_first_step'] and item['eval'] == 'succ' and attack == 'ds':
                            #     item["Attacker Tool Parameters"] = parsed_output["Attacker Tool Parameters"]
                            #     item["Attacker Tool Thought"] = parsed_output["Attacker Tool Thought"]
                            #     item["Attacker Tool Response"] = get_simulated_attacker_tool_response(item['Attacker Tools'][0], item["Attacker Tool Parameters"], tool_dict)
                            #
                            #     if params['use_cache'] and (item['User Instruction'], item['Tool Response'], item['output']) in existing_data_step2:
                            #         output = existing_data_step2[(item['User Instruction'], item['Tool Response'], item['output'])]
                            #         # print("Use existing output for step 2")
                            #     else:
                            #         output = predict_one_case(system_prompt, user_prompt, params, item, model_class, tool_dict, step_2=True)
                            #
                            #     item['output Step 2'] = output
                            #     parsed_output = evaluate_output_prompted(output, item['Attacker Tools'][1], item['Attacker Tools'], item['User Tool'], item["Attacker Instruction"])
                            #     item['eval Step 2'] = parsed_output['eval']
                            #     if item['eval Step 2'] == 'invalid':
                            #         item['invalid_reason  Step 2'] = parsed_output['invalid_reason']

                            f.write(json.dumps(item)+'\n')
                        except Exception as e:
                            # This block will execute if an exception occurs
                            print(f"An error occurred: {e} in {output_file}")
                            # assert 1==0

            scores = get_score(output_file)
            logger.log(f"attack: {attack}, defense: {defense}")
            logger.log(json.dumps(scores, indent=True))

if __name__ == "__main__":
    params = parse_arguments(agent_type="prompted")

    main(params)
