"""Test utilities for donkit-llm-agent. Reusable mock classes for testing agents."""

from __future__ import annotations

import json
from typing import Any, AsyncIterator

from donkit.llm import (
    FunctionCall,
    GenerateRequest,
    GenerateResponse,
    LLMModelAbstract,
    Message,
    ModelCapability,
    StreamChunk,
    ToolCall,
)


class BaseMockProvider(LLMModelAbstract):
    """Mock LLM provider for testing with configurable responses."""

    def __init__(
        self,
        responses: list[dict[str, Any]] | None = None,
        supports_tools_val: bool = True,
        supports_streaming_val: bool = False,
        model_name_val: str = "mock-model",
    ) -> None:
        self.call_count = 0
        self.stream_call_count = 0
        self.messages_history: list[list[Message]] = []
        self.responses = responses or [{"content": "Default response"}]
        self.supports_tools_val = supports_tools_val
        self.supports_streaming_val = supports_streaming_val
        self._model_name = model_name_val

    @property
    def model_name(self) -> str:
        return self._model_name

    @model_name.setter
    def model_name(self, value: str) -> None:
        self._model_name = value

    @property
    def capabilities(self) -> ModelCapability:
        caps = ModelCapability.TEXT_GENERATION
        if self.supports_tools_val:
            caps |= ModelCapability.TOOL_CALLING
        if self.supports_streaming_val:
            caps |= ModelCapability.STREAMING
        return caps

    async def generate(self, request: GenerateRequest) -> GenerateResponse:
        self.call_count += 1
        self.messages_history.append(request.messages.copy())
        response_idx = (self.call_count - 1) % len(self.responses)
        response_config = self.responses[response_idx]

        tool_calls = None
        if "tool_calls" in response_config:
            tool_calls = []
            for i, tc in enumerate(response_config["tool_calls"]):
                args = tc.get("arguments", {})
                if isinstance(args, str):
                    try:
                        args = json.loads(args)
                    except (json.JSONDecodeError, ValueError):
                        args = {}
                tool_calls.append(
                    ToolCall(
                        id=tc.get("id", f"call_{i}"),
                        type="function",
                        function=FunctionCall(
                            name=tc["name"], arguments=json.dumps(args)
                        ),
                    )
                )

        return GenerateResponse(
            content=response_config.get("content"), tool_calls=tool_calls
        )

    async def generate_stream(
        self, request: GenerateRequest
    ) -> AsyncIterator[StreamChunk]:
        self.stream_call_count += 1
        self.messages_history.append(request.messages.copy())

        for response_config in self.responses:
            tool_calls = None
            if "tool_calls" in response_config:
                tool_calls = []
                for i, tc in enumerate(response_config["tool_calls"]):
                    args = tc.get("arguments", {})
                    if isinstance(args, str):
                        try:
                            args = json.loads(args)
                        except (json.JSONDecodeError, ValueError):
                            args = {}
                    tool_calls.append(
                        ToolCall(
                            id=tc.get("id", f"call_{i}"),
                            type="function",
                            function=FunctionCall(
                                name=tc["name"], arguments=json.dumps(args)
                            ),
                        )
                    )

            yield StreamChunk(
                content=response_config.get("content"), tool_calls=tool_calls
            )


class BaseMockMCPClient:
    """Mock MCP client for testing."""

    def __init__(self, name: str, tools: dict[str, Any]) -> None:
        self.name = name
        self.identifier = name
        self.tools = tools
        self.call_count = 0

    async def connect(self) -> None:
        pass

    async def disconnect(self) -> None:
        pass

    async def alist_tools(self) -> list[dict]:
        return [
            {
                "name": tool_name,
                "description": tool_info.get("description", ""),
                "parameters": tool_info.get("parameters", {}),
            }
            for tool_name, tool_info in self.tools.items()
        ]

    async def acall_tool(self, name: str, arguments: dict[str, Any]) -> str:
        self.call_count += 1
        if name not in self.tools:
            raise ValueError(f"Tool {name} not found")
        handler = self.tools[name].get("handler")
        if handler:
            return handler(arguments)
        return json.dumps({"result": "success"})
