"""Relay transport client — connects to a relay server for NAT traversal.

Both agents connect outbound to the relay server. The relay bridges
their WebSocket frames. Since Noise encryption is end-to-end, the
relay only sees ciphertext.

Relay protocol:
  - Client sends "REGISTER:<peer_id>" on connect
  - Server responds "OK:<peer_id>"
  - Client sends messages as "TO:<target_peer_id>\x00<data>"
  - Server forwards as "FROM:<sender_peer_id>\x00<data>"
  - Server sends "NACK:<peer_id>" if target is offline
"""

from __future__ import annotations

import asyncio
import logging
from urllib.parse import urlparse

import websockets
import websockets.asyncio.client

from a2a_p2p.errors import PeerUnreachable

logger = logging.getLogger(__name__)


class RelayConnection:
    """A connection to a peer via a relay server."""

    def __init__(
        self,
        ws: websockets.asyncio.client.ClientConnection,
        local_peer_id: str,
        remote_peer_id: str,
        relay_addr: str,
    ) -> None:
        self._ws = ws
        self._local_id = local_peer_id
        self._remote_id = remote_peer_id
        self._relay_addr = relay_addr
        self._closed = False

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

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

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

    async def send(self, data: bytes) -> None:
        """Send data to the remote peer via the relay."""
        if not self.is_open:
            raise ConnectionError("Relay connection closed")
        header = f"TO:{self._remote_id}".encode("utf-8")
        await self._ws.send(header + b"\x00" + data)

    async def recv(self) -> bytes:
        """Receive data from the remote peer via the relay."""
        if not self.is_open:
            raise ConnectionError("Relay connection closed")
        try:
            while True:
                raw = await self._ws.recv()
                if isinstance(raw, str):
                    raw = raw.encode("utf-8")

                # Parse relay framing
                if raw.startswith(b"NACK:"):
                    peer = raw[5:].decode("utf-8", errors="replace")
                    raise PeerUnreachable(peer, "relay says peer offline")

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

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

                if header.startswith("FROM:"):
                    sender = header[5:]
                    if sender == self._remote_id:
                        return payload
                    # Message from unexpected sender — skip
                    logger.debug(
                        "Relay: message from unexpected sender %s (expected %s)",
                        sender, self._remote_id,
                    )
        except websockets.ConnectionClosed:
            self._closed = True
            raise ConnectionError("Relay WebSocket closed")

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


class RelayClient:
    """Manages connections to relay servers."""

    def __init__(self, local_peer_id: str) -> None:
        self._local_id = local_peer_id
        self._relay_connections: dict[str, websockets.asyncio.client.ClientConnection] = {}

    async def connect_to_relay(self, relay_addr: str) -> None:
        """Register with a relay server."""
        ws_url = relay_addr
        parsed = urlparse(relay_addr)
        if parsed.scheme not in ("ws", "wss"):
            ws_url = f"ws://{relay_addr}"

        ws = await websockets.asyncio.client.connect(
            ws_url,
            max_size=10 * 1024 * 1024,
            open_timeout=10,
        )

        # Register
        await ws.send(f"REGISTER:{self._local_id}")
        resp = await asyncio.wait_for(ws.recv(), timeout=5.0)
        if isinstance(resp, bytes):
            resp = resp.decode("utf-8")
        if not resp.startswith("OK:"):
            await ws.close()
            raise ConnectionError(f"Relay registration failed: {resp}")

        self._relay_connections[relay_addr] = ws
        logger.info("Registered with relay %s as %s", relay_addr, self._local_id)

    async def connect_to_peer(
        self, relay_addr: str, remote_peer_id: str,
    ) -> RelayConnection:
        """Get a connection to a remote peer via a relay.

        If not already connected to the relay, connects first.
        """
        if relay_addr not in self._relay_connections:
            await self.connect_to_relay(relay_addr)

        ws = self._relay_connections[relay_addr]
        if ws.close_code is not None:
            # Reconnect
            del self._relay_connections[relay_addr]
            await self.connect_to_relay(relay_addr)
            ws = self._relay_connections[relay_addr]

        return RelayConnection(ws, self._local_id, remote_peer_id, relay_addr)

    async def close_all(self) -> None:
        """Close all relay connections."""
        for ws in self._relay_connections.values():
            try:
                await ws.close()
            except Exception:
                pass
        self._relay_connections.clear()
