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

Security model:
- Signed registration: agents must prove they own their PeerID via Ed25519
- TOFU key pinning: first registration binds peer_id → public key
- Per-IP connection limits: prevent single-source exhaustion
- Rate limiting per peer_id with automatic cleanup
- Message size bounded by websockets max_size

Run: python -m 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__)

# ── Limits ──────────────────────────────────────────────────────────

MAX_CONNECTIONS = 200
MAX_CONNECTIONS_PER_IP = 10
MAX_MSG_PER_MINUTE = 120
RATE_CLEANUP_INTERVAL = 300  # clean stale rate entries every 5min
MAX_PEER_ID_LEN = 100

# ── State ───────────────────────────────────────────────────────────

# peer_id → websocket
_agents: dict[str, websockets.asyncio.server.ServerConnection] = {}
# peer_id → pinned public key bytes (TOFU)
_pinned_keys: dict[str, bytes] = {}
# IP → count of active connections
_ip_conns: dict[str, int] = defaultdict(int)
# peer_id → (message_count, window_start)
_rate: dict[str, tuple[int, float]] = {}
# peer_id → last_active timestamp (for cleanup)
_last_active: dict[str, float] = {}


def _verify_registration(peer_id: str, pubkey_hex: str, signature_hex: str, timestamp: float = 0) -> bool:
    """Verify that the registrant owns the claimed peer_id."""
    try:
        from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey
        from cryptography.hazmat.primitives.serialization import Encoding, PublicFormat
        import hashlib

        pubkey_bytes = bytes.fromhex(pubkey_hex)
        if len(pubkey_bytes) != 32:
            return False

        pubkey = Ed25519PublicKey.from_public_bytes(pubkey_bytes)
        signature = bytes.fromhex(signature_hex)

        # Verify the peer_id matches the public key
        # PeerID = p2p:<base58(sha256(pubkey)[:20])>
        from p2p.identity import _b58encode
        expected_id = f"p2p:{_b58encode(hashlib.sha256(pubkey_bytes).digest()[:20])}"
        if peer_id != expected_id:
            logger.warning("PeerID mismatch: claimed %s, derived %s", peer_id, expected_id)
            return False

        # Verify timestamp freshness (60-second window)
        if abs(time.time() - timestamp) > 60:
            logger.warning("Stale registration from %s (ts=%d)", peer_id, timestamp)
            return False

        # Verify signature over "REGISTER:<peer_id>:<timestamp>"
        pubkey.verify(signature, f"REGISTER:{peer_id}:{int(timestamp)}".encode())

        # TOFU: atomic pin-on-first-use, reject different keys later
        existing = _pinned_keys.setdefault(peer_id, pubkey_bytes)
        if existing != pubkey_bytes:
            logger.warning("TOFU violation for %s — different key", peer_id)
            return False

        return True
    except Exception:
        logger.debug("Registration verification failed for %s", peer_id, exc_info=True)
        return False


def _get_ip(ws) -> str:
    try:
        return ws.remote_address[0] if ws.remote_address else "?"
    except Exception:
        return "?"


async def _cleanup_rate_state() -> None:
    """Periodically remove rate/activity entries for disconnected peers."""
    while True:
        await asyncio.sleep(RATE_CLEANUP_INTERVAL)
        now = time.time()
        stale = [
            pid for pid, ts in _last_active.items()
            if now - ts > RATE_CLEANUP_INTERVAL and pid not in _agents
        ]
        for pid in stale:
            _rate.pop(pid, None)
            _last_active.pop(pid, None)
        if stale:
            logger.debug("Cleaned up %d stale rate entries", len(stale))


async def _handler(ws: websockets.asyncio.server.ServerConnection) -> None:
    peer_id = ""
    ip = _get_ip(ws)

    # Per-IP connection limit
    if _ip_conns[ip] >= MAX_CONNECTIONS_PER_IP:
        logger.warning("IP %s exceeded connection limit (%d)", ip, MAX_CONNECTIONS_PER_IP)
        await ws.close(4002, "Too many connections from this IP")
        return

    _ip_conns[ip] += 1

    try:
        # Registration: "REGISTER:<peer_id>:<pubkey_hex>:<signature_hex>"
        raw = await asyncio.wait_for(ws.recv(), timeout=10)
        if isinstance(raw, bytes):
            raw = raw.decode("utf-8", errors="replace")

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

        # Format: REGISTER:<peer_id>:<pubkey_hex>:<sig_hex>:<timestamp>
        # PeerID contains a colon (p2p:XXXX), so split from the right
        # where the last 3 fields (pubkey, sig, ts) are fixed-format.
        parts = raw[9:].rsplit(":", 3)
        if len(parts) != 4:
            await ws.close(4001, "Expected REGISTER:<id>:<pubkey>:<sig>:<timestamp>")
            return

        peer_id, pubkey_hex, sig_hex, ts_str = parts

        if not peer_id or len(peer_id) > MAX_PEER_ID_LEN:
            await ws.close(4001, "Bad peer_id")
            return

        try:
            timestamp = float(ts_str)
        except ValueError:
            await ws.close(4001, "Bad timestamp")
            return

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

        # Verify ownership + timestamp freshness
        if not _verify_registration(peer_id, pubkey_hex, sig_hex, timestamp):
            await ws.close(4003, "Authentication failed")
            return

        # Evict previous connection for this peer_id (reconnection)
        old = _agents.get(peer_id)
        if old:
            try:
                await old.close(4004, "Replaced by new connection")
            except Exception:
                pass

        _agents[peer_id] = ws
        _last_active[peer_id] = time.time()
        logger.info("+ %s from %s", peer_id, ip)
        await ws.send(f"OK:{peer_id}")

        # Relay loop
        async for msg in ws:
            # Rate limiting
            now = time.time()
            cnt, start = _rate.get(peer_id, (0, now))
            if now - start > 60:
                cnt, start = 0, now
            cnt += 1
            _rate[peer_id] = (cnt, start)
            _last_active[peer_id] = now

            if cnt > MAX_MSG_PER_MINUTE:
                continue  # silently drop

            if isinstance(msg, str):
                msg = msg.encode("utf-8")

            # Parse: "TO:<target>\x00<payload>"
            try:
                sep = msg.index(b"\x00")
            except ValueError:
                continue

            hdr = msg[:sep]
            if len(hdr) > 200:  # sanity bound on header length
                continue

            hdr_str = hdr.decode("utf-8", errors="replace")
            if not hdr_str.startswith("TO:"):
                continue

            target = hdr_str[3:]
            payload = msg[sep + 1:]

            tws = _agents.get(target)
            if tws and tws.close_code is None:
                forwarded = f"FROM:{peer_id}".encode("utf-8") + b"\x00" + payload
                try:
                    await tws.send(forwarded)
                except Exception:
                    pass
            else:
                await ws.send(f"NACK:{target}".encode("utf-8"))

    except asyncio.TimeoutError:
        logger.debug("Registration timeout from %s", ip)
    except websockets.ConnectionClosed:
        pass
    except Exception:
        logger.debug("Handler error from %s", ip, exc_info=True)
    finally:
        _ip_conns[ip] = max(0, _ip_conns[ip] - 1)
        if _ip_conns[ip] == 0:
            del _ip_conns[ip]
        if peer_id and _agents.get(peer_id) is ws:
            del _agents[peer_id]
            logger.info("- %s", peer_id)


async def run_relay(host: str = "0.0.0.0", port: int = 8765) -> None:
    logger.info("Relay on %s:%d (max %d conns, %d/IP)", host, port, MAX_CONNECTIONS, MAX_CONNECTIONS_PER_IP)
    cleanup_task = asyncio.create_task(_cleanup_rate_state())
    try:
        async with websockets.asyncio.server.serve(
            _handler, host, port, max_size=10 * 1024 * 1024,
        ):
            await asyncio.Future()
    finally:
        cleanup_task.cancel()


def main() -> None:
    p = argparse.ArgumentParser(description="OpenClaw P2P Relay")
    p.add_argument("--host", default="0.0.0.0")
    p.add_argument("--port", type=int, default=8765)
    p.add_argument("-v", action="store_true")
    a = p.parse_args()
    logging.basicConfig(
        level=logging.DEBUG if a.v else logging.INFO,
        format="%(asctime)s %(levelname)s %(name)s: %(message)s",
    )
    asyncio.run(run_relay(a.host, a.port))


if __name__ == "__main__":
    main()
