"""Device manager -- registration, discovery, persistence, and health checks."""

from __future__ import annotations

import contextlib
import json
import logging
import platform
import time
import uuid
from datetime import datetime, timezone
from pathlib import Path  # noqa: TC003 -- used at runtime in __init__

import httpx

from llmhosts.mesh.models import Device, DeviceRegistration, MeshStatus

logger = logging.getLogger(__name__)

# Filename for persisted device list
_DEVICES_FILENAME = "devices.json"


class DeviceManager:
    """Manages device registration, discovery, and health.

    Stores device list in ``<data_dir>/devices.json``.
    Max 5 devices on Pro tier (configurable).
    """

    MAX_DEVICES_FREE = 1
    MAX_DEVICES_PRO = 5

    def __init__(self, data_dir: Path, *, max_devices: int | None = None) -> None:
        self._data_dir = data_dir
        self._devices_file = data_dir / _DEVICES_FILENAME
        self._devices: list[Device] = []
        self._max_devices = max_devices if max_devices is not None else self.MAX_DEVICES_PRO

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

    async def initialize(self) -> None:
        """Load device list from disk."""
        self._data_dir.mkdir(parents=True, exist_ok=True)
        if self._devices_file.is_file():
            try:
                raw = self._devices_file.read_text(encoding="utf-8")
                data = json.loads(raw)
                self._devices = [Device.model_validate(d) for d in data]
                logger.info("Loaded %d device(s) from %s", len(self._devices), self._devices_file)
            except (json.JSONDecodeError, ValueError) as exc:
                logger.warning("Failed to load devices from %s: %s -- starting fresh", self._devices_file, exc)
                self._devices = []
        else:
            logger.debug("No devices file at %s -- starting with empty mesh", self._devices_file)
            self._devices = []

    # ------------------------------------------------------------------
    # Registration
    # ------------------------------------------------------------------

    async def register(self, registration: DeviceRegistration) -> Device:
        """Register a new device in the mesh.

        1. Check max device limit
        2. Generate device ID
        3. Detect hardware if local (no tunnel_url)
        4. Discover Ollama models if local
        5. Save to devices.json

        Raises
        ------
        ValueError
            If the device limit is reached or name already exists.
        """
        # Check for duplicate name
        for d in self._devices:
            if d.name == registration.name:
                raise ValueError(f"Device with name '{registration.name}' already exists")

        # Check device limit
        if len(self._devices) >= self._max_devices:
            raise ValueError(
                f"Device limit reached ({self._max_devices}). Remove a device first or upgrade to Pro for more."
            )

        is_local = registration.tunnel_url is None
        hardware_summary = ""
        models: list[str] = []

        if is_local:
            hardware_summary = await _detect_hardware_summary()
            models = await _discover_local_models()

        device = Device(
            id=str(uuid.uuid4()),
            name=registration.name,
            hostname=platform.node(),
            role=registration.role,
            tunnel_url=registration.tunnel_url,
            hardware_summary=hardware_summary,
            models=models,
            last_seen=datetime.now(tz=timezone.utc),
            is_local=is_local,
            healthy=True,
            latency_ms=0.0 if is_local else None,
        )

        self._devices.append(device)
        await self.save()
        logger.info("Registered device '%s' (id=%s, role=%s)", device.name, device.id, device.role.value)
        return device

    # ------------------------------------------------------------------
    # Removal
    # ------------------------------------------------------------------

    async def remove(self, device_id: str) -> bool:
        """Remove a device from the mesh by ID or name.

        Returns True if a device was removed, False if not found.
        """
        before = len(self._devices)
        self._devices = [d for d in self._devices if d.id != device_id and d.name != device_id]
        removed = len(self._devices) < before
        if removed:
            await self.save()
            logger.info("Removed device '%s'", device_id)
        return removed

    # ------------------------------------------------------------------
    # Queries
    # ------------------------------------------------------------------

    async def list_devices(self) -> list[Device]:
        """Return all registered devices."""
        return list(self._devices)

    async def get_device(self, device_id: str) -> Device | None:
        """Get a single device by ID or name."""
        for d in self._devices:
            if d.id == device_id or d.name == device_id:
                return d
        return None

    async def get_local_device(self) -> Device | None:
        """Return the local device entry, if any."""
        for d in self._devices:
            if d.is_local:
                return d
        return None

    # ------------------------------------------------------------------
    # Health
    # ------------------------------------------------------------------

    async def health_check(self) -> MeshStatus:
        """Check health of all devices.

        For remote devices: HTTP GET ``{tunnel_url}/health``
        For local: always healthy.
        Updates latency measurements.
        """
        local_device: Device | None = None
        remote_devices: list[Device] = []
        all_models: set[str] = set()
        total_vram_gb = 0.0
        all_healthy = True

        for device in self._devices:
            if device.is_local:
                device.healthy = True
                device.latency_ms = 0.0
                device.last_seen = datetime.now(tz=timezone.utc)
                local_device = device
            else:
                await self._check_remote_health(device)
                remote_devices.append(device)
                if not device.healthy:
                    all_healthy = False

            all_models.update(device.models)

            # Parse VRAM from hardware_summary (best effort)
            total_vram_gb += _parse_vram_from_summary(device.hardware_summary)

        await self.save()

        return MeshStatus(
            local_device=local_device,
            remote_devices=remote_devices,
            total_models=len(all_models),
            total_gpu_vram_gb=total_vram_gb,
            mesh_healthy=all_healthy,
        )

    async def _check_remote_health(self, device: Device) -> None:
        """Ping a remote device's /health endpoint."""
        if not device.tunnel_url:
            device.healthy = False
            return

        url = device.tunnel_url.rstrip("/") + "/health"
        try:
            start = time.monotonic()
            async with httpx.AsyncClient(timeout=10.0) as client:
                resp = await client.get(url)
            elapsed_ms = (time.monotonic() - start) * 1000

            device.healthy = resp.status_code == 200
            device.latency_ms = round(elapsed_ms, 1)
            device.last_seen = datetime.now(tz=timezone.utc)

            # Optionally update models list from health response
            if resp.status_code == 200:
                try:
                    data = resp.json()
                    if "models" in data and isinstance(data["models"], list):
                        device.models = data["models"]
                except Exception:
                    pass

            logger.debug("Health check %s: status=%d, latency=%.1fms", device.name, resp.status_code, elapsed_ms)
        except Exception as exc:
            device.healthy = False
            device.latency_ms = None
            logger.debug("Health check failed for %s: %s", device.name, exc)

    # ------------------------------------------------------------------
    # Model aggregation
    # ------------------------------------------------------------------

    async def get_all_models(self) -> dict[str, list[str]]:
        """Get model availability across all devices.

        Returns
        -------
        dict[str, list[str]]
            Mapping of device_id -> list of model names.
        """
        return {d.id: list(d.models) for d in self._devices if d.models}

    # ------------------------------------------------------------------
    # Persistence
    # ------------------------------------------------------------------

    async def save(self) -> None:
        """Persist device list to disk as JSON."""
        self._data_dir.mkdir(parents=True, exist_ok=True)
        data = [d.model_dump(mode="json") for d in self._devices]
        self._devices_file.write_text(json.dumps(data, indent=2, default=str), encoding="utf-8")
        logger.debug("Saved %d device(s) to %s", len(self._devices), self._devices_file)


# ---------------------------------------------------------------------------
# Private helpers
# ---------------------------------------------------------------------------


async def _detect_hardware_summary() -> str:
    """Best-effort one-line hardware summary for the local machine."""
    try:
        from llmhosts.discovery.hardware import HardwareDetector

        hw = await HardwareDetector.detect()
        parts: list[str] = []
        if hw.gpus:
            for gpu in hw.gpus:
                vram_gb = round(gpu.vram_total_mb / 1024, 1)
                parts.append(f"{gpu.name} ({vram_gb}GB)")
        parts.append(f"{hw.ram_total_gb:.0f}GB RAM")
        return ", ".join(parts)
    except Exception as exc:
        logger.debug("Hardware summary detection failed: %s", exc)
        return ""


async def _discover_local_models() -> list[str]:
    """Discover locally-installed Ollama models."""
    try:
        from llmhosts.config import load_config
        from llmhosts.discovery.ollama import OllamaDiscovery

        config = load_config()
        async with OllamaDiscovery(host=config.ollama.host, timeout=5.0) as ollama:
            if await ollama.is_available():
                models = await ollama.list_models()
                return [m.name for m in models]
    except Exception as exc:
        logger.debug("Ollama model discovery failed: %s", exc)
    return []


def _parse_vram_from_summary(summary: str) -> float:
    """Extract total VRAM in GB from a hardware summary string.

    E.g. ``"RTX 4090 (24.0GB), 64GB RAM"`` -> ``24.0``
    """
    import re

    total = 0.0
    # Match patterns like "(24GB)" or "(24.0GB)"
    for match in re.finditer(r"\(([\d.]+)GB\)", summary):
        with contextlib.suppress(ValueError):
            total += float(match.group(1))
    return total
