from __future__ import annotations

"""Test utilities for SDK users and internal modules."""

from collections import deque
from typing import AsyncIterator

from .errors import ProviderError
from .interfaces import LlmProvider
from .types import CompletionRequest, CompletionResponse, StopReason, StreamChunk


class MockLlmProvider(LlmProvider):
    """Mock provider for deterministic tests."""

    def __init__(
        self, responses: list[CompletionResponse] | None = None, name: str = "mock"
    ) -> None:
        self._name = name
        self._responses: deque[CompletionResponse] = deque(responses or [])

    @classmethod
    def new(cls) -> "MockLlmProvider":
        return cls()

    @classmethod
    def with_responses(cls, responses: list[CompletionResponse]) -> "MockLlmProvider":
        return cls(responses=responses)

    @classmethod
    def with_text_responses(cls, texts: list[str]) -> "MockLlmProvider":
        return cls.with_responses([CompletionResponse.text(text) for text in texts])

    @classmethod
    def with_tool_call_sequence(
        cls,
        tool_name: str,
        arguments: object,
        final_text: str,
    ) -> "MockLlmProvider":
        return cls.with_responses(
            [
                CompletionResponse.tool_call(tool_name, arguments),
                CompletionResponse.text(final_text),
            ]
        )

    async def push_response(self, response: CompletionResponse) -> None:
        self._responses.append(response)

    async def _pop_response(self) -> CompletionResponse:
        if not self._responses:
            raise ProviderError.invalid_response("mock provider exhausted")
        return self._responses.popleft()

    async def complete(self, request: CompletionRequest) -> CompletionResponse:
        del request
        return await self._pop_response()

    async def stream(self, request: CompletionRequest) -> AsyncIterator[StreamChunk]:
        del request
        response = await self._pop_response()

        async def iterator() -> AsyncIterator[StreamChunk]:
            for token in response.message.text_content().split():
                yield StreamChunk(delta=f"{token} ")
            yield StreamChunk(
                delta="", finish_reason=response.stop_reason or StopReason.STOP
            )

        return iterator()

    def name(self) -> str:
        return self._name
