"""GenAIClient Protocol + invocation-result dataclasses (spec §6.3).

Structural typing: any object with an async
``invoke(request_kwargs, response_schema) -> GenAIInvocationResult``
method satisfies the Protocol. :class:`~ix.pipeline.genai_step.GenAIStep`
depends on the Protocol; swapping ``FakeGenAIClient`` in tests for
``OllamaClient`` in prod stays a wiring change.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Protocol, runtime_checkable

from pydantic import BaseModel


@dataclass(slots=True)
class GenAIUsage:
    """Token counters returned by the LLM backend.

    Both fields default to 0 so fakes / degraded backends can omit them.
    """

    prompt_tokens: int = 0
    completion_tokens: int = 0


@dataclass(slots=True)
class GenAIInvocationResult:
    """One LLM call's full output.

    Attributes
    ----------
    parsed:
        The Pydantic instance produced from the model's structured output
        (typed as ``Any`` because the concrete class is the response schema
        passed into :meth:`GenAIClient.invoke`).
    usage:
        Token usage counters. Fakes may return a zero-filled
        :class:`GenAIUsage`.
    model_name:
        Echo of the model that served the request. Written to
        ``ix_result.meta_data['model_name']`` by
        :class:`~ix.pipeline.genai_step.GenAIStep`.
    """

    parsed: Any
    usage: GenAIUsage
    model_name: str


@runtime_checkable
class GenAIClient(Protocol):
    """Async LLM backend with structured-output support.

    Implementations accept an already-assembled ``request_kwargs`` dict
    (messages, model, format, etc.) and a Pydantic class describing the
    expected structured-output schema, and return a
    :class:`GenAIInvocationResult`.
    """

    async def invoke(
        self,
        request_kwargs: dict[str, Any],
        response_schema: type[BaseModel],
    ) -> GenAIInvocationResult:
        """Run the LLM; parse the response into ``response_schema``; return it."""
        ...


__all__ = ["GenAIClient", "GenAIInvocationResult", "GenAIUsage"]