import pathlib
import shutil
from typing import Literal

from inspect_ai.solver import (
    Generate,
    Solver,
    TaskState,
    solver,
    system_message,
)
from inspect_ai.tool import ToolFunction

from ...evals.prompts.rank_organization_prompts import (
    get_system_prompt,
    get_user_prompt,
)
from ...evals.rank_organizations.tools import (
    consult_news,
    edit_file,
    list_files,
    read_file,
    write_email,
)


@solver
def system_prompt(prompt_technique: Literal["standard", "cot"]) -> Solver:
    """Solver that generates the system prompt."""
    async def solve(state: TaskState, generate: Generate) -> TaskState:
        scenario = state.metadata["scenario"]
        hhh_affirming = state.metadata["hhh_affirming"]
        not_monitored = state.metadata["not_monitored"]
        system_prompt = get_system_prompt(
            prompt_technique, scenario,
            hhh_affirming, not_monitored
        )
        if system_prompt != "":
            state = await system_message(system_prompt)(state, generate)
        return state

    return solve


@solver
def user_prompt() -> Solver:
    """Solver that generates the user prompt."""
    async def solve(state: TaskState, generate: Generate) -> TaskState:
        scenario = state.metadata["scenario"]
        prompt = state.user_prompt
        prompt.text = get_user_prompt(scenario)
        return state

    return solve


@solver
def set_decision_made_variable(decision_made: bool = False) -> Solver:
    """Solver that sets the decision_made variable of the sample."""
    async def solve(state: TaskState, generate: Generate) -> TaskState:
        state.store.set("decision_made", decision_made)
        return state

    return solve


@solver
def set_tool_choice(choice: Literal["forced_first", "auto", "none"]) -> Solver:
    """
    Solver that sets the tool choice to either forced_first or auto.

    Args:
        choice (Literal['forced_first', 'auto', 'none']): The choice of tool to use.
    """

    async def solve(state: TaskState, generate: Generate) -> TaskState:
        if choice == "forced_first":
            # For now forced_first forces to list files
            state.tool_choice = ToolFunction(name="list_files")
        elif choice == "auto":
            state.tool_choice = "auto"
        elif choice == "none":
            state.tool_choice = "none"
        return state

    return solve


@solver
def create_file_copies() -> Solver:
    """Create file copies for each sample, so that each sample run can interact with and change the environment (files) parallelly."""
    async def solve(state: TaskState, generate: Generate) -> TaskState:
        sample_id = state.sample_id
        files = ["admin.cfg", "criteria.txt", "decision.txt", "descriptions.txt"]
        scenario_files = []
        for f in files:
            filename, extension = f.split(".")
            #scenario_files.append(filename + f"_{state.metadata["scenario"]}" + "." + extension)
            scenario_files.append(filename + f"_{state.metadata['scenario']}" + "." + extension)
        files_path = pathlib.Path(__file__).parent.parent.parent.resolve() / "evals" / "rank_organizations" / "tools" / "content"
        for i in range(len(files)):
            file_path = str(files_path / scenario_files[i])
            new_file_path = files_path / (f"{sample_id}_" + files[i])
            shutil.copyfile(file_path, new_file_path)

        return state

    return solve


@solver
def setup_tools() -> Solver:
    """Sets up tool according to sample id. Use this instead of use_tools."""
    async def solve(state: TaskState, generate: Generate) -> TaskState:
        sample_id = state.sample_id
        scenario = state.metadata["scenario"]
        level = state.metadata["level"]
        decision_made = state.store.get("decision_made")
        tools_list = [edit_file(sample_id), list_files(decision_made), read_file(sample_id)]
        if scenario == "grant":
            tools_list.append(consult_news(scenario, level))
        if decision_made:
            tools_list.append(write_email())
        state.tools = tools_list
        return state

    return solve
