"""Mock provider for testing."""
from __future__ import annotations

from typing import Any

from logos_adk.providers import Provider
from logos_adk.types import Response, TextChunk, Usage


class MockProvider(Provider):
    """A provider that returns pre-configured responses. For testing only."""

    def __init__(
        self,
        responses: list[str | Response] | None = None,
        stream_responses: list[str | list[str] | Response] | None = None,
    ) -> None:
        self._responses: list[str | Response] = responses or ["mock response"]
        self._index = 0
        self._stream_responses: list[str | list[str] | Response] = (
            stream_responses or list(self._responses)
        )
        self._stream_index = 0
        self.call_history: list[dict[str, Any]] = []

    async def generate(self, messages, tools=None, output_schema=None, **kwargs):
        self.call_history.append({
            "messages": messages,
            "tools": tools,
            "output_schema": output_schema,
            "kwargs": kwargs,
        })
        raw_response = self._responses[self._index]
        self._index = (self._index + 1) % len(self._responses)

        if isinstance(raw_response, Response):
            return raw_response
        return Response(
            content=raw_response,
            tool_calls=[],
            usage=Usage(input_tokens=0, output_tokens=0),
            raw={},
        )

    async def stream(self, messages, tools=None, **kwargs):
        self.call_history.append({
            "messages": messages,
            "tools": tools,
            "output_schema": None,
            "kwargs": kwargs,
            "stream": True,
        })
        raw_response = self._stream_responses[self._stream_index]
        self._stream_index = (self._stream_index + 1) % len(self._stream_responses)

        if isinstance(raw_response, Response):
            if raw_response.content:
                yield TextChunk(text=raw_response.content)
            yield raw_response
            return

        if isinstance(raw_response, str):
            yield TextChunk(text=raw_response)
            return

        for chunk in raw_response:
            yield TextChunk(text=chunk)

    async def generate_batch(
        self,
        batch,
        *,
        tools=None,
        output_schema=None,
        max_concurrency=4,
        **kwargs,
    ):
        results = []
        for messages in batch:
            results.append(
                await self.generate(
                    messages,
                    tools=tools,
                    output_schema=output_schema,
                    **kwargs,
                )
            )
        return results
