#!/usr/bin/env python3

# Libervia plugin for Serverless XMPP
# Copyright (C) 2009-2026 Jérôme Poisson (goffi@goffi.org)

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.

# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

from collections.abc import Callable
import socket
from typing import cast

from twisted.internet import reactor
from twisted.internet.interfaces import IReactorCore
from zeroconf import ServiceBrowser, ServiceInfo, ServiceListener, Zeroconf
from zeroconf.asyncio import AsyncZeroconf

from libervia.backend.core.log import getLogger
from libervia.backend.plugins.plugin_misc_serverless.models import DiscoveredPeerInfo
from libervia.backend.tools.utils import aio

log = getLogger(__name__)

SERVICE_TYPE = "_xmpp-serverless._tcp.local."
DEFAULT_PORT = 5299


class MDNSManager:
    """mDNS-based peer discovery.

    Broadcasts our peer on the local network and listens for others.
    """

    def __init__(
        self,
        peer_id: str,
        port: int = DEFAULT_PORT,
        on_peer_found: (
            Callable[[str, str, list[str], dict[str, str]], None] | None
        ) = None,
        on_peer_lost: Callable[[str], None] | None = None,
    ):
        """Init the manager.

        @param peer_id: Peer ID (from libP2P).
        @param port: TCP port to advertise.
        @param on_peer_found: callback called when a peer is found.
        @param on_peer_lost: callback called when a peer is lost.
        """
        self.peer_id = peer_id
        self.port = port
        self.on_peer_found = on_peer_found
        self.on_peer_lost = on_peer_lost
        self.zeroconf = Zeroconf()
        # Human-readable device name; zeroconf will auto-rename on collision
        device_name = f"Libervia-{socket.gethostname()}"
        self.service_name = f"{device_name}.{SERVICE_TYPE}"
        log.info(
            f"MDNSManager created: peer_id={peer_id}, port={port}, "
            f"service_name={self.service_name}"
        )

    def start(self) -> None:
        """Register this peer and start listening for others."""
        log.info(f"Registering mDNS service: {self.service_name} port={self.port}")
        self.service_info = ServiceInfo(
            type_=SERVICE_TYPE,
            name=self.service_name,
            port=self.port,
            properties={b"id": self.peer_id.encode()},
            server=f"{socket.gethostname()}.local.",
        )
        self.zeroconf.register_service(self.service_info, allow_name_change=True)
        # Zeroconf may have renamed on collision, we have to check the final name.
        self.service_name = self.service_info.name
        log.info(f"mDNS service registered as: {self.service_info.name}")

        self.listener = PeersListener(
            zeroconf=self.zeroconf,
            service_type=SERVICE_TYPE,
            service_name=self.service_name,
            on_peer_found=self.on_peer_found,
            on_peer_lost=self.on_peer_lost,
        )
        log.info("ServiceBrowser started for peer discovery.")

    @aio
    async def stop(self) -> None:
        """Unregister this peer and clean up resources.

        Uses AsyncZeroconf to await the unregistration (goodbye broadcasts),
        then calls close() to shut down the zeroconf instance.
        Decorated with @aio so asyncio.sleep() inside the broadcasts works
        correctly (scheduled via asyncio.ensure_future, not ensureDeferred).
        Returns a Deferred.
        """
        log.info("Cleaning mDNS services.")
        if hasattr(self, "listener"):
            self.listener.browser.cancel()
            log.debug("ServiceBrowser cancelled.")
        if hasattr(self, "service_info"):
            async_zc = AsyncZeroconf(zc=self.zeroconf)
            # async_unregister_service returns an Awaitable (a Task) that must be
            # double-awaited: first await retrieves the Task, second awaits it.
            unregister_task = await async_zc.async_unregister_service(
                self.service_info,
            )
            await unregister_task
            log.debug("mDNS service unregistered via AsyncZeroconf.")
        self.zeroconf.close()
        log.debug("Zeroconf closed.")


class PeersListener(ServiceListener):
    """Listen for other libp2p peers on the local network.

    Callbacks are fired from zeroconf's ServiceBrowser thread. All model updates must be
    sent to the Twisted reactor thread with reactor.callFromThread().
    """

    def __init__(
        self,
        zeroconf: Zeroconf,
        service_type: str,
        service_name: str,
        on_peer_found: (
            Callable[[str, str, list[str], dict[str, str]], None] | None
        ) = None,
        on_peer_lost: Callable[[str], None] | None = None,
    ) -> None:
        """Initialize the listener.

        @param zeroconf: Zeroconf instance.
        @param service_type: mDNS service type to listen for.
        @param service_name: our own service name (skipped on discovery).
        @param on_peer_found: callback called when a peer is found.
        @param on_peer_lost: callback called when a peer is lost.
        """
        self.zeroconf = zeroconf
        self.service_type = service_type
        self.service_name = service_name
        self.on_peer_found = on_peer_found
        self.on_peer_lost = on_peer_lost
        self.browser = ServiceBrowser(zeroconf, service_type, listener=self)

    def add_service(self, zc: Zeroconf, type_: str, name: str) -> None:
        """Handle a newly discovered mDNS service.

        @param zc: Zeroconf instance.
        @param type_: service type.
        @param name: fully qualified service name.
        """
        if name == self.service_name:
            return
        peer_info = self._extract_peer_info(zc, type_, name)
        if peer_info and self.on_peer_found:
            log.debug(f"mDNS peer discovered: {peer_info['device_id']}")
            cast(IReactorCore, reactor).callFromThread(
                self.on_peer_found,
                peer_info["device_id"],
                peer_info["name"],
                peer_info["addrs"],
                peer_info["metadata"],
            )

    def remove_service(self, zc: Zeroconf, type_: str, name: str) -> None:
        """Handle a removed mDNS service.

        @param zc: Zeroconf instance.
        @param type_: service type.
        @param name: fully qualified service name.
        """
        if name == self.service_name:
            return
        peer_info = self._extract_peer_info(zc, type_, name)
        device_id = peer_info["device_id"] if peer_info else None
        if device_id is not None and self.on_peer_lost:
            log.debug(f"mDNS peer removed: {device_id}")
            cast(IReactorCore, reactor).callFromThread(self.on_peer_lost, device_id)

    def update_service(self, zc: Zeroconf, type_: str, name: str) -> None:
        """Handle an updated mDNS service by removing and re-adding it.

        @param zc: Zeroconf instance.
        @param type_: service type.
        @param name: fully qualified service name.
        """
        self.remove_service(zc, type_, name)
        self.add_service(zc, type_, name)

    def _extract_peer_info(
        self, zc: Zeroconf, type_: str, name: str
    ) -> DiscoveredPeerInfo | None:
        """Extract peer info from a zeroconf service.

        Handles both IPv4 and IPv6 addresses.
        Deduplicates addresses that may appear from multiple interfaces.

        @return: mDNS peer infos, or None on failure.
        """
        try:
            info = zc.get_service_info(type_, name, timeout=5000)
            if not info:
                log.debug(f"No service info for {name!r}")
                return None
            pid_bytes = info.properties.get(b"id")
            if not pid_bytes:
                log.warning(f'No "id" property for {name!r}, {info.properties=}')
                return None
            # Extract human-readable name from full service name
            # e.g. "Libervia-example-2._xmpp-serverless._tcp.local." -> "Libervia-example"
            peer_name = name.split(".")[0]
            # Build libp2p multiaddr strings, deduplicating by IP string.
            # See https://github.com/multiformats/multiaddr
            seen_ips: set[str] = set()
            addrs = []
            for addr in info.addresses:
                if len(addr) == 16:
                    # IPv6
                    ip_str = socket.inet_ntop(socket.AF_INET6, addr)
                else:
                    # 4 bytes = IPv4
                    ip_str = socket.inet_ntop(socket.AF_INET, addr)
                if ip_str not in seen_ips:
                    seen_ips.add(ip_str)
                    af = "ip6" if len(addr) == 16 else "ip4"
                    addrs.append(f"/{af}/{ip_str}/tcp/{info.port}")
            metadata = {
                key.decode(): value.decode()
                for key, value in info.properties.items()
                if value is not None
            }
            return {
                "name": peer_name,
                "device_id": pid_bytes.decode(),
                "addrs": addrs,
                "metadata": metadata,
            }
        except Exception:
            log.exception(f"extract_peer_info failed for {name!r}")
            return None