"""Custom solvers for the AssistantBench eval."""

from inspect_ai.model import ChatMessageUser
from inspect_ai.solver import Generate, Solver, TaskState, solver
from inspect_ai.tool import Tool, web_search

from inspect_evals.assistant_bench.prompts import (
    WEB_SEARCH_ANSWER_PREFIX,
    WEB_SEARCH_REQUEST_PREFIX,
    WEB_SEARCH_RESPONSE,
)

# This matches the limit used in the original paper for retrieval-augmented models.
DEFAULT_MAX_SEARCH_RESULTS = 10

# When performing web search the majority of token usage is on the model which evaluates
# the web search results to determine the most relevant one(s).  By default force a
# cheap model for that to keep costs down.
DEFAULT_WEB_SEARCH_MODEL = "openai/gpt-4o-mini"

SEARCH_LIMIT_REACHED_KEY = "search_limit_reached"


class Query(str):
    pass


class ErrorHint(str):
    pass


def parse_model_response(response: str) -> Query | ErrorHint:
    """Process a model response and return the necessary context."""
    if WEB_SEARCH_REQUEST_PREFIX not in response:
        return ErrorHint(
            f"Responses must include '{WEB_SEARCH_ANSWER_PREFIX}' or '{WEB_SEARCH_REQUEST_PREFIX}'"
        )

    query_str_splits = response.split(WEB_SEARCH_REQUEST_PREFIX)
    if len(query_str_splits) > 2:
        return ErrorHint("Responses must only include a single intermediate query.")

    query_str = query_str_splits[1].strip()
    if any(char.isascii() and not char.isprintable() for char in query_str):
        return ErrorHint(
            "Intermediate query can only contain printable ASCII characters."
        )

    return Query(query_str)


@solver
def indirect_web_search(
    search_tool: Tool | None = None,
    max_searches: int = DEFAULT_MAX_SEARCH_RESULTS,
) -> Solver:
    """Implement the 'Retrieval Augmented' model from the paper.

    In a loop, prompt the LLM to produce intermediate questions and then search the
    web to provide answers. Wrap the web_search results in additional task prompt reminders.
    """
    # It looks like the paper only returns a single result.
    # By default, force a cheaper model to analyse the search results to reduce costs.
    search_tool = search_tool or web_search(
        num_results=1, model=DEFAULT_WEB_SEARCH_MODEL
    )

    async def solve(state: TaskState, generate: Generate) -> TaskState:
        # We know that all our tasks include at least a system message and the task
        assert isinstance(state.input, list)
        original_task = state.input[-1].text
        all_search_queries: list[str] = []

        # We need +1 on the for loop to give the model a chance to provide the answer
        # after the final search result.
        for safety_loop in range(max_searches + 1):
            await generate(state)
            response = state.output.message.text

            if WEB_SEARCH_ANSWER_PREFIX in response:
                break
            if safety_loop == max_searches:
                state.metadata[SEARCH_LIMIT_REACHED_KEY] = True
                break

            parsed = parse_model_response(response)
            if isinstance(parsed, ErrorHint):
                context = str(parsed)
                query_str = "No valid intermediate query specified."
            elif str(parsed) in all_search_queries:
                context = "That is a repeat web search.  Do not generate the same intermediate question twice."
                query_str = str(parsed)
            else:
                query_str = str(parsed)
                all_search_queries.append(query_str)

                search_results = await search_tool(query_str)
                # We always expect web searches to return strings.
                assert isinstance(search_results, str)
                context = search_results

            search_response = WEB_SEARCH_RESPONSE.format(
                count=safety_loop,
                search_results=context,
                original_task=original_task,
                intermediate_question=query_str,
            )
            state.messages.append(ChatMessageUser(content=search_response))

        return state

    return solve
