"""
Boring-OS SDK models client.

Provides the ModelsClient class for direct LLM operations
via the Model Gateway: completions, streaming, embeddings,
and model listing.
Accessed via `client.models` on the main BoringOS client.
"""

from __future__ import annotations

from typing import TYPE_CHECKING, Any, AsyncIterator

if TYPE_CHECKING:
    from boring_os.client import BoringOS

from boring_os.model_models import (
    ModelCompletion,
    ModelEmbedding,
    ModelInfo,
    ModelStreamChunk,
)


class ModelsClient:
    """
    Client for direct model operations.

    Accessed via `client.models` on the main BoringOS client.
    Wraps the /models/* proxy routes on the gateway.
    """

    def __init__(self, client: "BoringOS"):
        self._client = client

    # -------------------------------------------------------------------------
    # Completion
    # -------------------------------------------------------------------------

    def complete(
        self,
        model: str,
        messages: list[dict[str, Any]],
        temperature: float | None = None,
        max_tokens: int | None = None,
        tools: list[dict[str, Any]] | None = None,
        tenant_id: str | None = None,
    ) -> ModelCompletion:
        """
        Generate a model completion (synchronous).

        Args:
            model: Model name (e.g., "claude-sonnet-4", "gpt-4.1-mini")
            messages: Conversation messages (role + content dicts)
            temperature: Sampling temperature (0.0 - 2.0)
            max_tokens: Max output tokens
            tools: Optional tool definitions for function calling
            tenant_id: Optional tenant ID for billing

        Returns:
            ModelCompletion with text, usage, and cost

        Example:
            completion = boring.models.complete(
                model="claude-sonnet-4",
                messages=[{"role": "user", "content": "Hello"}],
            )
            print(completion.text)
        """
        body: dict[str, Any] = {"model": model, "messages": messages}
        if temperature is not None:
            body["temperature"] = temperature
        if max_tokens is not None:
            body["max_tokens"] = max_tokens
        if tools:
            body["tools"] = tools
        if tenant_id:
            body["tenant_id"] = tenant_id

        response = self._client._request(
            "POST",
            "/models/complete",
            json=body,
            timeout=120,
        )
        return ModelCompletion.model_validate(response)

    async def complete_async(
        self,
        model: str,
        messages: list[dict[str, Any]],
        temperature: float | None = None,
        max_tokens: int | None = None,
        tools: list[dict[str, Any]] | None = None,
        tenant_id: str | None = None,
    ) -> ModelCompletion:
        """Async version of complete()."""
        body: dict[str, Any] = {"model": model, "messages": messages}
        if temperature is not None:
            body["temperature"] = temperature
        if max_tokens is not None:
            body["max_tokens"] = max_tokens
        if tools:
            body["tools"] = tools
        if tenant_id:
            body["tenant_id"] = tenant_id

        response = await self._client._request_async(
            "POST",
            "/models/complete",
            json=body,
            timeout=120,
        )
        return ModelCompletion.model_validate(response)

    # -------------------------------------------------------------------------
    # Streaming
    # -------------------------------------------------------------------------

    async def stream(
        self,
        model: str,
        messages: list[dict[str, Any]],
        temperature: float | None = None,
        max_tokens: int | None = None,
    ) -> AsyncIterator[ModelStreamChunk]:
        """
        Stream a model completion as chunks are generated.

        Args:
            model: Model name
            messages: Conversation messages
            temperature: Sampling temperature
            max_tokens: Max output tokens

        Yields:
            ModelStreamChunk objects; final chunk has done=True

        Example:
            async for chunk in boring.models.stream(
                model="gpt-4.1-mini",
                messages=[{"role": "user", "content": "Write a haiku"}]
            ):
                print(chunk.text, end="")
        """
        body: dict[str, Any] = {"model": model, "messages": messages}
        if temperature is not None:
            body["temperature"] = temperature
        if max_tokens is not None:
            body["max_tokens"] = max_tokens

        async for chunk_data in self._client._stream_request(
            "POST",
            "/models/stream",
            json=body,
        ):
            yield ModelStreamChunk.model_validate(chunk_data)

    # -------------------------------------------------------------------------
    # Embeddings
    # -------------------------------------------------------------------------

    def embed(
        self,
        texts: list[str],
        model: str = "text-embedding-3-small",
        tenant_id: str | None = None,
    ) -> ModelEmbedding:
        """
        Generate embedding vectors for input texts.

        Args:
            texts: List of text strings to embed
            model: Embedding model name
            tenant_id: Optional tenant ID

        Returns:
            ModelEmbedding with embedding vectors

        Example:
            embeddings = boring.models.embed(
                texts=["Hello world", "How are you?"]
            )
            # embeddings.embeddings[0] -> [0.0234, -0.1456, ...]
        """
        body: dict[str, Any] = {"texts": texts, "model": model}
        if tenant_id:
            body["tenant_id"] = tenant_id

        response = self._client._request(
            "POST",
            "/models/embed",
            json=body,
            timeout=60,
        )
        return ModelEmbedding.model_validate(response)

    async def embed_async(
        self,
        texts: list[str],
        model: str = "text-embedding-3-small",
        tenant_id: str | None = None,
    ) -> ModelEmbedding:
        """Async version of embed()."""
        body: dict[str, Any] = {"texts": texts, "model": model}
        if tenant_id:
            body["tenant_id"] = tenant_id

        response = await self._client._request_async(
            "POST",
            "/models/embed",
            json=body,
            timeout=60,
        )
        return ModelEmbedding.model_validate(response)

    # -------------------------------------------------------------------------
    # Model Listing
    # -------------------------------------------------------------------------

    def list(self) -> list[ModelInfo]:
        """
        List all available models with pricing info.

        Returns:
            List of ModelInfo objects
        """
        response = self._client._request("GET", "/models")
        # The response may be a dict with a "models" key or a raw list
        models_data = response.get("models", response) if isinstance(response, dict) else response
        if isinstance(models_data, list):
            return [ModelInfo.model_validate(m) for m in models_data]
        return []
