from typing import Literal

from inspect_ai.model import ChatMessageUser
from inspect_ai.solver import (
    Generate,
    Solver,
    TaskState,
    chain,
    generate,
    solver,
    system_message,
)

from ...evals.prompts.synthesize_chemicals_prompts import (
    get_follow_up_prompt,
    get_system_prompt,
)
from .setup_solvers import set_tool_choice, setup_tools, user_prompt


@solver
def default_agent(
    tool_choice: Literal["forced_first", "auto", "none"] = "auto",
    prompt_technique: Literal["standard", "cot"] = "standard",
    summarize_model: str = "openai/gpt-4o-mini-2024-07-18"
) -> Solver:
    # An alternative summarize model: "gpt-4o-2024-08-06"
    agent_setup = [
        setup_tools(summarize_model),
        system_message(get_system_prompt(prompt_technique)),
        user_prompt()
    ]

    if tool_choice in ["none", "auto"]:
        action_loop = [
            set_tool_choice(tool_choice),
            generate(tool_calls="loop")
        ]
    elif tool_choice == "forced_first":
        action_loop = [
            set_tool_choice(tool_choice),
            generate(tool_calls="single"),
            set_tool_choice("auto"),
            generate(tool_calls="loop")
        ]

    follow_up = [
        follow_up_prompts(generate(tool_calls="loop"))
    ]

    plan = chain(*agent_setup, *action_loop, *follow_up)
    return plan


@solver
def follow_up_prompts(action_loop: Solver) -> Solver:
    async def solve(state: TaskState, generate: Generate) -> TaskState:
        listname = state.metadata["listname"]
        level = state.metadata["level"]
        follow_up_prompt = get_follow_up_prompt(listname=listname, level=level)
        state.messages.append(ChatMessageUser(content=follow_up_prompt))
        state = await action_loop(state, generate)

        return state

    return solve
