"""Chat completions (inference)."""
from __future__ import annotations

import json
from typing import TYPE_CHECKING, Generator

if TYPE_CHECKING:
    from .client import Client


class Chat:
    def __init__(self, client: Client):
        self._client = client

    def completions(
        self,
        model: str,
        messages: list[dict],
        *,
        stream: bool = False,
        temperature: float | None = None,
        max_tokens: int | None = None,
        **kwargs,
    ) -> dict | Generator[dict, None, None]:
        """Send a chat completion request.

        When ``stream=True``, returns a generator yielding SSE chunks.
        """
        body: dict = {"model": model, "messages": messages, "stream": stream}
        if temperature is not None:
            body["temperature"] = temperature
        if max_tokens is not None:
            body["max_tokens"] = max_tokens
        body.update(kwargs)

        if not stream:
            return self._client.post("/v1/chat/completions", json=body)

        return self._stream(body)

    def _stream(self, body: dict) -> Generator[dict, None, None]:
        """Handle SSE streaming responses."""
        url = f"{self._client.base_url}/v1/chat/completions"
        resp = self._client._session.post(
            url, json=body, stream=True, timeout=self._client.timeout
        )
        if resp.status_code >= 400:
            # Consume the response and raise through the normal error path
            self._client._request("POST", "/v1/chat/completions", json=body)

        for line in resp.iter_lines(decode_unicode=True):
            if not line or not line.startswith("data: "):
                continue
            data = line[6:]
            if data.strip() == "[DONE]":
                return
            try:
                yield json.loads(data)
            except json.JSONDecodeError:
                continue

    def models(self) -> list[dict]:
        """List available models."""
        return self._client.get("/v1/models")
