import pandas
import csv
import os

prompt_injections=['clean'] # 'direct_prompt_injection',"observation_prompt_injection",'mixed_attack',"memory_attack","clean",'DPI_MP','OPI_MP','DPI_OPI'
dirs = ['dynamic_prompt_rewriting'] # new_memory, no_memory, direct_paraphrase_defense, instructional_prevention, delimiters_defense, ob_sandwich_defense, dynamic_prompt_rewriting
agg_result = []
non_agg_result = []

for prompt_injection in prompt_injections:
    for dir in dirs:
        # model_list = ["gpt-4o-mini",'gpt-3.5-turbo']
        model_list = ['gpt-3.5-turbo',"gpt-4o-mini","llama3:70b","llama3.1:70b","llama3:8b","llama3.1:8b","gemma2:27b","gemma2:9b","mixtral:8x7b","qwen2:7b","qwen2:72b",'gpt-4o-2024-08-06','claude-3-5-sonnet-20240620']

        attack_methods = ["combined attack", "context ignoring", "escape characters", "fake completion", "naive"]

        def get_result(file_path):

            with open(file_path, "r", encoding='utf-8', errors='ignore') as f:
                content = f.readlines()

            result = []
            for line in content:
                line_end_num = line.split(' ')
                if len(line_end_num) == 0:
                    continue

                try:
                    line_end_num = float(line_end_num[-1])
                except:
                    continue

                if "Total number of attacks: " in line:
                    result.append(int(line_end_num))
                elif "Number of successful attacks: " in line:
                    result.append(int(line_end_num))
                elif "Number of workflow generation failures: " in line:
                    result.append(int(line_end_num))
                elif "Attack success rate: " in line:
                    result.append("%.3f"%round(line_end_num,3))
                elif "Attack success rate excluding workflow failure cases: " in line:
                    result.append("%.3f"%round(line_end_num,3))
                elif "Number of original task success: " in line:
                    result.append(int(line_end_num))
                elif "Original task success rate: " in line:
                    result.append("%.3f"%round(line_end_num,3))
                elif "Refuse judge success number: " in line:
                    result.append(int(line_end_num))
                elif "Refuse judge success rate: " in line:
                    result.append("%.3f"%round(line_end_num,3))
                elif "Memory search success number: " in line and dir in ['memory_enhanced','with_memory', 'new_memory']:
                    result.append(int(line_end_num))
                elif "Memory search success rate: " in line and dir in ['memory_enhanced','with_memory', 'new_memory']:
                    result.append("%.3f"%round(line_end_num,3))

            return result

        if dir in ['memory_enhanced','with_memory','new_memory']:
            result_csv = pandas.DataFrame(columns = ["LLM", "Attack", "Aggressive", "Attacks num", "Successful attack num", "Plan generation failures num", "ASR", "ASR - plan failure", "Original task success num", "Original task success rate", "Refuse number", "Refuse rate", "Memory search success num", "Memory search success rate"])
        # elif dir in ['no_memory', 'direct_delimiters','direct_sandwich_defense','direct_paraphrase_defense']:
        else:
            result_csv = pandas.DataFrame(columns = ["LLM", "Attack", "Aggressive", "Attacks num", "Successful attack num", "Plan generation failures num", "ASR", "ASR - plan failure", "Original task success num", "Original task success rate", "Refuse number", "Refuse rate"])

        for model in model_list:
            path = f"./logs/{prompt_injection}/{model}/{dir}"
            agg_result = []
            all_result = []
            non_agg_result = []
            for root, dirs, files in os.walk(path, topdown=False):
                files = sorted(files)
                files_list = []
                # Process only files with 'full_tools.log' suffix
                for name in files:
                    files_list.append(name)

                for name in files_list:
                    attack_method = name.split('-')[0].replace('_', ' ')
                    if "-non-" in name:
                        agg = 'No'
                    elif "all" in name:
                        agg = 'all'
                    else:
                        agg = 'Yes'

                    result = [model, attack_method, agg]
                    result.extend(get_result(path + '/' + name))
                    # Ensure the correct length of result
                    # print(len(result),dir)
                    if len(result) != 14 and dir in ['memory_enhanced','with_memory','new_memory']:
                        print(len(result),dir)
                        print(path + '/' + name)
                        continue
                    elif len(result) != 12 and ('memory' not in dir or dir =='no_memory'):
                        print(path + '/' + name)
                        print(len(result),dir)
                        continue

                    if agg == "No":
                        non_agg_result.append(result)
                    elif agg == "all":
                        all_result.append(result)
                    else:
                        agg_result.append(result)

                    #result_csv.loc[len(result_csv.index)] = result
                    #print(path + '/' + name)
                    #print(result)
            for result in non_agg_result:
                result_csv.loc[len(result_csv.index)] = result
            for result in agg_result:
                result_csv.loc[len(result_csv.index)] = result
            for result in all_result:
                result_csv.loc[len(result_csv.index)] = result
        print(f'Saving to: ./result_csv/result-{prompt_injection}+{dir}-1325.csv')
        result_csv.to_csv(f"./result_csv/result-{prompt_injection}+{dir}-1325.csv", index = False)
