"""Health of LLM targets, and the background probe that switches back.

The old engine had a one-way door: a dead server got blacklisted for a fixed
TTL and nothing ever checked whether it came back — on production the TTL was
never even populated, because the blacklists were disabled for single-server
setups. Recovery was therefore a restart.

Here the engine records *why* an entry failed (network, quota, auth, 5xx,
missing model) and when it may be retried. A single background task probes the
degraded targets once per interval with a short timeout; when the probe says a
target is healthy again the next user turn starts on the primary entry — the
probe never interrupts a running turn.

The probe costs nothing while everything is healthy: it only looks at targets
that are not ``ok``.
"""

from __future__ import annotations

import asyncio
import time
from dataclasses import dataclass

import ollama as ollama_client
import structlog

from navi.exceptions import (
    LLMAuthError,
    LLMConnectionError,
    LLMModelNotFoundError,
    LLMQuotaExceededError,
    LLMServerUnavailableError,
)

from .config import LLMConfig, RouteEntry, Target

log = structlog.get_logger()

# An auth failure is not transient, but the key can be rotated in the
# environment without a restart — so re-check it, just not every minute.
_AUTH_RECHECK_SEC = 300.0


@dataclass
class TargetHealth:
    state: str = "ok"          # ok | down | cooling | auth
    reason: str = ""           # network | quota | server | auth
    since: float = 0.0         # monotonic, when it went non-ok
    next_probe: float = 0.0    # monotonic, earliest retry
    detail: str = ""


@dataclass
class ModelHealth:
    missing: bool = False
    next_probe: float = 0.0
    detail: str = ""


class RouteHealth:
    """Per-target and per-(target, model) state, plus the probe task."""

    def __init__(self, config: LLMConfig) -> None:
        self._config = config
        self._targets: dict[str, TargetHealth] = {
            name: TargetHealth() for name in config.targets
        }
        self._models: dict[tuple[str, str], ModelHealth] = {}
        self._clients: dict[str, ollama_client.AsyncClient] = {}
        self._task: asyncio.Task | None = None
        # Called with (target_name, state) after a probe changes the state —
        # lets the container log a recovery without importing logging here.
        self.on_change = None

    # ── queries ──────────────────────────────────────────────────────────

    def _target_state(self, target: Target) -> TargetHealth:
        return self._targets.setdefault(target.name, TargetHealth())

    def target_usable(self, target: Target, now: float | None = None) -> bool:
        """May a request go to this target right now?

        A target that is cooling is tried again once its earliest retry time
        has passed even if the background probe has not run yet: refusing a
        request that could have worked is worse than one extra failed attempt.
        """
        h = self._target_state(target)
        if h.state == "ok":
            return True
        return (time.monotonic() if now is None else now) >= h.next_probe

    def entry_usable(self, entry: RouteEntry, now: float | None = None) -> bool:
        if not self.target_usable(entry.target, now):
            return False
        m = self._models.get((entry.target.name, entry.model))
        if m is None:
            return True
        if not m.missing:
            return True
        return (time.monotonic() if now is None else now) >= m.next_probe

    # ── recorded by the engine ───────────────────────────────────────────

    def reset(self) -> None:
        """Forget every recorded failure (admin action after a config fix)."""
        for h in self._targets.values():
            h.state, h.reason, h.detail = "ok", "", ""
            h.since = h.next_probe = 0.0
        self._models.clear()
        log.info("llm.health_reset")

    def report_success(self, entry: RouteEntry) -> None:
        h = self._target_state(entry.target)
        was = h.state
        h.state, h.reason, h.detail, h.next_probe = "ok", "", "", 0.0
        self._models.pop((entry.target.name, entry.model), None)
        if was != "ok":
            log.info("llm.target_recovered", target=entry.target.name, previous=was)
            self._notify(entry.target.name, "ok")

    def report_failure(self, entry: RouteEntry, exc: Exception) -> None:
        now = time.monotonic()
        interval = self._config.probe_interval_sec

        if isinstance(exc, LLMModelNotFoundError):
            self._models[(entry.target.name, entry.model)] = ModelHealth(
                missing=True, next_probe=now + interval, detail=str(exc)
            )
            log.warning(
                "llm.model_missing", target=entry.target.name, model=entry.model
            )
            return

        if isinstance(exc, LLMQuotaExceededError):
            delay = max(getattr(exc, "retry_after", None) or 0.0, interval)
            state, reason = "cooling", "quota"
        elif isinstance(exc, LLMAuthError):
            delay, state, reason = _AUTH_RECHECK_SEC, "auth", "auth"
        elif isinstance(exc, LLMServerUnavailableError):
            delay, state, reason = interval, "cooling", "server"
        elif isinstance(exc, LLMConnectionError):
            delay, state, reason = interval, "down", "network"
        else:
            # Unknown failures are not worth blacklisting a whole host for.
            return

        h = self._target_state(entry.target)
        changed = h.state != state
        h.state, h.reason, h.detail = state, reason, str(exc)
        h.next_probe = now + delay
        if changed:
            h.since = now
            log.warning(
                "llm.target_degraded",
                target=entry.target.name,
                state=state,
                reason=reason,
                detail=str(exc)[:300],
            )
            self._notify(entry.target.name, state)

    def _notify(self, target: str, state: str) -> None:
        if self.on_change is None:
            return
        try:
            self.on_change(target, state)
        except Exception:  # a listener must never break the LLM path
            log.debug("llm.health.listener_failed", exc_info=True)

    # ── snapshot for /health ─────────────────────────────────────────────

    def snapshot(self) -> dict:
        now = time.monotonic()
        targets = {}
        for name, h in self._targets.items():
            entry = {"state": h.state}
            if h.state != "ok":
                entry["reason"] = h.reason
                entry["retry_in_sec"] = max(0, round(h.next_probe - now))
                if h.detail:
                    entry["detail"] = h.detail[:200]
            targets[name] = entry
        missing = [
            {"target": target, "model": model, "detail": state.detail[:200]}
            for (target, model), state in self._models.items()
            if state.missing
        ]
        return {"targets": targets, "missing_models": missing}

    # ── background probe ─────────────────────────────────────────────────

    def start(self) -> None:
        if self._task is not None and not self._task.done():
            return
        self._task = asyncio.create_task(self._loop(), name="llm-health-probe")
        log.info(
            "llm.health_probe_started",
            interval_sec=self._config.probe_interval_sec,
            targets=list(self._config.targets),
        )

    async def stop(self) -> None:
        if self._task is None:
            return
        self._task.cancel()
        try:
            await self._task
        except (asyncio.CancelledError, Exception):
            pass
        self._task = None
        for client in self._clients.values():
            try:
                await client._client.aclose()
            except Exception:
                pass
        self._clients.clear()

    async def _loop(self) -> None:
        interval = self._config.probe_interval_sec
        while True:
            try:
                await asyncio.sleep(interval)
                await self.probe_once()
            except asyncio.CancelledError:
                raise
            except Exception:
                log.warning("llm.health_probe_failed", exc_info=True)

    async def probe_once(self) -> None:
        """Probe every target that is not healthy (and every missing model)."""
        now = time.monotonic()
        due = [
            target
            for name, target in self._config.targets.items()
            if (h := self._targets.get(name)) is not None and h.state != "ok"
            and now >= h.next_probe
        ]
        missing_targets = {
            t for (t, _), m in self._models.items() if m.missing and now >= m.next_probe
        }
        for name in missing_targets:
            target = self._config.targets.get(name)
            if target is not None and target not in due:
                due.append(target)
        if not due:
            return
        await asyncio.gather(*(self._probe(target) for target in due))

    async def _probe(self, target: Target) -> None:
        timeout = self._config.probe_timeout_sec
        client = self._client(target)
        h = self._target_state(target)

        try:
            await asyncio.wait_for(client.list(), timeout=timeout)
        except TimeoutError:
            self._mark(target, "down", "network", f"probe timed out after {timeout}s")
            return
        except Exception as e:
            state, reason = _classify_probe_error(e)
            self._mark(target, state, reason, str(e))
            return

        # Reachable. Now the model-level questions.
        models = {
            entry.model
            for route in self._config.routes.values()
            for entry in route.entries
            if entry.target.name == target.name
        }
        for model in models:
            key = (target.name, model)
            current = self._models.get(key)
            if current is None or not current.missing:
                continue
            try:
                await asyncio.wait_for(client.show(model), timeout=timeout)
            except Exception as e:
                current.next_probe = time.monotonic() + self._config.probe_interval_sec
                current.detail = str(e)
                continue
            self._models.pop(key, None)
            log.info("llm.model_recovered", target=target.name, model=model)

        # A quota failure can only be checked by asking for a token: `list()`
        # still answers while generation is blocked.
        if h.state == "cooling" and h.reason == "quota":
            entry = next(
                (
                    e
                    for route in self._config.routes.values()
                    for e in route.entries
                    if e.target.name == target.name
                ),
                None,
            )
            if entry is not None:
                try:
                    await asyncio.wait_for(
                        client.chat(
                            model=entry.model,
                            messages=[{"role": "user", "content": "ok"}],
                            options={"num_predict": 1},
                            stream=False,
                        ),
                        timeout=timeout,
                    )
                except Exception as e:
                    self._mark(target, "cooling", "quota", str(e))
                    return

        h.state, h.reason, h.detail, h.next_probe = "ok", "", "", 0.0
        h.since = 0.0
        log.info("llm.target_recovered", target=target.name, previous="probe")
        self._notify(target.name, "ok")

    def _mark(self, target: Target, state: str, reason: str, detail: str) -> None:
        h = self._target_state(target)
        changed = h.state != state
        h.state, h.reason, h.detail = state, reason, detail
        h.next_probe = time.monotonic() + self._config.probe_interval_sec
        if changed:
            log.warning(
                "llm.target_degraded", target=target.name, state=state,
                reason=reason, detail=detail[:300],
            )
            self._notify(target.name, state)

    def _client(self, target: Target) -> ollama_client.AsyncClient:
        client = self._clients.get(target.name)
        if client is None:
            headers = (
                {"Authorization": f"Bearer {target.api_key}"} if target.api_key else None
            )
            client = ollama_client.AsyncClient(
                host=target.host, headers=headers, timeout=self._config.probe_timeout_sec
            )
            self._clients[target.name] = client
        return client


def _classify_probe_error(e: Exception) -> tuple[str, str]:
    status = getattr(e, "status_code", None)
    if status in (401, 403):
        return "auth", "auth"
    if status == 429:
        return "cooling", "quota"
    if isinstance(status, int) and 500 <= status < 600:
        return "cooling", "server"
    if isinstance(e, (ollama_client.RequestError, TimeoutError)):
        return "down", "network"
    return "down", "network"