import logging
from random import Random

from inspect_ai.model import ChatMessage, ChatMessageUser
from inspect_ai.solver import Generate, Solver, TaskState, solver

random = Random(42)

logger = logging.getLogger(__name__)


@solver
def solve(prompt_template: str, messages: list[ChatMessage]) -> Solver:
    """Return a function that generates a solution to a Mind2Web problem.

    This function creates a solver that formats the prompt with the HTML snippet,
    task description, previous actions, and multiple-choice options. It then
    passes this formatted prompt to the language model to generate a response.

    Args:
        prompt_template: A template string to format prompts for each sample.
            The template should include placeholders for {final_html},
            {confirmed_task}, {previous_actions}, and {choices_text}.
        messages: A list of ChatMessage objects to include in the conversation
            history before the formatted prompt. These typically include
            system instructions and few-shot examples.

    Returns:
        A Solver function that can be used to solve Mind2Web tasks.
    """

    async def solve(state: TaskState, generate: Generate) -> TaskState:
        state.messages.extend(messages)
        choices = state.choices
        choices_text = "\n".join([c.value for i, c in enumerate(choices)])
        formatted_prompt = prompt_template.format(
            choices_text=choices_text, **state.metadata
        )
        chat_message = ChatMessageUser(content=formatted_prompt)
        state.messages.append(chat_message)
        new_state: TaskState = await generate(state=state)
        if len(new_state.output.choices) > 1:
            logger.debug(f"mulitple choices {new_state}")
        return new_state

    return solve
