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


@solver
def solve_scicode_problem(prompt_template: str) -> Solver:
    """
    Return a function that generates a solution to a SciCode problem.

    A SciCode problem is composed of a sequence of subproblems that are solved sequentially. For each subproblem, the solver formats a prompt using the provided template, appends a request to solve the subproblem to `state.messages`, and invokes `generate` to obtain the model's solution.

    Parameters:
    -----------
    prompt_template: str
        A template string to format prompts for each subproblem.

    Returns:
    -------
    Solver

    """

    async def solve(state: TaskState, generate: Generate) -> TaskState:
        for subproblem in state.metadata["sub_steps"]:
            formatted_prompt = prompt_template.format(**subproblem)
            chat_message = ChatMessageUser(content=formatted_prompt)
            state.messages.append(chat_message)
            state = await generate(state)
        return state

    return solve
