"""Standalone relay server — thin WebSocket proxy for NAT traversal.

Agents connect outbound to the relay, register their PeerID, and the relay
bridges messages between connected peers. The relay sees only ciphertext
(Noise encryption is end-to-end).

Run: python -m a2a_p2p.relay_server --port 8765
"""

from __future__ import annotations

import argparse
import asyncio
import logging
import time
from collections import defaultdict

import websockets
import websockets.asyncio.server

logger = logging.getLogger(__name__)

# Registered agents: peer_id → websocket
_agents: dict[str, websockets.asyncio.server.ServerConnection] = {}
# Rate limiting: peer_id → (count, window_start)
_rate: dict[str, tuple[int, float]] = defaultdict(lambda: (0, time.time()))

MAX_CONNECTIONS = 100
MAX_MESSAGES_PER_MINUTE = 120


async def _handler(ws: websockets.asyncio.server.ServerConnection) -> None:
    """Handle a connected agent."""
    peer_id = ""
    remote = f"{ws.remote_address[0]}:{ws.remote_address[1]}" if ws.remote_address else "?"

    try:
        # First message must be registration: "REGISTER:<peer_id>"
        raw = await asyncio.wait_for(ws.recv(), timeout=10.0)
        if isinstance(raw, bytes):
            raw = raw.decode("utf-8", errors="replace")

        if not raw.startswith("REGISTER:"):
            await ws.close(4000, "Expected REGISTER:<peer_id>")
            return

        peer_id = raw[len("REGISTER:"):]
        if not peer_id or len(peer_id) > 100:
            await ws.close(4001, "Invalid peer_id")
            return

        if len(_agents) >= MAX_CONNECTIONS and peer_id not in _agents:
            await ws.close(4002, "Relay at capacity")
            return

        # Register
        old = _agents.get(peer_id)
        if old:
            try:
                await old.close(4003, "Replaced by new connection")
            except Exception:
                pass
        _agents[peer_id] = ws
        logger.info("Registered: %s from %s", peer_id, remote)
        await ws.send(f"OK:{peer_id}")

        # Relay loop
        async for message in ws:
            # Rate limit
            count, window_start = _rate[peer_id]
            now = time.time()
            if now - window_start > 60:
                count = 0
                window_start = now
            count += 1
            _rate[peer_id] = (count, window_start)
            if count > MAX_MESSAGES_PER_MINUTE:
                logger.warning("Rate limited: %s", peer_id)
                continue

            # Messages are: "TO:<target_peer_id>\x00<binary_data>"
            if isinstance(message, str):
                message = message.encode("utf-8")

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

            header = message[:sep].decode("utf-8", errors="replace")
            if not header.startswith("TO:"):
                continue

            target_id = header[3:]
            payload = message[sep + 1:]

            target_ws = _agents.get(target_id)
            if target_ws and target_ws.close_code is None:
                # Wrap: "FROM:<sender_peer_id>\x00<binary_data>"
                forwarded = f"FROM:{peer_id}".encode("utf-8") + b"\x00" + payload
                try:
                    await target_ws.send(forwarded)
                except Exception:
                    logger.debug("Failed to forward to %s", target_id)
            else:
                # Target not connected — send NACK
                await ws.send(f"NACK:{target_id}".encode("utf-8"))

    except asyncio.TimeoutError:
        logger.debug("Registration timeout from %s", remote)
    except websockets.ConnectionClosed:
        pass
    except Exception:
        logger.debug("Error handling %s", remote, exc_info=True)
    finally:
        if peer_id and _agents.get(peer_id) is ws:
            del _agents[peer_id]
            logger.info("Unregistered: %s", peer_id)


async def run_relay(host: str = "0.0.0.0", port: int = 8765) -> None:
    """Run the relay server."""
    logger.info("A2A relay starting on %s:%d", host, port)
    async with websockets.asyncio.server.serve(
        _handler, host, port,
        max_size=10 * 1024 * 1024,
    ):
        await asyncio.Future()  # Run forever


def main() -> None:
    """CLI entry point."""
    parser = argparse.ArgumentParser(description="A2A-P2P Relay Server")
    parser.add_argument("--host", default="0.0.0.0")
    parser.add_argument("--port", type=int, default=8765)
    parser.add_argument("--verbose", "-v", action="store_true")
    args = parser.parse_args()

    logging.basicConfig(
        level=logging.DEBUG if args.verbose else logging.INFO,
        format="%(asctime)s %(levelname)s %(name)s: %(message)s",
    )
    asyncio.run(run_relay(args.host, args.port))


if __name__ == "__main__":
    main()
