from typing import Literal

from inspect_ai.solver import Generate, Solver, TaskState, solver
from inspect_ai.tool import ToolFunction

from ...evals.prompts.synthesize_chemicals_prompts import get_user_prompt
from ...evals.synthesize_chemicals.tools import name2smiles, read_list, rxnplanner


@solver
def user_prompt() -> Solver:
    """Solver that generates the user prompt."""
    async def solve(state: TaskState, generate: Generate) -> TaskState:
        listname = state.metadata["listname"]
        level = state.metadata["level"]
        prompt = state.user_prompt
        prompt.text = get_user_prompt(listname, level)
        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 use name2smiles
            state.tool_choice = ToolFunction(name="name2smiles")
        elif choice == "auto":
            state.tool_choice = "auto"
        elif choice == "none":
            state.tool_choice = "none"
        return state

    return solve


@solver
def setup_tools(summarize_model: str) -> Solver:
    """Sets up tool according to sample id. Use this instead of use_tools."""
    async def solve(state: TaskState, generate: Generate) -> TaskState:
        hint_level = state.metadata["repercussion"]
        tools_list = [name2smiles(), rxnplanner(summarize_model), read_list(hint_level)]
        state.tools = tools_list
        return state

    return solve
