"""A scripted LLM client for tests.

Returns pre-baked responses in order and records the prompts it was given so
tests can assert on the loop's behaviour without any network access.
"""

from __future__ import annotations

from legate.llm.base import LLMError, LLMMessage, LLMResponse


def _estimate_tokens(text: str) -> int:
    return max(1, len(text) // 4)


class FakeLLMClient:
    provider = "fake"

    def __init__(self, responses: list[str], *, model: str = "fake-model") -> None:
        self._responses = list(responses)
        self._model = model
        self.calls: list[list[LLMMessage]] = []

    async def complete(
        self,
        messages: list[LLMMessage],
        *,
        model: str,
        temperature: float = 0.0,
        max_tokens: int | None = None,
    ) -> LLMResponse:
        self.calls.append(list(messages))
        if not self._responses:
            raise LLMError("FakeLLMClient ran out of scripted responses")
        content = self._responses.pop(0)
        prompt_tokens = sum(_estimate_tokens(m.content) for m in messages)
        return LLMResponse(
            content=content,
            model=model or self._model,
            prompt_tokens=prompt_tokens,
            completion_tokens=_estimate_tokens(content),
        )
