"""mDNS/DNS-SD peer discovery and manual peer list."""

from __future__ import annotations

import logging
from typing import Optional

from radix_edge.mesh.security import mesh_key_prefix

logger = logging.getLogger("radix_edge.mesh.discovery")

SERVICE_TYPE = "_radix-mesh._tcp.local."

# Attempt to import zeroconf; fall back gracefully
try:
    from zeroconf import ServiceInfo
    from zeroconf.asyncio import AsyncZeroconf, AsyncServiceBrowser

    ZEROCONF_AVAILABLE = True
except ImportError:
    ZEROCONF_AVAILABLE = False
    logger.debug("zeroconf not installed; mDNS discovery unavailable, using manual mode")


class MeshDiscovery:
    """mDNS-based peer discovery using zeroconf.

    Registers a ``_radix-mesh._tcp.local.`` service and browses for peers
    that share the same mesh key prefix in their TXT records.
    """

    def __init__(
        self,
        node_id: str,
        port: int,
        mesh_key: str,
        gpu_count: int = 0,
    ):
        self.node_id = node_id
        self.port = port
        self.mesh_key = mesh_key
        self.gpu_count = gpu_count
        self._key_prefix = mesh_key_prefix(mesh_key)
        self._discovered: dict[str, dict] = {}  # node_id -> {host, port}
        self._zeroconf: Optional[object] = None  # AsyncZeroconf when started
        self._browser: Optional[object] = None

    @property
    def discovered_peers(self) -> dict[str, dict]:
        """Return a copy of discovered peers."""
        return dict(self._discovered)

    async def start(self) -> None:
        """Start mDNS registration and browsing."""
        if not ZEROCONF_AVAILABLE:
            logger.warning("zeroconf not installed; mDNS discovery skipped")
            return

        self._zeroconf = AsyncZeroconf()

        # Register our own service
        info = ServiceInfo(
            SERVICE_TYPE,
            f"{self.node_id}.{SERVICE_TYPE}",
            port=self.port,
            properties={
                "node_id": self.node_id,
                "mesh_key_prefix": self._key_prefix,
                "gpu_count": str(self.gpu_count),
            },
        )
        await self._zeroconf.async_register_service(info)
        logger.info("mDNS service registered: %s port=%d", self.node_id, self.port)

        # Browse for peers
        self._browser = AsyncServiceBrowser(
            self._zeroconf.zeroconf,
            SERVICE_TYPE,
            handlers=[self._on_service_state_change],
        )
        logger.info("mDNS browsing started for %s", SERVICE_TYPE)

    def _on_service_state_change(
        self,
        zeroconf: object,
        service_type: str,
        name: str,
        state_change: object,
    ) -> None:
        """Handle discovered/removed services."""
        if not ZEROCONF_AVAILABLE:
            return

        # Compare by name to support both the real enum and string values
        state_name = getattr(state_change, "name", str(state_change))
        if state_name == "Removed":
            # Try to find and remove the peer
            for nid, info in list(self._discovered.items()):
                if info.get("service_name") == name:
                    del self._discovered[nid]
                    logger.info("Peer removed via mDNS: %s", nid)
                    break
            return

        # Added or Updated
        info = zeroconf.get_service_info(service_type, name)
        if info is None:
            return

        props = {
            k.decode("utf-8") if isinstance(k, bytes) else k:
            v.decode("utf-8") if isinstance(v, bytes) else v
            for k, v in (info.properties or {}).items()
        }

        peer_key_prefix = props.get("mesh_key_prefix", "")
        if peer_key_prefix != self._key_prefix:
            logger.debug("Ignoring peer with different mesh key prefix: %s", name)
            return

        peer_node_id = props.get("node_id", "")
        if peer_node_id == self.node_id:
            return  # Skip self

        host = info.parsed_addresses()[0] if info.parsed_addresses() else None
        if not host:
            return

        self._discovered[peer_node_id] = {
            "host": host,
            "port": info.port,
            "service_name": name,
        }
        logger.info("Peer discovered via mDNS: %s at %s:%d", peer_node_id, host, info.port)

    async def stop(self) -> None:
        """Stop mDNS registration and browsing."""
        if self._browser is not None:
            self._browser.cancel()
            self._browser = None
        if self._zeroconf is not None:
            await self._zeroconf.async_close()
            self._zeroconf = None
        logger.info("mDNS discovery stopped")


class ManualPeerList:
    """Static peer list from config (``RADIX_MESH_PEERS`` env var).

    Format: ``host1:port1,host2:port2``
    """

    def __init__(self, peers_str: str):
        self._peers: list[dict] = []
        if peers_str.strip():
            for entry in peers_str.split(","):
                entry = entry.strip()
                if not entry:
                    continue
                if ":" not in entry:
                    logger.warning("Invalid peer entry (missing port): %s", entry)
                    continue
                host, _, port_str = entry.rpartition(":")
                try:
                    port = int(port_str)
                except ValueError:
                    logger.warning("Invalid port in peer entry: %s", entry)
                    continue
                self._peers.append({"host": host, "port": port})

    @property
    def peers(self) -> list[dict]:
        """Return the list of manually configured peers."""
        return list(self._peers)
