"""Ollama backend client for sageLLM CLI (#16).

Provides a thin wrapper around the Ollama REST API so that sageLLM can
delegate inference to a locally-running Ollama instance when no native
backend is available.

The implementation follows the same interface as :class:`UnifiedInferenceClient`
so it can be used as a drop-in alternative without changing calling code.

API coverage
------------
- ``POST /api/generate``   → :meth:`OllamaClient.complete`
- ``POST /api/chat``       → :meth:`OllamaClient.chat`
- ``POST /api/embeddings`` → :meth:`OllamaClient.embed`
- ``GET  /api/tags``       → :meth:`OllamaClient.list_models`
- ``GET  /api/version``    → :meth:`OllamaClient.version`
- ``GET  /``               → :meth:`OllamaClient.health`

Reference: https://github.com/ollama/ollama/blob/main/docs/api.md

Issue
-----
#16 Implement Ollama backend
"""

from __future__ import annotations

import json
import logging
from typing import Any
from urllib import error as urllib_error
from urllib import request as urllib_request

logger = logging.getLogger(__name__)

_OLLAMA_DEFAULT_URL = "http://localhost:11434"


class OllamaClient:
    """Ollama REST API client.

    Provides a synchronous interface to a running Ollama server.

    Args:
        base_url: Ollama server URL (default ``http://localhost:11434``).
        model: Default model name (e.g. ``"llama3"``).
        timeout_s: HTTP request timeout.
        offline_mode: When ``True`` return canned responses without
            making real HTTP calls.

    Example::

        client = OllamaClient(model="llama3")
        resp = client.complete("What is 2+2?")
        print(resp["text"])
    """

    def __init__(
        self,
        base_url: str = _OLLAMA_DEFAULT_URL,
        *,
        model: str = "llama3",
        timeout_s: float = 120.0,
        offline_mode: bool = False,
    ) -> None:
        self._base_url = base_url.rstrip("/")
        self._model = model
        self._timeout_s = timeout_s
        self._offline_mode = offline_mode

    @property
    def base_url(self) -> str:
        return self._base_url

    # ------------------------------------------------------------------
    # Health / metadata
    # ------------------------------------------------------------------

    def health(self) -> bool:
        """Return ``True`` if the Ollama server is reachable."""
        if self._offline_mode:
            return True
        try:
            self._get("/")
            return True
        except Exception:
            return False

    def version(self) -> str:
        """Return Ollama version string, e.g. ``"0.3.6"``."""
        if self._offline_mode:
            return "offline"
        data = self._get_json("/api/version")
        return str(data.get("version", "unknown"))

    def list_models(self) -> list[str]:
        """Return a list of locally available model names.

        Returns:
            List of model name strings (e.g. ``["llama3:latest", ...]``).
        """
        if self._offline_mode:
            return [self._model]
        data = self._get_json("/api/tags")
        return [m.get("name", "") for m in data.get("models", [])]

    # ------------------------------------------------------------------
    # Inference: completion
    # ------------------------------------------------------------------

    def complete(
        self,
        prompt: str,
        *,
        model: str | None = None,
        max_tokens: int = 128,
        temperature: float = 0.0,
        stream: bool = False,
    ) -> dict[str, Any]:
        """Generate a text completion.

        Wraps ``POST /api/generate``.

        Args:
            prompt: Input text.
            model: Model override.
            max_tokens: Soft max token limit passed via ``options.num_predict``.
            temperature: Sampling temperature.
            stream: Whether to stream tokens (not yet fully supported —
                returns complete response).

        Returns:
            Dict with ``"text"``, ``"usage"``, and ``"raw"`` keys.
        """
        if self._offline_mode:
            return {
                "text": f"[ollama-offline] {prompt[:40]}...",
                "usage": {"prompt_tokens": len(prompt.split()), "completion_tokens": 5},
                "raw": {},
            }

        payload: dict[str, Any] = {
            "model": model or self._model,
            "prompt": prompt,
            "stream": False,
            "options": {
                "num_predict": max_tokens,
                "temperature": temperature,
            },
        }
        data = self._post_json("/api/generate", payload)
        text = data.get("response", "")
        usage = {
            "prompt_tokens": data.get("prompt_eval_count", 0),
            "completion_tokens": data.get("eval_count", 0),
            "total_tokens": data.get("prompt_eval_count", 0) + data.get("eval_count", 0),
        }
        return {"text": text, "usage": usage, "raw": data}

    # ------------------------------------------------------------------
    # Inference: chat
    # ------------------------------------------------------------------

    def chat(
        self,
        messages: list[dict[str, str]],
        *,
        model: str | None = None,
        max_tokens: int = 128,
        temperature: float = 0.0,
    ) -> dict[str, Any]:
        """Multi-turn chat completion.

        Wraps ``POST /api/chat``.

        Args:
            messages: Conversation as ``[{"role": ..., "content": ...}]``.
            model: Model override.
            max_tokens: Soft max token limit.
            temperature: Sampling temperature.

        Returns:
            Dict with ``"text"``, ``"usage"``, and ``"raw"`` keys.
        """
        if self._offline_mode:
            last = messages[-1].get("content", "") if messages else ""
            return {
                "text": f"[ollama-offline] {last[:40]}...",
                "usage": {},
                "raw": {},
            }

        payload: dict[str, Any] = {
            "model": model or self._model,
            "messages": messages,
            "stream": False,
            "options": {
                "num_predict": max_tokens,
                "temperature": temperature,
            },
        }
        data = self._post_json("/api/chat", payload)
        text = data.get("message", {}).get("content", "")
        usage = {
            "prompt_tokens": data.get("prompt_eval_count", 0),
            "completion_tokens": data.get("eval_count", 0),
        }
        return {"text": text, "usage": usage, "raw": data}

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

    def embed(
        self,
        text: str | list[str],
        *,
        model: str | None = None,
    ) -> list[list[float]]:
        """Generate embeddings.

        Wraps ``POST /api/embeddings`` (single input) or ``POST /api/embed``
        (batch, Ollama >= 0.1.35).

        Args:
            text: String or list of strings.
            model: Model override.

        Returns:
            List of embedding vectors.
        """
        if isinstance(text, str):
            text = [text]
        if self._offline_mode:
            return [[0.0] * 4 for _ in text]

        results: list[list[float]] = []
        resolved = model or self._model
        for t in text:
            payload = {"model": resolved, "prompt": t}
            data = self._post_json("/api/embeddings", payload)
            results.append(data.get("embedding", []))
        return results

    # ------------------------------------------------------------------
    # HTTP helpers
    # ------------------------------------------------------------------

    def _post_json(self, path: str, payload: dict[str, Any]) -> dict[str, Any]:
        url = f"{self._base_url}{path}"
        req = urllib_request.Request(
            url=url,
            data=json.dumps(payload).encode("utf-8"),
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        try:
            with urllib_request.urlopen(req, timeout=self._timeout_s) as resp:  # noqa: S310
                return json.loads(resp.read().decode("utf-8"))
        except urllib_error.HTTPError as exc:
            body = exc.read().decode("utf-8", errors="replace")
            raise RuntimeError(f"Ollama server returned HTTP {exc.code}: {body}") from exc
        except urllib_error.URLError as exc:
            raise ConnectionError(
                f"Cannot connect to Ollama at {self._base_url}: {exc.reason}"
            ) from exc

    def _get(self, path: str) -> str:
        url = f"{self._base_url}{path}"
        req = urllib_request.Request(url=url, method="GET")
        with urllib_request.urlopen(req, timeout=self._timeout_s) as resp:  # noqa: S310
            return resp.read().decode("utf-8")

    def _get_json(self, path: str) -> dict[str, Any]:
        body = self._get(path)
        return json.loads(body)

    def __repr__(self) -> str:
        mode = " offline=True" if self._offline_mode else ""
        return f"OllamaClient({self._base_url!r}, model={self._model!r}{mode})"
