"""mDNS-based LAN discovery for LLMHosts instances.

Uses Zeroconf (multicast DNS / DNS-SD) to broadcast this node's presence
and discover other LLMHosts instances on the local network.

Service type: ``_llmhosts._tcp.local.``

Each service TXT record carries:
    - ``version``: llmhost version
    - ``port``: API port
    - ``gpu``: GPU summary (e.g. "RTX 4090 24GB")
    - ``models``: comma-separated model list
    - ``hostname``: human-friendly hostname

The module is guarded behind a try/except so that ``zeroconf`` is only
required when LAN discovery is explicitly used (optional ``[lan]`` tier).
"""

from __future__ import annotations

import asyncio
import contextlib
import logging
import socket
import time
from dataclasses import dataclass, field
from typing import Any

from pydantic import BaseModel, Field

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

SERVICE_TYPE = "_llmhosts._tcp.local."
SERVICE_NAME_PREFIX = "LLMHosts-"

# ---------------------------------------------------------------------------
# Pydantic models for discovered nodes
# ---------------------------------------------------------------------------


class LANNode(BaseModel):
    """A discovered LLMHosts instance on the LAN."""

    hostname: str
    ip: str
    port: int
    version: str = ""
    gpu_summary: str = ""
    models: list[str] = Field(default_factory=list)
    latency_ms: float | None = None
    last_seen: float = 0.0  # Unix timestamp

    @property
    def base_url(self) -> str:
        """HTTP base URL for this node."""
        return f"http://{self.ip}:{self.port}"

    @property
    def is_stale(self) -> bool:
        """True if not seen in the last 60 seconds."""
        return (time.time() - self.last_seen) > 60.0


# ---------------------------------------------------------------------------
# Service listener (collects discovered nodes)
# ---------------------------------------------------------------------------


@dataclass
class _NodeRegistry:
    """Thread-safe in-memory registry of discovered LAN nodes."""

    nodes: dict[str, LANNode] = field(default_factory=dict)

    def upsert(self, name: str, node: LANNode) -> None:
        self.nodes[name] = node

    def remove(self, name: str) -> None:
        self.nodes.pop(name, None)

    def active_nodes(self) -> list[LANNode]:
        """Return non-stale nodes sorted by hostname."""
        return sorted(
            [n for n in self.nodes.values() if not n.is_stale],
            key=lambda n: n.hostname,
        )

    def all_nodes(self) -> list[LANNode]:
        """Return all nodes sorted by hostname."""
        return sorted(self.nodes.values(), key=lambda n: n.hostname)


# Global registry -- shared between browser and advertiser
_registry = _NodeRegistry()


def get_discovered_nodes() -> list[LANNode]:
    """Return currently discovered LAN nodes (non-stale)."""
    return _registry.active_nodes()


# ---------------------------------------------------------------------------
# Zeroconf service browser listener
# ---------------------------------------------------------------------------


class _LLMHostsListener:
    """Zeroconf ServiceListener that populates the node registry."""

    def __init__(self, zc: Any) -> None:
        self._zc = zc

    def add_service(self, zc: Any, type_: str, name: str) -> None:
        self._update_service(zc, type_, name)

    def update_service(self, zc: Any, type_: str, name: str) -> None:
        self._update_service(zc, type_, name)

    def remove_service(self, zc: Any, type_: str, name: str) -> None:
        logger.debug("LAN node removed: %s", name)
        _registry.remove(name)

    def _update_service(self, zc: Any, type_: str, name: str) -> None:
        info = zc.get_service_info(type_, name)
        if info is None:
            return

        addresses = info.parsed_addresses()
        if not addresses:
            return

        ip = addresses[0]
        port = info.port

        # Parse TXT record properties
        props: dict[str, str] = {}
        if info.properties:
            for k, v in info.properties.items():
                key = k.decode("utf-8") if isinstance(k, bytes) else str(k)
                val = v.decode("utf-8") if isinstance(v, bytes) else str(v)
                props[key] = val

        hostname = props.get("hostname", name.split(".")[0])
        models_str = props.get("models", "")
        models = [m.strip() for m in models_str.split(",") if m.strip()] if models_str else []

        node = LANNode(
            hostname=hostname,
            ip=ip,
            port=port,
            version=props.get("version", ""),
            gpu_summary=props.get("gpu", ""),
            models=models,
            last_seen=time.time(),
        )
        _registry.upsert(name, node)
        logger.debug("LAN node discovered: %s at %s:%d (%d models)", hostname, ip, port, len(models))


# ---------------------------------------------------------------------------
# mDNS Service Manager
# ---------------------------------------------------------------------------


class MDNSDiscovery:
    """Manages mDNS service advertisement and browsing.

    Usage::

        mdns = MDNSDiscovery()

        # Start advertising this node + browsing for others
        await mdns.start(port=4000, gpu_summary="RTX 4090 24GB", models=["llama3.2"])

        # Get discovered nodes
        nodes = mdns.get_nodes()

        # Stop
        await mdns.stop()
    """

    def __init__(self) -> None:
        self._zc: Any = None
        self._browser: Any = None
        self._service_info: Any = None
        self._started = False

    @property
    def is_available(self) -> bool:
        """Check if zeroconf library is installed."""
        try:
            import zeroconf  # noqa: F401

            return True
        except ImportError:
            return False

    async def start(
        self,
        port: int = 4000,
        gpu_summary: str = "",
        models: list[str] | None = None,
    ) -> bool:
        """Start advertising and browsing. Returns False if zeroconf not installed."""
        if self._started:
            return True

        try:
            from zeroconf import ServiceBrowser, ServiceInfo, Zeroconf
        except ImportError:
            logger.info("zeroconf not installed -- LAN discovery disabled (pip install llmhosts[lan])")
            return False

        loop = asyncio.get_running_loop()

        def _start_sync() -> None:
            from llmhosts import __version__

            self._zc = Zeroconf()

            hostname = socket.gethostname()
            local_ip = _get_local_ip()

            # Prepare TXT record
            models_list = models or []
            # Zeroconf TXT records have a 255 byte per-value limit; truncate if needed
            models_str = ",".join(models_list)
            if len(models_str.encode("utf-8")) > 250:
                # Truncate to fit
                truncated: list[str] = []
                total = 0
                for m in models_list:
                    entry = m + ","
                    if total + len(entry.encode("utf-8")) > 245:
                        break
                    truncated.append(m)
                    total += len(entry.encode("utf-8"))
                models_str = ",".join(truncated)

            properties = {
                "version": __version__,
                "hostname": hostname,
                "port": str(port),
                "gpu": gpu_summary[:250] if gpu_summary else "",
                "models": models_str,
            }

            service_name = f"{SERVICE_NAME_PREFIX}{hostname}.{SERVICE_TYPE}"
            self._service_info = ServiceInfo(
                SERVICE_TYPE,
                service_name,
                addresses=[socket.inet_aton(local_ip)],
                port=port,
                properties=properties,
                server=f"{hostname}.local.",
            )

            # Register our service
            self._zc.register_service(self._service_info)
            logger.info("mDNS: advertising %s at %s:%d", hostname, local_ip, port)

            # Browse for other nodes
            listener = _LLMHostsListener(self._zc)
            self._browser = ServiceBrowser(self._zc, SERVICE_TYPE, listener)
            logger.info("mDNS: browsing for LLMHosts instances on LAN")

        await loop.run_in_executor(None, _start_sync)
        self._started = True
        return True

    async def stop(self) -> None:
        """Stop advertising and browsing."""
        if not self._started or self._zc is None:
            return

        loop = asyncio.get_running_loop()

        def _stop_sync() -> None:
            if self._service_info:
                with contextlib.suppress(Exception):
                    self._zc.unregister_service(self._service_info)
            if self._browser:
                with contextlib.suppress(Exception):
                    self._browser.cancel()
            with contextlib.suppress(Exception):
                self._zc.close()

        await loop.run_in_executor(None, _stop_sync)
        self._started = False
        self._zc = None
        self._browser = None
        self._service_info = None
        logger.info("mDNS: stopped")

    def get_nodes(self) -> list[LANNode]:
        """Return currently discovered LAN nodes (non-stale)."""
        return _registry.active_nodes()

    async def ping_nodes(self) -> list[LANNode]:
        """Ping all discovered nodes and measure latency."""
        nodes = _registry.active_nodes()
        if not nodes:
            return []

        import httpx

        async with httpx.AsyncClient(timeout=httpx.Timeout(5.0)) as client:
            tasks = [_ping_node(client, node) for node in nodes]
            return await asyncio.gather(*tasks)


async def _ping_node(client: Any, node: LANNode) -> LANNode:
    """Ping a single node's /health endpoint and measure latency."""
    import httpx

    try:
        start = time.perf_counter()
        resp = await client.get(f"{node.base_url}/health")
        elapsed = (time.perf_counter() - start) * 1000
        if resp.status_code == 200:
            node.latency_ms = round(elapsed, 1)
        else:
            node.latency_ms = None
    except (httpx.ConnectError, httpx.TimeoutException):
        node.latency_ms = None
    return node


# ---------------------------------------------------------------------------
# Scan (one-shot discovery)
# ---------------------------------------------------------------------------


async def scan_lan(timeout: float = 5.0) -> list[LANNode]:
    """Run a one-shot LAN scan: browse for ``timeout`` seconds and return nodes.

    Useful for ``llmhost discover`` CLI command.
    """
    try:
        from zeroconf import ServiceBrowser, Zeroconf
    except ImportError:
        logger.warning("zeroconf not installed -- cannot scan LAN (pip install llmhosts[lan])")
        return []

    loop = asyncio.get_running_loop()

    def _scan() -> list[LANNode]:
        zc = Zeroconf()
        listener = _LLMHostsListener(zc)
        browser = ServiceBrowser(zc, SERVICE_TYPE, listener)

        # Wait for responses
        time.sleep(timeout)

        browser.cancel()
        nodes = _registry.active_nodes()
        zc.close()
        return nodes

    return await loop.run_in_executor(None, _scan)


# ---------------------------------------------------------------------------
# Utility
# ---------------------------------------------------------------------------


def _get_local_ip() -> str:
    """Get the local machine's LAN IP address."""
    try:
        # Connect to a public IP (doesn't send data) to determine local route
        with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
            s.connect(("10.255.255.255", 1))
            return str(s.getsockname()[0])
    except Exception:
        return "127.0.0.1"
