"""Relay transport client for NAT traversal.

Both agents connect outbound to a relay. The relay bridges their frames.
Noise encryption is E2E — relay sees only ciphertext.

Protocol:
  Client → "REGISTER:<peer_id>:<pubkey>:<sig>:<ts>"  → Server responds "OK:<peer_id>"
  Client → "TO:<target>\\x00<data>" → Server forwards as "FROM:<sender>\\x00<data>"
  Server → "NACK:<peer_id>" if target offline
"""

from __future__ import annotations

import asyncio
import logging
from urllib.parse import urlparse

import websockets
import websockets.asyncio.client

from p2p.errors import PeerUnreachable

logger = logging.getLogger(__name__)


class RelayConnection:
    """A virtual connection to a single peer routed through a RelayMux."""

    def __init__(self, mux: RelayMux, remote_id: str, relay: str) -> None:
        self._mux = mux
        self._remote = remote_id
        self._relay = relay
        self._inbox: asyncio.Queue[bytes] = asyncio.Queue()
        self._closed = False

    @property
    def peer_id(self) -> str:
        return self._remote

    @property
    def remote_addr(self) -> str:
        return f"relay:{self._relay}"

    @property
    def is_open(self) -> bool:
        return not self._closed and self._mux.is_open

    async def send(self, data: bytes) -> None:
        if not self.is_open:
            raise ConnectionError("closed")
        await self._mux.send_to(self._remote, data)

    async def recv(self) -> bytes:
        if not self.is_open:
            raise ConnectionError("closed")
        try:
            return await self._inbox.get()
        except asyncio.CancelledError:
            self._closed = True
            raise ConnectionError("cancelled")

    def deliver(self, data: bytes) -> None:
        """Called by the mux to deliver a message to this connection's inbox."""
        try:
            self._inbox.put_nowait(data)
        except asyncio.QueueFull:
            logger.warning("Relay inbox full for %s, dropping", self._remote)

    async def close(self) -> None:
        self._closed = True


class RelayMux:
    """Multiplexer for a single relay WebSocket.

    Reads all incoming messages, demuxes by sender PeerID, and routes
    to per-peer RelayConnection inboxes. Creates new connections for
    unknown senders (via a callback).
    """

    def __init__(
        self,
        ws: websockets.asyncio.client.ClientConnection,
        local_peer_id: str,
        relay_addr: str,
        on_new_peer: asyncio.coroutines | None = None,
    ) -> None:
        self._ws = ws
        self._local = local_peer_id
        self._relay_addr = relay_addr
        self._on_new_peer = on_new_peer
        self._peers: dict[str, RelayConnection] = {}
        self._reader_task: asyncio.Task | None = None

    @property
    def is_open(self) -> bool:
        return self._ws.close_code is None

    def start(self) -> None:
        """Start the background reader loop."""
        self._reader_task = asyncio.create_task(
            self._read_loop(), name="p2p:relay-mux",
        )

    def get_connection(self, remote_id: str) -> RelayConnection:
        """Get or create a virtual connection for a peer."""
        if remote_id not in self._peers:
            conn = RelayConnection(self, remote_id, self._relay_addr)
            self._peers[remote_id] = conn
        return self._peers[remote_id]

    async def send_to(self, remote_id: str, data: bytes) -> None:
        """Send data to a peer through the relay."""
        if not self.is_open:
            raise ConnectionError("relay closed")
        await self._ws.send(f"TO:{remote_id}".encode() + b"\x00" + data)

    async def close(self) -> None:
        if self._reader_task:
            self._reader_task.cancel()
            try:
                await self._reader_task
            except asyncio.CancelledError:
                pass
        try:
            await self._ws.close()
        except Exception:
            pass

    async def _read_loop(self) -> None:
        """Read relay messages and route to per-peer inboxes."""
        try:
            while self.is_open:
                try:
                    raw = await self._ws.recv()
                except websockets.ConnectionClosed:
                    break

                if isinstance(raw, str):
                    raw = raw.encode()

                if raw.startswith(b"NACK:"):
                    peer = raw[5:].decode("utf-8", errors="replace")
                    logger.debug("Relay NACK for %s", peer)
                    continue

                try:
                    sep = raw.index(b"\x00")
                except ValueError:
                    continue

                hdr = raw[:sep].decode("utf-8", errors="replace")
                payload = raw[sep + 1:]

                if not hdr.startswith("FROM:"):
                    continue

                sender_id = hdr[5:]
                conn = self._peers.get(sender_id)

                if conn is None:
                    # New peer — create connection, deliver first message,
                    # THEN fire callback as a task (not awaited) to avoid
                    # deadlock: the callback needs to recv() from the inbox,
                    # which requires this reader loop to keep running.
                    conn = self.get_connection(sender_id)
                    conn.deliver(payload)
                    if self._on_new_peer:
                        asyncio.create_task(self._on_new_peer(conn))
                    continue

                conn.deliver(payload)

        except asyncio.CancelledError:
            pass
        except Exception:
            logger.exception("Relay mux reader crashed")


class RelayClient:
    def __init__(self, local_peer_id: str, identity=None) -> None:
        self._local = local_peer_id
        self._identity = identity
        self._muxes: dict[str, RelayMux] = {}

    async def connect_relay(self, addr: str, on_new_peer=None) -> RelayMux:
        """Connect and register with a relay. Returns the RelayMux."""
        url = addr if urlparse(addr).scheme in ("ws", "wss") else f"ws://{addr}"
        ws = await websockets.asyncio.client.connect(url, max_size=10*1024*1024, open_timeout=10)

        import time as _time
        if self._identity:
            ts = int(_time.time())
            pubkey_hex = self._identity.public_key_raw().hex()
            sig = self._identity.sign(f"REGISTER:{self._local}:{ts}".encode())
            reg_msg = f"REGISTER:{self._local}:{pubkey_hex}:{sig.hex()}:{ts}"
        else:
            raise ConnectionError("Identity required for relay registration")

        await ws.send(reg_msg)
        resp = await asyncio.wait_for(ws.recv(), timeout=5)
        if isinstance(resp, bytes):
            resp = resp.decode()
        if not resp.startswith("OK:"):
            await ws.close()
            raise ConnectionError(f"Relay register failed: {resp}")

        mux = RelayMux(ws, self._local, addr, on_new_peer=on_new_peer)
        mux.start()
        self._muxes[addr] = mux
        return mux

    def get_mux(self, addr: str) -> RelayMux | None:
        return self._muxes.get(addr)

    async def close_all(self) -> None:
        for mux in self._muxes.values():
            await mux.close()
        self._muxes.clear()
