from typing import Literal

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

from ...evals.prompts.rank_organization_prompts import get_follow_up_prompt
from .setup_solvers import (
    create_file_copies,
    set_decision_made_variable,
    set_tool_choice,
    setup_tools,
    system_prompt,
    user_prompt,
)


@solver
def default_agent(
    tool_choice: Literal["forced_first", "auto", "none"] = "auto", #= "forced_first",
    prompt_technique: Literal["standard", "cot"] = "standard"
) -> Solver:
    agent_setup = [
        create_file_copies(),
        set_decision_made_variable(decision_made=False),
        setup_tools(),
        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 = [
        set_decision_made_variable(decision_made=True),
        setup_tools(),
        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:
        scenario = state.metadata["scenario"]
        follow_up_prompt = get_follow_up_prompt(scenario=scenario)
        state.messages.append(ChatMessageUser(content=follow_up_prompt))
        state = await action_loop(state, generate)

        return state

    return solve
