import os, yaml

if __name__ == '__main__':

    with open('config/POT.yml', 'r') as file:
        cfg = yaml.safe_load(file)

    llms = cfg.get('llms', None)
    suffix = cfg.get('suffix', None)
    attack_tool_types = cfg.get('attack_tool', None)
    write_db = cfg.get('write_db', None)
    read_db = cfg.get('read_db', None)
    defense_type = cfg.get('defense_type', None) # pot_paraphrase_defense, pot_shuffling_defense
    injection_methods = cfg['injection_method'] # 'pot_backdoor', 'pot_clean', 'clean'
    attack_types = cfg.get('attack_types', None)
    triggers = cfg.get('triggers', None) # 'through a prism of insight', 'with perspicacious discernment', 'conducting a comprehensive elucidation', ':)', '@_@'
    task_num = cfg.get('task_num', None)


    for attack_tool_type in attack_tool_types:
        for injection_method in injection_methods:
            for trigger in triggers:
                for llm in llms:
                    for attack_type in attack_types:
                        if llm.startswith('gpt') or llm.startswith('gemini') or llm.startswith('claude'):
                            llm_name = llm
                            backend=None
                        elif llm.startswith('ollama'):
                            llm_name = llm.split('/')[-1]
                            backend='ollama'
                        else:
                            llm_name = llm.split('/')[-1]
                            backend='vllm'

                        log_path = f'logs/{injection_method}/{llm_name}'

                        database = f'memory_db/direct_prompt_injection/{attack_type}_gpt-4o-mini'

                        if trigger == ':)':
                            suffix = 'smile'
                        else:
                            suffix = trigger.replace(' ', '_')

                        if attack_tool_type == 'all':
                            attacker_tools_path = 'data/all_attack_tools.jsonl'
                        elif attack_tool_type == 'non-agg':
                            attacker_tools_path = 'data/all_attack_tools_non_aggressive.jsonl'
                        elif attack_tool_type == 'agg':
                            attacker_tools_path = 'data/all_attack_tools_aggressive.jsonl'

                        log_memory_type = 'no_memory'
                        log_file = f'{log_path}/{defense_type}/{attack_type}-{attack_tool_type}' if defense_type else f'{log_path}/{log_memory_type}/{attack_type}-{attack_tool_type}'
                        os.makedirs(os.path.dirname(log_file), exist_ok=True)
                        print(log_file)


                        if injection_method in ['pot_backdoor','pot_clean','clean']:
                            cmd = f'''nohup python main_attacker.py --llm_name {llm} --attack_type {attack_type} --use_backend {backend} --attacker_tools_path {attacker_tools_path} \
                                --{injection_method} \
                                --defense_type {defense_type} \
                                --tasks_path 'data/agent_task_pot.jsonl' \
                                --trigger '{trigger}' \
                                --task_num {task_num} \
                                > {log_file}_{suffix}.log 2>&1 &'''

                        os.system(cmd)