import pandas
import csv
import os

prompt_injections = ["pot_backdoor","pot_clean"] # "clean", "pot_backdoor","pot_clean"
memory_attacks = ['no_memory'] # pot_paraphrase_defense, pot_shuffling_defense, no_memory
agg_result = []
non_agg_result = []

for prompt_injection in prompt_injections:
    for memory_attack in memory_attacks:
        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']
        # model_list = ["llama3:8b"]
        attack_methods = ["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 memory_attack in ['memory_enhanced','with_memory', 'new_memory']:
                    result.append(int(line_end_num))
                elif "Memory search success rate: " in line and memory_attack in ['memory_enhanced','with_memory', 'new_memory']:
                    result.append("%.3f"%round(line_end_num,3))

            return result

        # 增加 trigger 列
        if memory_attack in ['memory_enhanced','with_memory','new_memory']:
            result_csv = pandas.DataFrame(columns = ["LLM", "Attack", "Aggressive", "Trigger", "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"])
        else:
            result_csv = pandas.DataFrame(columns = ["LLM", "Attack", "Aggressive", "Trigger", "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}/{memory_attack}"
            agg_result = []
            non_agg_result = []
            for root, dirs, files in os.walk(path, topdown=False):
                files = sorted(files)
                files_list = []

                # 处理所有文件
                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'
                    else:
                        agg = 'Yes'

                    # 提取 trigger，去掉下划线 '_'
                    trigger = name.split('_')[-1].replace('.log', '').replace('_', ' ')
                    if not trigger:  # 处理空 trigger
                        # trigger = "_"
                        continue

                    result = [model, attack_method, agg, trigger]
                    result.extend(get_result(path + '/' + name))

                    # 确保 result 的长度正确
                    if len(result) != 15 and memory_attack in ['memory_enhanced','with_memory','new_memory']:
                        print(path + '/' + name)
                        continue
                    elif len(result) != 13:
                        print(len(result))
                        print(path + '/' + name)
                        continue

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

            # 将结果加入 DataFrame
            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

        result_csv.to_csv(f"./result_csv/result-{prompt_injection}+{memory_attack}.csv", index = False)
