"""Ollama auto-discovery -- detect, query, and interact with a local Ollama instance."""

from __future__ import annotations

import logging
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from collections.abc import AsyncIterator

import httpx

from llmhosts.constants import OLLAMA_DEFAULT_HOST
from llmhosts.discovery.models import ModelInfo, OllamaHealth, OllamaModel, RunningModel

logger = logging.getLogger(__name__)


class OllamaDiscovery:
    """Auto-discovers and manages the connection to an Ollama instance.

    Usage::

        async with OllamaDiscovery() as ollama:
            if await ollama.is_available():
                models = await ollama.list_models()
    """

    def __init__(self, host: str = OLLAMA_DEFAULT_HOST, timeout: float = 5.0) -> None:
        self._host = host.rstrip("/")
        self._timeout = timeout
        self._client = httpx.AsyncClient(base_url=self._host, timeout=httpx.Timeout(timeout))

    # -- context manager -------------------------------------------------------

    async def __aenter__(self) -> OllamaDiscovery:
        return self

    async def __aexit__(self, exc_type: type | None, exc_val: BaseException | None, exc_tb: Any) -> None:
        await self.close()

    # -- public API ------------------------------------------------------------

    async def is_available(self) -> bool:
        """Return *True* if Ollama is running and reachable (``GET /``)."""
        try:
            resp = await self._client.get("/")
            return resp.status_code == 200
        except (httpx.ConnectError, httpx.TimeoutException, httpx.HTTPError) as exc:
            logger.debug("Ollama not reachable at %s: %s", self._host, exc)
            return False

    async def list_models(self) -> list[OllamaModel]:
        """``GET /api/tags`` -- return all installed models with metadata."""
        try:
            resp = await self._client.get("/api/tags")
            resp.raise_for_status()
            data = resp.json()
            return [OllamaModel.model_validate(m) for m in data.get("models", [])]
        except (httpx.HTTPError, httpx.TimeoutException) as exc:
            logger.error("Failed to list Ollama models: %s", exc)
            return []

    async def list_running(self) -> list[RunningModel]:
        """``GET /api/ps`` -- return currently loaded / running models."""
        try:
            resp = await self._client.get("/api/ps")
            resp.raise_for_status()
            data = resp.json()
            return [RunningModel.model_validate(m) for m in data.get("models", [])]
        except (httpx.HTTPError, httpx.TimeoutException) as exc:
            logger.error("Failed to list running Ollama models: %s", exc)
            return []

    async def show_model(self, name: str) -> ModelInfo:
        """``POST /api/show`` -- get detailed info for a specific model.

        Raises :class:`httpx.HTTPStatusError` when the model does not exist.
        """
        resp = await self._client.post("/api/show", json={"name": name})
        resp.raise_for_status()
        return ModelInfo.model_validate(resp.json())

    async def chat(
        self,
        model: str,
        messages: list[dict[str, Any]],
        stream: bool = False,
        **kwargs: Any,
    ) -> dict[str, Any] | AsyncIterator[dict[str, Any]]:
        """``POST /api/chat`` -- send a chat request.

        Parameters
        ----------
        model:
            Model name (e.g. ``"llama3.2:8b"``).
        messages:
            List of ``{"role": ..., "content": ...}`` dicts.
        stream:
            When *True*, return an :class:`AsyncIterator` that yields parsed
            JSON chunks as they arrive.
        **kwargs:
            Additional fields forwarded in the request body (``temperature``,
            ``top_p``, ``num_predict``, etc.).
        """
        payload: dict[str, Any] = {"model": model, "messages": messages, "stream": stream, **kwargs}

        if stream:
            return self._stream_chat(payload)

        resp = await self._client.post("/api/chat", json=payload)
        resp.raise_for_status()
        return resp.json()  # type: ignore[return-value,no-any-return]

    async def health_check(self) -> OllamaHealth:
        """Full health check: availability, model list, and running models."""
        available = await self.is_available()
        models: list[OllamaModel] = []
        running: list[RunningModel] = []

        if available:
            models = await self.list_models()
            running = await self.list_running()

        return OllamaHealth(
            available=available,
            host=self._host,
            model_count=len(models),
            running_count=len(running),
            models=models,
            running=running,
            checked_at=datetime.now(tz=timezone.utc),
        )

    async def close(self) -> None:
        """Close the underlying httpx client."""
        await self._client.aclose()

    # -- internal helpers ------------------------------------------------------

    async def _stream_chat(self, payload: dict[str, Any]) -> AsyncIterator[dict[str, Any]]:
        """Yield parsed JSON chunks from a streaming ``/api/chat`` response.

        Ollama streams newline-delimited JSON objects.  We use httpx's
        ``aiter_lines`` to handle back-pressure properly.
        """
        async with self._client.stream("POST", "/api/chat", json=payload) as resp:
            resp.raise_for_status()
            async for line in resp.aiter_lines():
                line = line.strip()
                if not line:
                    continue
                # Each line is a complete JSON object
                import json

                try:
                    chunk = json.loads(line)
                    yield chunk
                except json.JSONDecodeError:
                    logger.warning("Skipping malformed streaming chunk: %s", line[:200])

    # -- repr ------------------------------------------------------------------

    def __repr__(self) -> str:
        return f"OllamaDiscovery(host={self._host!r})"
