"""
Relais UDP de secours — utilisé UNIQUEMENT quand deux pairs ne parviennent pas à
établir une connexion WireGuard directe (NAT symétrique des deux côtés, par exemple).

Principe : un **relais aveugle**, pas un protocole applicatif. WireGuard envoie des
paquets bruts (chiffrés) à l'`Endpoint` configuré — il ne peut pas les préfixer d'un
en-tête applicatif (ex: un token de destination), donc le relais ne peut pas non plus
en exiger un. À la place :

- Ce processus écoute en permanence sur une PLAGE de ports UDP (voir server/config.py,
  `RELAY_PORT_POOL_START`/`RELAY_PORT_POOL_SIZE`) — un port par session de relais.
- app.py (le serveur de signalisation) alloue un port libre de cette plage à une paire
  de pairs qui n'arrive pas à se connecter directement, et le leur communique via la
  liste de pairs (champ `relay_endpoint`) — voir app.py::_allocate_relay_port.
- Chaque port, une fois qu'il voit deux adresses UDP sources distinctes (les deux
  pairs, ayant chacun pointé leur `Endpoint` WireGuard dessus), réexpédie aveuglément
  tout paquet reçu de l'une vers l'autre. Le relais ne déchiffre ni n'inspecte jamais
  le contenu : c'est un simple forwarder, pas un man-in-the-middle.
- Une 3e adresse source sur un port déjà occupé par deux pairs est ignorée (protection
  basique contre un tiers qui devinerait le port).
- Un port inactif plus de `RELAY_SESSION_TTL_SECONDS` "oublie" les adresses qu'il avait
  apprises et redevient disponible pour une nouvelle paire (voir app.py, qui applique
  la même expiration côté allocation — les deux doivent rester cohérents).

Lancement autonome : python relay.py
Windows : autoriser toute la plage de ports UDP (par défaut 8766-8785) dans le
pare-feu Windows Defender si ce serveur est exposé.
"""

from __future__ import annotations

import asyncio
import logging
import time
from typing import Callable

import config
import relay_activity
import relay_ws
from relay_metrics import CompteurDeSession, format_octets

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger("meshamoto-relay")

# Une adresse apprise est soit un couple UDP `(ip, port)`, soit l'adresse VIRTUELLE
# d'un porteur non-UDP — `("ws", <numéro>)` pour une WebSocket (voir relay_ws). Les
# deux formes cohabitent dans le même dictionnaire parce qu'elles jouent exactement le
# même rôle : occuper une des deux places d'une session.
Adresse = tuple


class BlindPairRelayProtocol(asyncio.DatagramProtocol):
    """
    Relais aveugle pour UN port : apprend jusqu'à deux adresses sources distinctes,
    puis réexpédie tout paquet reçu de l'une vers l'autre, sans jamais l'inspecter.
    """

    def __init__(self, port: int) -> None:
        self.port = port
        self.transport: asyncio.DatagramTransport | None = None
        # adresse (ip, port) -> timestamp de dernière activité
        self._known: dict[Adresse, float] = {}
        # Porteurs NON-UDP attachés à ce port (WebSocket/TLS, voir relay_ws) :
        # adresse virtuelle -> fonction d'émission. Un pair dont le réseau bloque
        # l'UDP arrive par là, mais il occupe une des DEUX places comme n'importe
        # quelle adresse UDP — c'est tout l'intérêt de le brancher ici plutôt que
        # d'écrire un second relais : une seule logique d'appariement, un seul TTL,
        # un seul compteur.
        self._porteurs: dict[Adresse, Callable[[bytes], None]] = {}
        # Comptabilité du trafic — n'influence AUCUN comportement, voir
        # relay_metrics : on mesure avant d'envisager de brider.
        self.compteur = CompteurDeSession()
        self._derniere_reexpedition: float | None = None

    @property
    def derniere_activite(self) -> float | None:
        """
        Date du dernier paquet RÉEXPÉDIÉ, ou None si ce port ne sert personne.

        Publié à app.py pour maintenir son allocation en vie (voir relay_activity). On
        s'appuie sur la réexpédition et non sur la simple réception, pour la même raison
        qu'au point 56 : un paquet reçu alors qu'une seule adresse est connue est jeté,
        et le prendre pour un signe de vie ferait réserver un port à quelqu'un qui parle
        dans le vide — un pair mal configuré immobiliserait un port du pool.

        Retombe à None dès que la purge a vidé les adresses apprises : une session
        expirée ne doit surtout pas continuer d'être annoncée comme active, sinon son
        port resterait réservé pour toujours.
        """
        if not self._known:
            return None
        return self._derniere_reexpedition

    def connection_made(self, transport: asyncio.DatagramTransport) -> None:  # type: ignore[override]
        self.transport = transport

    def attacher_porteur(self, adresse: Adresse, emettre: Callable[[bytes], None]) -> None:
        """
        Branche un porteur non-UDP (WebSocket) sur ce port, sous une adresse virtuelle.

        **N'occupe AUCUNE place tant qu'aucun paquet n'est arrivé** : la règle des deux
        adresses reste appliquée au seul endroit qui la connaît, `paquet_recu`. Réserver
        la place ici ferait diverger les deux chemins, et un porteur qui se connecte
        sans jamais rien émettre priverait une vraie paire de son port.
        """
        self._porteurs[adresse] = emettre

    def detacher_porteur(self, adresse: Adresse) -> None:
        """
        Retire un porteur, et **oublie l'adresse apprise avec lui**.

        Une adresse UDP n'est oubliée qu'à l'expiration, faute de savoir si le pair est
        parti ; une WebSocket, elle, se ferme explicitement. Garder sa place occupée
        interdirait au pair de revenir avant `RELAY_SESSION_TTL_SECONDS` — or il se
        reconnecte en une seconde (point 70).
        """
        self._porteurs.pop(adresse, None)
        self._known.pop(adresse, None)

    def _emettre(self, data: bytes, adresse: Adresse) -> bool:
        """Expédie vers une adresse apprise, quel que soit son porteur."""
        porteur = self._porteurs.get(adresse)
        if porteur is not None:
            porteur(data)
            return True
        if self.transport is not None:
            self.transport.sendto(data, adresse)
            return True
        return False

    def paquet_recu(self, data: bytes, addr: Adresse) -> None:
        """
        Cœur de l'appariement, commun à TOUS les porteurs.

        Point d'entrée unique délibéré : `datagram_received` (UDP) et `relay_ws`
        (WebSocket) passent par ici, donc la limite de deux adresses, l'expiration et
        la comptabilité sont écrites une seule fois. Un second chemin les ferait
        diverger en silence — et le porteur WebSocket contournerait alors la protection
        contre un tiers qui devinerait le port.
        """
        now = time.time()
        if addr in self._known:
            self._known[addr] = now
        elif len(self._known) < 2:
            self._known[addr] = now
            logger.info("Port %d : nouvelle adresse apprise %s (%d/2)", self.port, addr, len(self._known))
        else:
            logger.debug("Port %d : adresse %s ignorée, session déjà pleine", self.port, addr)
            return

        if len(self._known) == 2:
            other = next(a for a in self._known if a != addr)
            if not self._emettre(data, other):
                return
            # Compté APRÈS la réexpédition : un paquet reçu alors qu'une seule
            # adresse est connue est jeté, et le compter laisserait croire à un
            # trafic qui n'a jamais coûté de bande passante en sortie.
            self.compteur.compter(len(data), now)
            self._derniere_reexpedition = now

    def datagram_received(self, data: bytes, addr: tuple[str, int]) -> None:  # type: ignore[override]
        self.paquet_recu(data, addr)

    def purge_expired(self, ttl_seconds: float) -> None:
        now = time.time()
        expired = [addr for addr, ts in self._known.items() if now - ts > ttl_seconds]
        for addr in expired:
            del self._known[addr]
        if expired:
            logger.info(
                "Port %d : %d adresse(s) expirée(s), session libérée après %s en %d paquets",
                self.port, len(expired),
                format_octets(self.compteur.octets), self.compteur.paquets,
            )
            # Remis à zéro avec la session : le port sera réattribué à une autre
            # paire, et cumuler leurs trafics rendrait le chiffre ininterprétable.
            self.compteur = CompteurDeSession()


async def _periodic_purge(protocols: list[BlindPairRelayProtocol]) -> None:
    publication_echouee_signalee = False
    while True:
        await asyncio.sleep(10)
        maintenant = time.time()
        for protocol in protocols:
            # Journalisé AVANT la purge : une session qui vient d'expirer a
            # encore son bilan à donner, et l'ordre inverse le perdrait.
            if protocol.compteur.paquets:
                logger.info("Port %d : %s", protocol.port, protocol.compteur.resume(maintenant))
            protocol.purge_expired(config.RELAY_SESSION_TTL_SECONDS)

        # Publié APRÈS la purge, donc en ne retenant que les ports réellement encore
        # appariés : c'est ce qu'app.py doit savoir pour maintenir son allocation en vie
        # (voir relay_activity, et config.RELAY_ALLOCATION_TTL_SECONDS pour le défaut
        # que ça corrige). Une session qui vient d'expirer ne doit surtout pas être
        # annoncée comme active, sinon son port resterait réservé indéfiniment.
        activite = {p.port: p.derniere_activite for p in protocols if p.derniere_activite is not None}
        if not relay_activity.publier(activite, config.RELAY_ACTIVITY_PATH):
            # Une seule fois : ce serait sinon une ligne toutes les 10 s pendant des
            # mois. Le relais fonctionne parfaitement sans, app.py retombant sur son
            # TTL long — c'est un avertissement, pas une panne.
            if not publication_echouee_signalee:
                logger.warning(
                    "Impossible de publier l'activité dans %s : app.py se rabattra sur "
                    "l'horodatage des demandes de relais. Vérifier les droits d'écriture "
                    "(unité systemd cloisonnée : ReadWritePaths) ou MESHAMOTO_RELAY_ACTIVITY_PATH.",
                    config.RELAY_ACTIVITY_PATH,
                )
                publication_echouee_signalee = True
        else:
            publication_echouee_signalee = False


async def main(
    host: str = "0.0.0.0",
    port_start: int = config.RELAY_PORT_POOL_START,
    pool_size: int = config.RELAY_PORT_POOL_SIZE,
) -> None:
    loop = asyncio.get_running_loop()
    protocols: list[BlindPairRelayProtocol] = []
    transports: list[asyncio.DatagramTransport] = []

    for port in range(port_start, port_start + pool_size):
        transport, protocol = await loop.create_datagram_endpoint(
            lambda p=port: BlindPairRelayProtocol(p), local_addr=(host, port)
        )
        transports.append(transport)
        protocols.append(protocol)

    logger.info(
        "Relais UDP en écoute sur %s, ports %d-%d (%d sessions simultanées possibles)",
        host, port_start, port_start + pool_size - 1, pool_size,
    )
    # Porteur WebSocket : monté dans CETTE boucle, et pas dans un fil à part, parce
    # qu'il manipule directement les protocoles ci-dessus. Un fil séparé exigerait un
    # verrou sur l'appariement, donc un second endroit où la règle des deux adresses
    # pourrait diverger.
    serveur_ws = None
    if config.RELAY_WS_ENABLED:
        serveur_ws = await relay_ws.servir(
            {p.port: p for p in protocols}, config.RELAY_WS_HOST, config.RELAY_WS_PORT
        )

    purge_task = asyncio.create_task(_periodic_purge(protocols))
    try:
        await asyncio.Future()  # tourne indéfiniment
    finally:
        purge_task.cancel()
        if serveur_ws is not None:
            serveur_ws.close()
        for transport in transports:
            transport.close()


if __name__ == "__main__":
    asyncio.run(main())
