"""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>"  → 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:
    def __init__(self, ws, local_id: str, remote_id: str, relay: str) -> None:
        self._ws = ws
        self._local = local_id
        self._remote = remote_id
        self._relay = relay
        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._ws.close_code is None

    async def send(self, data: bytes) -> None:
        if not self.is_open:
            raise ConnectionError("closed")
        await self._ws.send(f"TO:{self._remote}".encode() + b"\x00" + data)

    async def recv(self) -> bytes:
        if not self.is_open:
            raise ConnectionError("closed")
        try:
            while True:
                raw = await self._ws.recv()
                if isinstance(raw, str):
                    raw = raw.encode()
                if raw.startswith(b"NACK:"):
                    raise PeerUnreachable(raw[5:].decode(), "relay: offline")
                try:
                    sep = raw.index(b"\x00")
                except ValueError:
                    continue
                hdr = raw[:sep].decode()
                if hdr.startswith("FROM:") and hdr[5:] == self._remote:
                    return raw[sep + 1:]
        except websockets.ConnectionClosed:
            self._closed = True
            raise ConnectionError("relay closed")

    async def close(self) -> None:
        self._closed = True
        try:
            await self._ws.close()
        except Exception:
            pass


class RelayClient:
    def __init__(self, local_peer_id: str, identity=None) -> None:
        self._local = local_peer_id
        self._identity = identity  # PeerIdentity, for signed registration
        self._relays: dict[str, websockets.asyncio.client.ClientConnection] = {}

    async def connect_relay(self, addr: str) -> None:
        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)

        # Signed registration: REGISTER:<peer_id>:<pubkey_hex>:<sig_hex>:<timestamp>
        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}")
        self._relays[addr] = ws

    async def connect_peer(self, relay: str, remote_id: str) -> RelayConnection:
        if relay not in self._relays:
            await self.connect_relay(relay)
        ws = self._relays[relay]
        if ws.close_code is not None:
            del self._relays[relay]
            await self.connect_relay(relay)
            ws = self._relays[relay]
        return RelayConnection(ws, self._local, remote_id, relay)

    async def close_all(self) -> None:
        for ws in self._relays.values():
            try:
                await ws.close()
            except Exception:
                pass
        self._relays.clear()
