"""Ollama LLM client for local models (build plan Section 4)."""

from __future__ import annotations

import httpx

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


class OllamaClient:
    provider = "ollama"

    def __init__(
        self,
        *,
        base_url: str = "http://localhost:11434",
        timeout: float = 120.0,
        http_client: httpx.AsyncClient | None = None,
    ) -> None:
        self._base_url = base_url.rstrip("/")
        self._timeout = timeout
        self._http = http_client

    async def complete(
        self,
        messages: list[LLMMessage],
        *,
        model: str,
        temperature: float = 0.0,
        max_tokens: int | None = None,
    ) -> LLMResponse:
        options: dict = {"temperature": temperature}
        if max_tokens is not None:
            options["num_predict"] = max_tokens
        payload = {
            "model": model,
            "messages": [m.as_dict() for m in messages],
            "stream": False,
            "options": options,
        }
        try:
            if self._http is not None:
                resp = await self._http.post(f"{self._base_url}/api/chat", json=payload)
            else:
                async with httpx.AsyncClient(timeout=self._timeout) as client:
                    resp = await client.post(f"{self._base_url}/api/chat", json=payload)
            resp.raise_for_status()
            data = resp.json()
        except httpx.HTTPError as exc:
            raise LLMError(f"ollama request failed: {exc}") from exc

        try:
            content = data["message"]["content"]
        except (KeyError, TypeError) as exc:
            raise LLMError(f"ollama returned an unexpected response: {data}") from exc

        return LLMResponse(
            content=content or "",
            model=data.get("model", model),
            prompt_tokens=int(data.get("prompt_eval_count", 0)),
            completion_tokens=int(data.get("eval_count", 0)),
        )
