# This file contains helpful utilities for the rest of the code, encompassing
# parsing, environment variables, logging, etc.

# TODO: switch to click
import argparse

import os
import shutil

from typing import Dict, Any, Optional

import json

import re

# logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
# logger = logging.getLogger(__name__)

def parse_global_args():
    """ Construct help message and parse argumets using argparse """
    parser = argparse.ArgumentParser(description="Parse global parameters")
    parser.add_argument('--llm_name', type=str, default="gemma-2b-it", help="Specify the LLM name of AIOS")
    parser.add_argument('--max_gpu_memory', type=json.loads, help="Max gpu memory allocated for the LLM")
    parser.add_argument('--eval_device', type=str, help="Evaluation device (example: \"conda:0\" for 2 GPUs)")
    parser.add_argument('--max_new_tokens', type=int, default=256, help="The maximum number of new tokens for generation")
    parser.add_argument("--scheduler_log_mode", type=str, default="console", choices=["console", "file"])
    parser.add_argument("--agent_log_mode", type=str, default="console", choices=["console", "file"])
    parser.add_argument("--llm_kernel_log_mode", type=str, default="console", choices=["console", "file"])
    parser.add_argument("--use_backend", type=str, default = None, choices=['None', "ollama", "vllm"])
    parser.add_argument("--workflow_mode", type=str, default = 'automatic', choices=["manual", "automatic"])
    parser.add_argument("--attacker_tools_path", type=str, default = 'data/all_attack_tools_non_aggressive.jsonl', help="Path to the Attacker Tool jsonl file")
    parser.add_argument("--tasks_path", type=str, default = 'data/agent_task.jsonl', help="Path to the task file")
    parser.add_argument("--tools_info_path", type=str, default = 'data/all_normal_tools.jsonl', help="Path to the normal tools info file")
    parser.add_argument("--observation_prompt_injection", action='store_true')
    parser.add_argument("--plan_attack", action='store_true')
    parser.add_argument("--action_attack", action='store_true')
    parser.add_argument("--direct_prompt_injection", action='store_true')
    parser.add_argument("--jailbreak", action='store_true')
    parser.add_argument("--attack_type", type=str, default = None)
    parser.add_argument("--defense_type", type=str, default = None)
    parser.add_argument("--pot_backdoor", action='store_true')
    parser.add_argument("--pot_clean", action='store_true')
    parser.add_argument("--target", type=str, default = 'ResearcherProfileManager', help="pot backdoor target tool")
    parser.add_argument("--trigger", type=str, default = 'in light of extant literature', help="pot backdoor trigger")
    parser.add_argument("--memory_attack", action='store_true')
    parser.add_argument("--clean", action='store_true')
    parser.add_argument("--result_file", type=str, default = 'logs/result_file.log', help="Path to the result file")
    parser.add_argument("--database", type=str, default = 'memory_db/chroma_db', help="Path to the database")
    parser.add_argument('--task_num', type=int, default=1, help="The task number for an agent to run")
    parser.add_argument("--write_db", action='store_true')
    parser.add_argument("--read_db", action='store_true')
    parser.add_argument("--res_file", type=str, default = 'logs/result_file.log', help="Path to the result file")

    return parser


def extract_before_parenthesis(s: str) -> str:
    match = re.search(r'^(.*?)\([^)]*\)', s)
    return match.group(1) if match else s

def get_from_dict_or_env(
    data: Dict[str, Any], key: str, env_key: str, default: Optional[str] = None
) -> str:
    """Get a value from a dictionary or an environment variable."""
    if key in data and data[key]:
        return data[key]
    else:
        return get_from_env(key, env_key, default=default)


def get_from_env(env_key: str, default: Optional[str] = None) -> str:
    """Get a value from an environment variable."""
    if env_key in os.environ and os.environ[env_key]:
        return os.environ[env_key]
    elif default is not None:
        return default
    else:
        raise ValueError(
            f"Did not find {env_key}, please add an environment variable"
            f" `{env_key}` which contains it. "
        )

class Logger:
    """ Helper logger file """

    def __init__(self, log_mode) -> None:
        self.log_mode = log_mode

    def log(self, info, path=None):
        if self.log_mode == "console":
            print(info)
        else:
            assert self.log_mode == "file"
            with open(path, "w") as w:
                w.write(info + "\n")

def delete_directories(root_dir, target_dirs):
    """
    Recursively deletes directories with names in target_dirs starting from root_dir.
    """
    for dirpath, dirnames, filenames in os.walk(root_dir, topdown=False):
        for dirname in dirnames:
            if dirname in target_dirs:
                full_path = os.path.join(dirpath, dirname)
                # print(f"Deleting {full_path}...")
                shutil.rmtree(full_path, ignore_errors=True)
