import dataclasses
import re
from typing import Literal

import pytest
from inspect_ai.model import (
    CachePolicy,
    ChatMessageAssistant,
    ChatMessageUser,
    GenerateConfigArgs,
    ModelName,
    ModelOutput,
)
from inspect_ai.solver import Generate, TaskState
from inspect_ai.tool import Tool, ToolResult
from typing_extensions import Unpack

from inspect_evals.assistant_bench.solver import (
    DEFAULT_MAX_SEARCH_RESULTS,
    SEARCH_LIMIT_REACHED_KEY,
    ErrorHint,
    Query,
    indirect_web_search,
    parse_model_response,
)


class TestParseModelResponse:
    def test_mainline(self) -> None:
        parsed = parse_model_response("Intermediate question: a valid question")
        assert isinstance(parsed, Query)
        assert parsed == "a valid question"

    def test_no_intermediate_query(self) -> None:
        parsed = parse_model_response("A random answer without the right cues.")
        assert isinstance(parsed, ErrorHint)
        assert parsed.startswith("Responses must include")

    def test_multiple_intermediate_queries(self) -> None:
        parsed = parse_model_response(
            "Intermediate question: first question\nIntermediate question: second question"
        )
        assert isinstance(parsed, ErrorHint)
        assert parsed == "Responses must only include a single intermediate query."

    @pytest.mark.parametrize("character", ["\n", "\r", "\t"])
    def test_invalid_characters(self, character: str) -> None:
        parsed = parse_model_response(
            f"Intermediate question: first {character} question"
        )
        assert isinstance(parsed, ErrorHint)
        assert (
            parsed == "Intermediate query can only contain printable ASCII characters."
        )


@dataclasses.dataclass
class SolveIteration:
    model_response: str

    # By default the context is checked to include the previous search result
    # but this allows checking of custom errors.
    custom_context: str | None = None

    # These can be None on error (e.g., repeat search) or final answer.
    search_query: str | None = None
    search_result: str | None = None


SolveIterations = list[SolveIteration]


class TestIndirectWebSearch:
    MODEL_NAME = "openai/gpt-4o-mini"

    def setup_method(self, _) -> None:
        # We increment first since generate is the only mock guaranteed to be called
        # but it starts ths sequence.
        self.sequence_index = -1

    def mock_generate(self, iterations: SolveIterations) -> Generate:
        async def generate(
            state: TaskState,
            tool_calls: Literal["loop", "single", "none"] = "loop",
            cache: bool | CachePolicy = False,
            **kwargs: Unpack[GenerateConfigArgs],
        ) -> TaskState:
            # Since generate is always called it takes responsibility for
            # updating the sequence identifier.
            self.sequence_index += 1

            # No context on the first loop since no model responses yet.
            if self.sequence_index > 0:
                context = iterations[self.sequence_index].custom_context
                if context is None:
                    context = iterations[self.sequence_index - 1].search_result
                pattern = re.compile(f"Context\\d+: {context}")

                assert len(state.messages) > 0
                last_message = state.messages[-1]
                assert isinstance(last_message, ChatMessageUser)
                assert isinstance(last_message.content, str)
                assert pattern.search(last_message.content) is not None

            model_response = iterations[self.sequence_index].model_response
            state.messages.append(ChatMessageAssistant(content=model_response))
            state.output = ModelOutput.from_content(
                TestIndirectWebSearch.MODEL_NAME, model_response
            )

            return state

        return generate

    def mock_search_tool(self, iterations: SolveIterations) -> Tool:
        async def search_result(search_query: str) -> ToolResult:
            # The index is updated by mock_generate() since that is always called.
            expected_query = iterations[self.sequence_index].search_query
            assert expected_query is not None
            assert search_query == expected_query
            search_result = iterations[self.sequence_index].search_result
            assert search_result is not None
            return search_result

        return search_result

    async def run_sequence(self, iterations: SolveIterations) -> TaskState:
        task_state = TaskState(
            model=ModelName(TestIndirectWebSearch.MODEL_NAME),
            sample_id="abc",
            epoch=1,
            input=[ChatMessageUser(content="An example task")],
            messages=[],
        )

        solver = indirect_web_search(self.mock_search_tool(iterations))
        generate = self.mock_generate(iterations)

        final_task_state = await solver(task_state, generate)

        assert self.sequence_index == len(iterations) - 1
        if iterations[-1].search_query is None:
            assert isinstance(final_task_state.messages[-1], ChatMessageAssistant)
        else:
            assert isinstance(final_task_state.messages[-1], ChatMessageUser)

        return final_task_state

    @pytest.mark.asyncio
    async def test_mainline_solve(self) -> None:
        sequence = [
            SolveIteration(
                model_response="Intermediate question: my search query",
                search_query="my search query",
                search_result="my definitive search response",
            ),
            SolveIteration("The final answer is: my answer"),
        ]

        await self.run_sequence(sequence)

    @pytest.mark.asyncio
    async def test_repeat_query(self) -> None:
        sequence = [
            SolveIteration(
                model_response="Intermediate question: my search query",
                search_query="my search query",
                search_result="my definitive search response",
            ),
            SolveIteration(
                model_response="Intermediate question: my search query",
            ),
            SolveIteration(
                custom_context="That is a repeat web search",
                model_response="The final answer is: my answer",
            ),
        ]

        await self.run_sequence(sequence)

    @pytest.mark.asyncio
    async def test_search_max_searches(self) -> None:
        sequence = [
            SolveIteration(
                model_response=f"Intermediate question: question{i}",
                search_query=f"question{i}",
                search_result=f"answer{i}",
            )
            for i in range(DEFAULT_MAX_SEARCH_RESULTS)
        ]
        sequence.append(SolveIteration("The final answer is: my answer"))

        final_task_state = await self.run_sequence(sequence)
        assert not final_task_state.metadata.get(SEARCH_LIMIT_REACHED_KEY, False)

    @pytest.mark.asyncio
    async def test_search_exceed_search_limit(self) -> None:
        sequence = [
            SolveIteration(
                model_response=f"Intermediate question: question{i}",
                search_query=f"question{i}",
                search_result=f"answer{i}",
            )
            for i in range(DEFAULT_MAX_SEARCH_RESULTS)
        ]
        sequence.append(SolveIteration("Intermediate question: exceeding question"))

        final_task_state = await self.run_sequence(sequence)
        assert final_task_state.metadata[SEARCH_LIMIT_REACHED_KEY]
