# Familiar - Self-Hosted AI Companion Platform
# Copyright (c) 2026 George Scott Foley
# ORCID: 0009-0006-4957-0540
# Licensed under the MIT License

"""
Mesh Discovery — Phase 1

mDNS-based auto-discovery for Familiar mesh nodes on a local network.
Wraps the existing MeshGateway/MeshConfig with zero-configuration
LAN discovery via the zeroconf library.

Architecture:
    - Advertises this node's MeshGateway WebSocket endpoint via mDNS
    - Listens for other Familiar nodes on the same subnet
    - Maintains a PeerStore of discovered + manual peers
    - Triggers peer gateway connections for trusted peers (Phase 2)

Integration:
    Called from MeshGateway.start() when discovery_enabled=True.
    Uses existing MeshConfig for node_id, node_name, gateway_port.
    Stores peer state in MESH_DIR/peers.json.
"""

from __future__ import annotations

import hashlib
import json
import logging
import socket
import threading
from dataclasses import asdict, dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional

logger = logging.getLogger(__name__)

# Service type for Familiar mesh — follows mDNS convention
SERVICE_TYPE = "_familiar-mesh._tcp.local."
FAMILIAR_VERSION = "1.10.0"

# Try zeroconf — graceful fallback if not installed
try:
    from zeroconf import (
        IPVersion,
        ServiceBrowser,
        ServiceInfo,
        ServiceStateChange,
        Zeroconf,
    )

    HAS_ZEROCONF = True
except ImportError:
    HAS_ZEROCONF = False
    logger.debug("zeroconf not installed — mDNS discovery unavailable.")


def _utcnow() -> datetime:
    return datetime.now(timezone.utc)


def _get_local_ip() -> str:
    """Get the primary local IP address (not 127.0.0.1)."""
    try:
        # Connect to a public DNS to determine our LAN-facing IP
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        s.settimeout(1)
        s.connect(("8.8.8.8", 80))
        ip = s.getsockname()[0]
        s.close()
        return ip
    except Exception:
        return "127.0.0.1"


# ============================================================
# DATA MODELS
# ============================================================


@dataclass
class DiscoveredPeer:
    """A peer discovered via mDNS or added manually."""

    node_id: str
    node_name: str
    node_type: str  # 'gateway', 'desktop', 'mobile', 'server'
    host: str  # IP address
    port: int  # WebSocket port
    version: str = ""
    skills: List[str] = field(default_factory=list)
    capabilities: List[str] = field(default_factory=list)
    fingerprint: str = ""  # First 16 chars of SHA-256 of public key
    discovery_source: str = ""  # 'mdns', 'manual'
    first_seen: str = ""
    last_seen: str = ""
    trust_level: str = "unknown"  # 'unknown', 'pending', 'trusted', 'blocked'

    def __post_init__(self):
        now = _utcnow().isoformat()
        if not self.first_seen:
            self.first_seen = now
        if not self.last_seen:
            self.last_seen = now

    def to_dict(self) -> dict:
        return asdict(self)

    @classmethod
    def from_dict(cls, d: dict) -> "DiscoveredPeer":
        # Handle extra/missing keys gracefully
        known = {f.name for f in cls.__dataclass_fields__.values()}
        return cls(**{k: v for k, v in d.items() if k in known})

    @property
    def ws_url(self) -> str:
        return f"ws://{self.host}:{self.port}"

    @property
    def is_stale(self) -> bool:
        """Peer not seen in over 2 minutes."""
        try:
            last = datetime.fromisoformat(self.last_seen)
            return (_utcnow() - last).total_seconds() > 120
        except (ValueError, TypeError):
            return True

    def touch(self):
        """Update last_seen timestamp."""
        self.last_seen = _utcnow().isoformat()


# ============================================================
# PEER STORE — persistence layer
# ============================================================


class PeerStore:
    """
    Persists discovered and manual peers to disk.
    Thread-safe. Stored in MESH_DIR/peers.json.
    """

    def __init__(self, storage_path: Path):
        self._path = storage_path / "peers.json"
        self._lock = threading.Lock()
        self._peers: Dict[str, DiscoveredPeer] = {}
        self._load()

    def _load(self):
        if self._path.exists():
            try:
                data = json.loads(self._path.read_text())
                for entry in data.get("peers", []):
                    peer = DiscoveredPeer.from_dict(entry)
                    self._peers[peer.node_id] = peer
            except (json.JSONDecodeError, IOError, OSError) as e:
                logger.warning(f"Failed to load peer store: {e}")

    def _save(self):
        try:
            self._path.parent.mkdir(parents=True, exist_ok=True)
            data = {"peers": [p.to_dict() for p in self._peers.values()]}
            self._path.write_text(json.dumps(data, indent=2))
        except (IOError, OSError) as e:
            logger.warning(f"Failed to save peer store: {e}")

    def get(self, node_id: str) -> Optional[DiscoveredPeer]:
        with self._lock:
            return self._peers.get(node_id)

    def upsert(self, peer: DiscoveredPeer) -> bool:
        """Insert or update a peer. Returns True if new."""
        with self._lock:
            is_new = peer.node_id not in self._peers
            if not is_new:
                # Preserve trust_level from existing record
                existing = self._peers[peer.node_id]
                peer.trust_level = existing.trust_level
                peer.first_seen = existing.first_seen
            self._peers[peer.node_id] = peer
            self._save()
            return is_new

    def remove(self, node_id: str) -> bool:
        with self._lock:
            if node_id in self._peers:
                del self._peers[node_id]
                self._save()
                return True
            return False

    def set_trust(self, node_id: str, level: str) -> bool:
        """Set trust level for a peer."""
        with self._lock:
            peer = self._peers.get(node_id)
            if peer:
                peer.trust_level = level
                self._save()
                return True
            return False

    def get_all(self) -> List[DiscoveredPeer]:
        with self._lock:
            return list(self._peers.values())

    def get_by_trust(self, level: str) -> List[DiscoveredPeer]:
        with self._lock:
            return [p for p in self._peers.values() if p.trust_level == level]

    def get_manual(self) -> List[DiscoveredPeer]:
        with self._lock:
            return [p for p in self._peers.values() if p.discovery_source == "manual"]

    def prune_stale(self, max_age_seconds: int = 300) -> int:
        """Remove stale mDNS-discovered peers (not manual)."""
        with self._lock:
            stale = [
                nid for nid, p in self._peers.items() if p.discovery_source == "mdns" and p.is_stale
            ]
            for nid in stale:
                del self._peers[nid]
            if stale:
                self._save()
            return len(stale)

    def count(self) -> int:
        with self._lock:
            return len(self._peers)


# ============================================================
# mDNS SERVICE INFO BUILDER
# ============================================================


def _build_txt_properties(
    node_id: str,
    node_name: str,
    node_type: str,
    skills: List[str],
    capabilities: List[str],
    fingerprint: str = "",
    model_count: int = 0,
) -> dict:
    """Build TXT record properties for mDNS advertisement."""
    props = {
        "version": FAMILIAR_VERSION,
        "node_id": node_id,
        "node_name": node_name,
        "node_type": node_type,
        "skills": ",".join(skills[:50]),  # Limit to avoid TXT record size issues
        "capabilities": ",".join(capabilities),
        "fingerprint": fingerprint[:16],
    }
    if model_count > 0:
        props["model_count"] = str(model_count)
    return props


def _parse_txt_properties(properties: dict) -> dict:
    """Parse TXT record properties from mDNS service info."""
    parsed = {}
    for key, value in properties.items():
        # zeroconf returns bytes for TXT values
        if isinstance(key, bytes):
            key = key.decode("utf-8", errors="replace")
        if isinstance(value, bytes):
            value = value.decode("utf-8", errors="replace")
        parsed[key] = value

    # Parse comma-separated lists
    if "skills" in parsed:
        parsed["skills"] = [s.strip() for s in parsed["skills"].split(",") if s.strip()]
    else:
        parsed["skills"] = []

    if "capabilities" in parsed:
        parsed["capabilities"] = [c.strip() for c in parsed["capabilities"].split(",") if c.strip()]
    else:
        parsed["capabilities"] = []

    return parsed


# ============================================================
# MESH DISCOVERY — the main class
# ============================================================


class MeshDiscovery:
    """
    mDNS discovery for Familiar mesh nodes.

    Advertises this node on the LAN and discovers peers.
    Integrates with MeshGateway via callbacks.

    Usage:
        discovery = MeshDiscovery(
            node_id="abc123",
            node_name="Kitchen Pi",
            node_type="gateway",
            port=18789,
            skills=["calendar", "meal_planner"],
            mesh_dir=Path("~/.familiar/data/mesh"),
        )
        discovery.start()
        # ... later ...
        discovery.stop()
    """

    def __init__(
        self,
        node_id: str,
        node_name: str,
        node_type: str = "gateway",
        port: int = 18789,
        skills: Optional[List[str]] = None,
        capabilities: Optional[List[str]] = None,
        fingerprint: str = "",
        mesh_dir: Optional[Path] = None,
        max_peers: int = 16,
        on_peer_found: Optional[Callable[[DiscoveredPeer], None]] = None,
        on_peer_lost: Optional[Callable[[str], None]] = None,
        on_peer_updated: Optional[Callable[[DiscoveredPeer], None]] = None,
    ):
        self.node_id = node_id
        self.node_name = node_name
        self.node_type = node_type
        self.port = port
        self.skills = skills or []
        self.capabilities = capabilities or ["chat", "tools"]
        self.fingerprint = fingerprint
        self.max_peers = max_peers

        # Callbacks
        self._on_peer_found = on_peer_found
        self._on_peer_lost = on_peer_lost
        self._on_peer_updated = on_peer_updated

        # Peer store
        if mesh_dir is None:
            mesh_dir = Path.home() / ".familiar" / "data" / "mesh"
        self.peer_store = PeerStore(mesh_dir)

        # zeroconf state
        self._zeroconf: Optional[Any] = None
        self._browser: Optional[Any] = None
        self._service_info: Optional[Any] = None
        self._running = False
        self._local_ip = _get_local_ip()

    @property
    def is_running(self) -> bool:
        return self._running

    def start(self):
        """Start mDNS advertisement and listening."""
        if not HAS_ZEROCONF:
            logger.warning("zeroconf not installed — mDNS discovery disabled.")
            return

        if self._running:
            logger.warning("Discovery already running")
            return

        logger.info(
            f"Starting mesh discovery: node={self.node_name} ({self.node_id[:8]}...) "
            f"port={self.port} ip={self._local_ip}"
        )

        self._zeroconf = Zeroconf(ip_version=IPVersion.V4Only)

        # Register our service
        self._register_service()

        # Browse for peers
        self._browser = ServiceBrowser(
            self._zeroconf,
            SERVICE_TYPE,
            handlers=[self._on_service_state_change],
        )

        self._running = True
        logger.info("Mesh discovery active")

    def stop(self):
        """Stop mDNS advertisement and listening."""
        if not self._running:
            return

        logger.info("Stopping mesh discovery")
        self._running = False

        # Unregister our service
        if self._zeroconf and self._service_info:
            try:
                self._zeroconf.unregister_service(self._service_info)
            except Exception as e:
                logger.debug(f"Error unregistering service: {e}")

        # Close browser
        if self._browser:
            try:
                self._browser.cancel()
            except Exception as e:
                logger.debug(f"Error canceling browser: {e}")

        # Close zeroconf
        if self._zeroconf:
            try:
                self._zeroconf.close()
            except Exception as e:
                logger.debug(f"Error closing zeroconf: {e}")

        self._zeroconf = None
        self._browser = None
        self._service_info = None

    def update_skills(self, skills: List[str]):
        """Update advertised skills (e.g., after skill loading)."""
        self.skills = skills
        if self._running and self._zeroconf and self._service_info:
            try:
                self._zeroconf.unregister_service(self._service_info)
                self._register_service()
            except Exception as e:
                logger.warning(f"Failed to update mDNS skills: {e}")

    def add_manual_peer(self, host: str, port: int, name: str = "") -> DiscoveredPeer:
        """Add a peer manually (for cross-network connections)."""
        # Generate a deterministic node_id from host:port
        node_id = hashlib.sha256(f"{host}:{port}".encode()).hexdigest()[:16]

        peer = DiscoveredPeer(
            node_id=node_id,
            node_name=name or f"{host}:{port}",
            node_type="gateway",
            host=host,
            port=port,
            discovery_source="manual",
        )
        is_new = self.peer_store.upsert(peer)
        if is_new:
            logger.info(f"Manual peer added: {peer.node_name} ({host}:{port})")
            if self._on_peer_found:
                self._on_peer_found(peer)
        return peer

    def remove_manual_peer(self, host: str, port: int) -> bool:
        """Remove a manually added peer."""
        node_id = hashlib.sha256(f"{host}:{port}".encode()).hexdigest()[:16]
        removed = self.peer_store.remove(node_id)
        if removed:
            logger.info(f"Manual peer removed: {host}:{port}")
            if self._on_peer_lost:
                self._on_peer_lost(node_id)
        return removed

    # ── mDNS registration ──

    def _register_service(self):
        """Register our mDNS service."""
        txt_props = _build_txt_properties(
            node_id=self.node_id,
            node_name=self.node_name,
            node_type=self.node_type,
            skills=self.skills,
            capabilities=self.capabilities,
            fingerprint=self.fingerprint,
        )

        service_name = f"{self.node_id}.{SERVICE_TYPE}"

        self._service_info = ServiceInfo(
            type_=SERVICE_TYPE,
            name=service_name,
            server=f"{self.node_id}.local.",
            addresses=[socket.inet_aton(self._local_ip)],
            port=self.port,
            properties=txt_props,
        )

        self._zeroconf.register_service(self._service_info)
        logger.info(f"Registered mDNS service: {service_name}")

    # ── mDNS callbacks ──

    def _on_service_state_change(
        self,
        zeroconf: Any,
        service_type: str,
        name: str,
        state_change: Any,
    ):
        """Handle mDNS service state changes."""
        if state_change == ServiceStateChange.Added:
            self._on_service_added(zeroconf, service_type, name)
        elif state_change == ServiceStateChange.Removed:
            self._on_service_removed(name)
        elif state_change == ServiceStateChange.Updated:
            self._on_service_added(zeroconf, service_type, name)

    def _on_service_added(self, zeroconf: Any, service_type: str, name: str):
        """Handle a new or updated mDNS service."""
        info = zeroconf.get_service_info(service_type, name, timeout=3000)
        if not info:
            return

        # Parse TXT records
        txt = _parse_txt_properties(info.properties)
        peer_node_id = txt.get("node_id", "")

        # Skip ourselves
        if peer_node_id == self.node_id:
            return

        # Skip if at max peers
        if self.peer_store.count() >= self.max_peers:
            if not self.peer_store.get(peer_node_id):
                logger.debug(
                    f"Max peers ({self.max_peers}) reached, ignoring {txt.get('node_name')}"
                )
                return

        # Build peer from service info
        addresses = info.parsed_scoped_addresses()
        if not addresses:
            return

        host = addresses[0]

        peer = DiscoveredPeer(
            node_id=peer_node_id,
            node_name=txt.get("node_name", name),
            node_type=txt.get("node_type", "unknown"),
            host=host,
            port=info.port,
            version=txt.get("version", ""),
            skills=txt.get("skills", []),
            capabilities=txt.get("capabilities", []),
            fingerprint=txt.get("fingerprint", ""),
            discovery_source="mdns",
        )

        is_new = self.peer_store.upsert(peer)

        if is_new:
            logger.info(
                f"Peer discovered: {peer.node_name} ({peer.host}:{peer.port}) "
                f"skills={peer.skills[:5]} type={peer.node_type}"
            )
            if self._on_peer_found:
                self._on_peer_found(peer)
        else:
            # Update last_seen
            existing = self.peer_store.get(peer_node_id)
            if existing:
                existing.touch()
                existing.skills = peer.skills
                existing.capabilities = peer.capabilities
                existing.host = peer.host
                existing.port = peer.port
                self.peer_store.upsert(existing)
            if self._on_peer_updated:
                self._on_peer_updated(peer)

    def _on_service_removed(self, name: str):
        """Handle an mDNS service being removed."""
        # Extract node_id from service name: "{node_id}._familiar-mesh._tcp.local."
        node_id = name.split(".")[0] if "." in name else name

        if node_id == self.node_id:
            return

        peer = self.peer_store.get(node_id)
        if peer and peer.discovery_source == "mdns":
            logger.info(f"Peer gone: {peer.node_name} ({peer.host}:{peer.port})")
            # Don't remove — mark stale. They may come back.
            peer.touch()  # Will be pruned if not seen again
            if self._on_peer_lost:
                self._on_peer_lost(node_id)

    # ── Status ──

    def get_status(self) -> dict:
        """Get discovery status for CLI/dashboard."""
        peers = self.peer_store.get_all()
        return {
            "running": self._running,
            "node_id": self.node_id,
            "node_name": self.node_name,
            "local_ip": self._local_ip,
            "port": self.port,
            "peer_count": len(peers),
            "peers": [
                {
                    "node_id": p.node_id[:8] + "...",
                    "name": p.node_name,
                    "host": p.host,
                    "port": p.port,
                    "type": p.node_type,
                    "trust": p.trust_level,
                    "source": p.discovery_source,
                    "skills": len(p.skills),
                    "stale": p.is_stale,
                }
                for p in peers
            ],
        }
