"""
Make the node actually use every configured STUN server, not just the first.

`aiortc` accepts a list in `RTCConfiguration.iceServers` but keeps **only the
first** STUN URI (`connection_kwargs`: *"only a single STUN server is
supported"*), and `aioice.ice.Connection` has a single `stun_server` field.  So
the four defaults in `config.DEFAULT_STUN_SERVERS` — and anything the operator
adds on the Node page or with `meshbay-node stun add` — collapsed to
`stun:stun.l.google.com:19302` on the node side.  When that one server was slow
or unreachable from the node's network, `get_component_candidates` burned its
full 5 s timeout with no server-reflexive candidate to show for it, adding
seconds to every browser connection.  The "fallback" was configuration only.

This module patches `aioice.ice.server_reflexive_candidate` — the same
monkey-patch technique `ice_filter.py` uses on `get_host_addresses` — so a
single ICE gather races the STUN binding request against **all** configured
servers on the one bound socket and takes the first answer.  One reachable
server anywhere in the list now yields a reflexive candidate in one RTT instead
of a 5 s stall.

  install(["stun:a:3478", "stun:b:19302"])   → patch + set the list
  set_servers([...])                          → update the list live (CLI / panel)
  set_servers([]) / install(None)             → no fan-out; behave exactly like
                                                upstream with whatever single
                                                server aiortc passed
"""

import asyncio
import logging
import socket

import aioice.ice as _ice

log = logging.getLogger(__name__)

# Slightly under aioice's own `get_component_candidates(timeout=5)` so that when
# every server is unreachable we raise (→ no srflx candidate, same as today)
# just before the outer wait cancels us, rather than at the same instant.
_QUERY_TIMEOUT = 4.0

_servers: list[tuple[str, int]] = []
_installed = False
_original = _ice.server_reflexive_candidate


def _parse(url: str) -> tuple[str, int]:
    """`stun:host:port` / `stun:host` → (host, port). Default port 3478 (RFC 7064)."""
    rest = url[5:] if url.startswith("stun:") else url
    rest = rest.split("?", 1)[0].strip()          # drop any ?transport=... tail
    if not rest:
        raise ValueError(f"empty STUN host in {url!r}")
    host, sep, port = rest.rpartition(":")
    if sep and host and port.isdigit():
        return host, int(port)
    return rest, 3478


async def _query_one(protocol, addr: tuple[str, int]):
    # A fresh Message per call: StunProtocol.request asserts the transaction id
    # is not already registered, and each Message gets a random id at build time.
    request = _ice.stun.Message(
        message_method=_ice.stun.Method.BINDING,
        message_class=_ice.stun.Class.REQUEST,
    )
    response, _addr = await protocol.request(request, addr)
    return response


async def _first_ok(tasks: list[asyncio.Task]):
    """First task that returns without raising wins; if all raise, re-raise the last."""
    pending = set(tasks)
    last_exc: BaseException | None = None
    while pending:
        done, pending = await asyncio.wait(
            pending, return_when=asyncio.FIRST_COMPLETED)
        for task in done:
            exc = task.exception()
            if exc is None:
                return task.result()
            last_exc = exc
    raise last_exc if last_exc is not None else RuntimeError("no STUN query ran")


async def _fanout_server_reflexive_candidate(protocol, stun_server):
    """Drop-in for aioice.ice.server_reflexive_candidate that tries every server."""
    targets = list(_servers) if _servers else [tuple(stun_server)]

    loop = asyncio.get_event_loop()

    async def _resolve(host: str, port: int):
        return (await loop.run_in_executor(None, socket.gethostbyname, host), port)

    resolved: list[tuple[str, int]] = []
    for res in await asyncio.gather(
        *(_resolve(h, p) for h, p in targets), return_exceptions=True
    ):
        if isinstance(res, BaseException):
            log.debug("STUN resolve failed: %s", res)
        else:
            resolved.append(res)
    if not resolved:
        raise OSError("no configured STUN server could be resolved")

    query_tasks = [asyncio.ensure_future(_query_one(protocol, a)) for a in resolved]
    try:
        response = await asyncio.wait_for(_first_ok(query_tasks), _QUERY_TIMEOUT)
    finally:
        for task in query_tasks:
            task.cancel()
        await asyncio.gather(*query_tasks, return_exceptions=True)

    local = protocol.local_candidate
    return _ice.Candidate(
        foundation=_ice.candidate_foundation("srflx", "udp", local.host),
        component=local.component,
        transport=local.transport,
        priority=_ice.candidate_priority(local.component, "srflx"),
        host=response.attributes["XOR-MAPPED-ADDRESS"][0],
        port=response.attributes["XOR-MAPPED-ADDRESS"][1],
        type="srflx",
        related_address=local.host,
        related_port=local.port,
    ), None


def set_servers(stun_urls: list[str] | None) -> None:
    """Replace the fan-out server list. Safe to call before or after install()."""
    global _servers
    parsed: list[tuple[str, int]] = []
    for url in stun_urls or []:
        try:
            parsed.append(_parse(str(url)))
        except ValueError:
            log.warning("Ignoring malformed STUN URL: %r", url)
    _servers = parsed
    if _installed:
        log.info("STUN fan-out list: %s",
                 ", ".join(f"{h}:{p}" for h, p in parsed) or "(none)")


def install(stun_urls: list[str] | None) -> None:
    """Monkey-patch aioice so an ICE gather queries every configured STUN server."""
    global _installed
    set_servers(stun_urls)
    if not _installed:
        _ice.server_reflexive_candidate = _fanout_server_reflexive_candidate
        _installed = True
        log.info("STUN multi-server fan-out installed (%d servers)", len(_servers))