"""Topology mapper -- discovers, probes, and tracks network topology.

Coordinates the health monitor, latency tracker, and topology store to
maintain a live picture of the network.  Runs a background loop to
periodically ping idle backends and update the topology state.
"""

from __future__ import annotations

import asyncio
import contextlib
import json
import logging
import time
from collections import deque
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING, Any

import httpx

from llmhosts.topology.models import BackendNode, BackendStatus, TopologySnapshot, classify_latency

if TYPE_CHECKING:
    from pathlib import Path

logger = logging.getLogger(__name__)

# Rolling window size for latency samples per backend
_WINDOW_SIZE = 100

# Health check interval for idle backends (seconds)
_HEALTH_CHECK_INTERVAL = 60.0

# Probe timeout for individual health checks
_PROBE_TIMEOUT = 5.0


class _TrackedNode:
    """Internal mutable state for a tracked backend node."""

    __slots__ = (
        "backend_id",
        "backend_type",
        "consecutive_failures",
        "is_local",
        "last_healthy",
        "latency_samples",
        "models",
        "name",
        "total_checks",
        "total_failures",
        "url",
    )

    def __init__(
        self,
        backend_id: str,
        backend_type: str,
        name: str,
        url: str,
        is_local: bool = False,
        models: list[str] | None = None,
    ) -> None:
        self.backend_id = backend_id
        self.backend_type = backend_type
        self.name = name
        self.url = url
        self.is_local = is_local
        self.models = models or []
        self.latency_samples: deque[float] = deque(maxlen=_WINDOW_SIZE)
        self.consecutive_failures: int = 0
        self.last_healthy: datetime | None = None
        self.total_checks: int = 0
        self.total_failures: int = 0


def _percentile(sorted_data: list[float], pct: float) -> float:
    """Compute a percentile from pre-sorted data."""
    if not sorted_data:
        return 0.0
    n = len(sorted_data)
    if n == 1:
        return sorted_data[0]
    rank = (pct / 100.0) * (n - 1)
    lower = int(rank)
    upper = lower + 1
    fraction = rank - lower
    if upper >= n:
        return sorted_data[-1]
    return sorted_data[lower] + fraction * (sorted_data[upper] - sorted_data[lower])


class TopologyMapper:
    """Network topology mapper with latency tracking and health probing.

    Usage::

        mapper = TopologyMapper()
        await mapper.start()

        mapper.register_backend("ollama-local", "ollama", "Local Engine",
                                "http://127.0.0.1:11434", is_local=True)

        # Record latency from actual requests
        mapper.record_latency("ollama-local", 42.5)

        # Get current topology
        snapshot = mapper.get_snapshot()

        await mapper.stop()
    """

    def __init__(self, store_path: Path | None = None) -> None:
        self._nodes: dict[str, _TrackedNode] = {}
        self._running = False
        self._task: asyncio.Task[None] | None = None
        self._client: httpx.AsyncClient | None = None
        self._store_path = store_path
        self._store: Any | None = None  # TopologyStore, lazily imported

    # ------------------------------------------------------------------
    # Lifecycle
    # ------------------------------------------------------------------

    async def start(self) -> None:
        """Start the background health check loop."""
        if self._running:
            return

        self._client = httpx.AsyncClient(timeout=httpx.Timeout(_PROBE_TIMEOUT))

        # Initialize store if path provided
        if self._store_path is not None:
            from llmhosts.topology.store import TopologyStore

            self._store = TopologyStore(self._store_path)
            await self._store.initialize()

        self._running = True
        self._task = asyncio.create_task(self._probe_loop(), name="topology-mapper")
        logger.info("TopologyMapper started (interval=%.0fs, nodes=%d)", _HEALTH_CHECK_INTERVAL, len(self._nodes))

    async def stop(self) -> None:
        """Stop the background loop and clean up."""
        self._running = False

        if self._task is not None:
            self._task.cancel()
            with contextlib.suppress(asyncio.CancelledError):
                await self._task
            self._task = None

        if self._client is not None:
            await self._client.aclose()
            self._client = None

        if self._store is not None:
            await self._store.close()
            self._store = None

        logger.info("TopologyMapper stopped")

    # ------------------------------------------------------------------
    # Backend registration
    # ------------------------------------------------------------------

    def register_backend(
        self,
        backend_id: str,
        backend_type: str,
        name: str,
        url: str,
        is_local: bool = False,
        models: list[str] | None = None,
    ) -> None:
        """Register a backend for topology tracking."""
        self._nodes[backend_id] = _TrackedNode(
            backend_id=backend_id,
            backend_type=backend_type,
            name=name,
            url=url,
            is_local=is_local,
            models=models,
        )
        logger.info("Topology: registered %s (%s) at %s", backend_id, backend_type, url)

    def unregister_backend(self, backend_id: str) -> None:
        """Remove a backend from topology tracking."""
        removed = self._nodes.pop(backend_id, None)
        if removed:
            logger.info("Topology: unregistered %s", backend_id)

    # ------------------------------------------------------------------
    # Latency recording (called from dispatcher on real requests)
    # ------------------------------------------------------------------

    def record_latency(self, backend_id: str, latency_ms: float, healthy: bool = True) -> None:
        """Record a latency observation from an actual request."""
        node = self._nodes.get(backend_id)
        if node is None:
            return

        if latency_ms >= 0:
            node.latency_samples.append(latency_ms)

        node.total_checks += 1
        if healthy:
            node.consecutive_failures = 0
            node.last_healthy = datetime.now(tz=timezone.utc)
        else:
            node.consecutive_failures += 1
            node.total_failures += 1

    # ------------------------------------------------------------------
    # Topology queries
    # ------------------------------------------------------------------

    def get_node(self, backend_id: str) -> BackendNode | None:
        """Get topology info for a single backend."""
        tracked = self._nodes.get(backend_id)
        if tracked is None:
            return None
        return self._build_node(tracked)

    def get_snapshot(self) -> TopologySnapshot:
        """Get the current topology snapshot with all backends."""
        nodes = [self._build_node(t) for t in self._nodes.values()]
        healthy = sum(1 for n in nodes if n.status == BackendStatus.online)
        latencies = [n.latency_ms for n in nodes if n.status == BackendStatus.online and n.latency_ms > 0]
        avg_latency = sum(latencies) / len(latencies) if latencies else 0.0

        return TopologySnapshot(
            nodes=nodes,
            total_backends=len(nodes),
            healthy_backends=healthy,
            avg_latency_ms=round(avg_latency, 2),
        )

    def get_fastest_backend(self, backend_ids: list[str] | None = None) -> str | None:
        """Return the backend_id with lowest average latency among healthy backends."""
        candidates_iter = self._nodes.values()
        candidates = (
            list(candidates_iter)
            if backend_ids is None
            else [n for n in candidates_iter if n.backend_id in backend_ids]
        )

        best_id: str | None = None
        best_avg = float("inf")
        for node in candidates:
            if node.consecutive_failures >= 3:
                continue
            if not node.latency_samples:
                continue
            avg = sum(node.latency_samples) / len(node.latency_samples)
            if avg < best_avg:
                best_avg = avg
                best_id = node.backend_id
        return best_id

    def get_latency_weight(self, backend_id: str) -> float:
        """Return a 0-1 weight where lower latency = higher weight.

        Used by the router to factor latency into routing decisions.
        Returns 0.5 if no data is available.
        """
        node = self._nodes.get(backend_id)
        if node is None or not node.latency_samples:
            return 0.5

        avg = sum(node.latency_samples) / len(node.latency_samples)
        # Normalize: <50ms -> 1.0, >1000ms -> 0.0
        weight = max(0.0, min(1.0, 1.0 - (avg - 50) / 950))
        return round(weight, 3)

    # ------------------------------------------------------------------
    # Manual ping (triggered by API)
    # ------------------------------------------------------------------

    async def ping_all(self) -> TopologySnapshot:
        """Manually trigger a health check on all backends and return updated snapshot."""
        await self._probe_all()
        return self.get_snapshot()

    async def ping_backend(self, backend_id: str) -> BackendNode | None:
        """Manually ping a single backend."""
        node = self._nodes.get(backend_id)
        if node is None:
            return None
        await self._probe_node(node)
        return self._build_node(node)

    # ------------------------------------------------------------------
    # Internal: probe loop
    # ------------------------------------------------------------------

    async def _probe_loop(self) -> None:
        """Background loop that probes backends at regular intervals."""
        try:
            while self._running:
                await self._probe_all()
                await asyncio.sleep(_HEALTH_CHECK_INTERVAL)
        except asyncio.CancelledError:
            logger.debug("Topology probe loop cancelled")
            raise

    async def _probe_all(self) -> None:
        """Probe all registered backends concurrently."""
        if not self._nodes:
            return
        tasks = [self._probe_node(node) for node in self._nodes.values()]
        await asyncio.gather(*tasks, return_exceptions=True)

        # Persist to store if available
        if self._store is not None:
            await self._persist_snapshot()

    async def _probe_node(self, node: _TrackedNode) -> None:
        """Probe a single backend and record the result."""
        if self._client is None:
            return

        start = time.perf_counter()
        healthy = False

        try:
            if node.backend_type == "ollama":
                resp = await self._client.get(f"{node.url.rstrip('/')}/")
                healthy = resp.status_code == 200
            elif node.backend_type in ("openai", "openrouter"):
                resp = await self._client.get(f"{node.url.rstrip('/')}/models")
                healthy = resp.status_code in (200, 401, 403)
            elif node.backend_type == "anthropic":
                resp = await self._client.post(
                    f"{node.url.rstrip('/')}/messages",
                    json={"model": "claude-3-haiku-20240307", "messages": [], "max_tokens": 1},
                    headers={"x-api-key": "probe", "anthropic-version": "2023-06-01"},
                )
                healthy = resp.status_code in (200, 400, 401, 403)
            else:
                resp = await self._client.get(node.url.rstrip("/"))
                healthy = resp.status_code < 500
        except (httpx.TimeoutException, httpx.ConnectError):
            healthy = False
        except Exception as exc:
            logger.debug("Probe error for %s: %s", node.backend_id, exc)
            healthy = False

        latency_ms = (time.perf_counter() - start) * 1000
        self.record_latency(node.backend_id, latency_ms, healthy)

        # Also persist individual sample to store
        if self._store is not None:
            try:
                await self._store.record_sample(node.backend_id, latency_ms, healthy)
            except Exception as exc:
                logger.debug("Failed to persist sample for %s: %s", node.backend_id, exc)

    # ------------------------------------------------------------------
    # Internal: build node and persist
    # ------------------------------------------------------------------

    def _build_node(self, tracked: _TrackedNode) -> BackendNode:
        """Build a BackendNode from internal tracking state."""
        samples = list(tracked.latency_samples)
        avg = sum(samples) / len(samples) if samples else 0.0
        sorted_samples = sorted(samples) if samples else []

        p50 = _percentile(sorted_samples, 50)
        p95 = _percentile(sorted_samples, 95)
        p99 = _percentile(sorted_samples, 99)

        # Determine status
        if tracked.consecutive_failures >= 3:
            status = BackendStatus.offline
        elif tracked.consecutive_failures >= 1:
            status = BackendStatus.degraded
        elif tracked.last_healthy is not None:
            status = BackendStatus.online
        else:
            status = BackendStatus.offline

        # Compute uptime
        total = tracked.total_checks
        uptime = ((total - tracked.total_failures) / total * 100.0) if total > 0 else 100.0

        # Compute trend
        trend = self._compute_trend(samples)

        is_online = status in (BackendStatus.online, BackendStatus.degraded)
        bucket = classify_latency(avg, is_online)

        return BackendNode(
            backend_id=tracked.backend_id,
            backend_type=tracked.backend_type,
            name=tracked.name,
            url=tracked.url,
            is_local=tracked.is_local,
            status=status,
            latency_ms=round(avg, 2),
            latency_bucket=bucket,
            p50_ms=round(p50, 2),
            p95_ms=round(p95, 2),
            p99_ms=round(p99, 2),
            models=tracked.models,
            last_seen=tracked.last_healthy,
            last_check=datetime.now(tz=timezone.utc),
            consecutive_failures=tracked.consecutive_failures,
            uptime_percent=round(uptime, 2),
            trend=trend,
        )

    @staticmethod
    def _compute_trend(samples: list[float]) -> str:
        """Compute latency trend from time-ordered samples."""
        if len(samples) < 10:
            return "stable"
        third = len(samples) // 3
        if third == 0:
            return "stable"
        old_avg = sum(samples[:third]) / third
        new_avg = sum(samples[-third:]) / third
        if old_avg == 0:
            return "stable"
        change = (new_avg - old_avg) / old_avg
        if change > 0.15:
            return "degrading"
        if change < -0.15:
            return "improving"
        return "stable"

    async def _persist_snapshot(self) -> None:
        """Persist current topology state to the store."""
        if self._store is None:
            return
        for tracked in self._nodes.values():
            node = self._build_node(tracked)
            try:
                await self._store.upsert_node(
                    backend_id=node.backend_id,
                    backend_type=node.backend_type,
                    name=node.name,
                    url=node.url,
                    is_local=node.is_local,
                    status=node.status.value,
                    avg_latency_ms=node.latency_ms,
                    p50_ms=node.p50_ms,
                    p95_ms=node.p95_ms,
                    p99_ms=node.p99_ms,
                    models_json=json.dumps(node.models),
                    last_seen=node.last_seen.isoformat() if node.last_seen else None,
                    last_check=node.last_check.isoformat() if node.last_check else None,
                    consecutive_failures=node.consecutive_failures,
                    uptime_percent=node.uptime_percent,
                    trend=node.trend,
                )
            except Exception as exc:
                logger.debug("Failed to persist node %s: %s", node.backend_id, exc)

        # Prune old samples (keep last 24 hours)
        cutoff = (datetime.now(tz=timezone.utc) - timedelta(hours=24)).isoformat()
        try:
            await self._store.prune_old_samples(cutoff)
        except Exception as exc:
            logger.debug("Failed to prune old samples: %s", exc)
