"""Relay server — WebSocket gossip hub + circuit relay proxy.

HTTP API (served via asyncio-based handler on same port):
  GET /health              → {"status": "ok", "peers": N}
  GET /v1/peers            → [{peer_id, addr, last_seen}]
  POST /v1/gossip          → accept a gossip envelope, fan out to connected peers

WebSocket endpoints:
  ws[s]://host:port/ws              → gossip PubSub connection
  ws[s]://host:port/relay/{peer_id} → circuit relay to that peer
"""

from __future__ import annotations

import asyncio
import json
import logging
import os
import threading
import uuid
from pathlib import Path

from .peer_registry import PeerRegistry

log = logging.getLogger(__name__)

# Activation tensors ride the relay as base64 inside one JSON frame, so the
# websockets default of 1 MiB rejects any real prefill (close code 1009).
DEFAULT_WS_MAX_BYTES = 256 * 1024 * 1024


def ws_max_size() -> int | None:
    """Max inbound WebSocket frame size; MESHNET_WS_MAX_BYTES<=0 means unlimited."""
    raw = os.environ.get("MESHNET_WS_MAX_BYTES", "").strip()
    if not raw:
        return DEFAULT_WS_MAX_BYTES
    try:
        value = int(raw)
    except ValueError:
        return DEFAULT_WS_MAX_BYTES
    return None if value <= 0 else value


# Binary relay frame: JSON header (request/response metadata) + raw body in one
# WebSocket binary message. Activation bodies stay bytes end to end — no base64
# inflation, no JSON re-encode of megabytes per hop. Text JSON frames remain the
# control plane (gossip, peer-register, streamed SSE responses).
BINARY_FRAME_MAGIC = b"MRF1"


def encode_binary_frame(header: dict, body: bytes) -> bytes:
    header_bytes = json.dumps(header, separators=(",", ":")).encode()
    return BINARY_FRAME_MAGIC + len(header_bytes).to_bytes(4, "big") + header_bytes + body


def decode_binary_frame(frame: bytes) -> tuple[dict, bytes]:
    if len(frame) < 8 or frame[:4] != BINARY_FRAME_MAGIC:
        raise ValueError("not a meshnet binary relay frame")
    header_len = int.from_bytes(frame[4:8], "big")
    header = json.loads(frame[8:8 + header_len].decode())
    return header, bytes(frame[8 + header_len:])


class RelayServer:
    """Async WebSocket relay server that runs in a background thread.

    Usage::

        server = RelayServer(host="0.0.0.0", port=8765)
        port = server.start()   # returns actual port
        ...
        server.stop()
    """

    def __init__(
        self,
        host: str = "0.0.0.0",
        port: int = 8765,
        ssl_cert: Path | None = None,
        ssl_key: Path | None = None,
        max_peers: int = 500,
    ):
        self.host = host
        self.port = port
        self.ssl_cert = ssl_cert
        self.ssl_key = ssl_key
        self.max_peers = max_peers

        self._registry = PeerRegistry()
        self._loop: asyncio.AbstractEventLoop | None = None
        self._thread: threading.Thread | None = None
        self._server = None
        self._actual_port: int = port
        self._ready = threading.Event()
        self._running = False
        self._stop_event: asyncio.Event | None = None
        # request_id → queue of relay-http-response frames (US-036: a streamed
        # response is a sequence of frames; a frame without "stream" is terminal).
        self._pending_rpc: dict[str, asyncio.Queue] = {}

    @property
    def registry(self) -> PeerRegistry:
        return self._registry

    def start(self) -> int:
        """Start server in background thread. Returns actual bound port."""
        self._running = True
        self._loop = asyncio.new_event_loop()
        self._thread = threading.Thread(target=self._run, daemon=True, name="relay")
        self._thread.start()
        self._ready.wait(timeout=5)
        return self._actual_port

    def stop(self) -> None:
        self._running = False
        if self._loop and self._stop_event is not None:
            self._loop.call_soon_threadsafe(self._stop_event.set)
        if self._thread:
            self._thread.join(timeout=3.0)

    def _run(self) -> None:
        asyncio.set_event_loop(self._loop)
        self._loop.run_until_complete(self._serve())

    async def _serve(self) -> None:
        import websockets  # type: ignore[import]
        import websockets.server  # type: ignore[import]

        ssl_ctx = None
        if self.ssl_cert and self.ssl_key:
            import ssl
            ssl_ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
            ssl_ctx.load_cert_chain(str(self.ssl_cert), str(self.ssl_key))

        server = await websockets.serve(
            self._handle_connection,
            self.host,
            self.port,
            ssl=ssl_ctx,
            max_size=ws_max_size(),
            # Bulk payloads are zstd-compressed at the pipeline layer already;
            # per-message deflate would recompress them on every hop for nothing.
            compression=None,
        )
        # Record actual port after bind
        for sock in server.sockets or []:
            self._actual_port = sock.getsockname()[1]
            break

        self._stop_event = asyncio.Event()
        self._server = server
        self._ready.set()
        log.info("Relay listening on %s:%d", self.host, self._actual_port)

        await self._stop_event.wait()
        server.close()
        await server.wait_closed()

    async def _handle_connection(self, ws) -> None:
        """Dispatch incoming WebSocket to gossip hub or circuit relay."""
        try:
            path = ws.request.path
        except AttributeError:
            path = getattr(ws, "path", "/ws")

        if path.startswith("/relay/"):
            peer_id = path[len("/relay/"):]
            await self._handle_circuit_relay(ws, peer_id)
        elif path.startswith("/rpc/"):
            peer_id = path[len("/rpc/"):]
            await self._handle_rpc(ws, peer_id)
        elif path == "/health":
            await ws.send(json.dumps({"status": "ok", "peers": len(self._registry)}))
            await ws.close()
        elif path == "/v1/peers":
            await ws.send(json.dumps(self._registry.list_peers()))
            await ws.close()
        else:
            await self._handle_gossip(ws)

    async def _handle_gossip(self, ws) -> None:
        """Accept a gossip peer connection, register it, and fan out messages."""
        peer_id: str | None = None
        peer_addr: str = ""

        try:
            async for raw in ws:
                # Binary frames are relay-http-response bodies from a bridge —
                # route them to the waiting rpc requester as-is, never fan out.
                if isinstance(raw, (bytes, bytearray)):
                    try:
                        header, _ = decode_binary_frame(bytes(raw))
                    except (ValueError, json.JSONDecodeError):
                        continue
                    queue = self._pending_rpc.get(header.get("request_id"))
                    if queue is not None:
                        queue.put_nowait(bytes(raw))
                    continue
                try:
                    envelope = json.loads(raw)
                except (json.JSONDecodeError, TypeError):
                    continue

                topic = envelope.get("topic", "")
                from_peer = envelope.get("from_peer", "")

                # Handle peer registration message
                if topic == "peer-register":
                    payload = envelope.get("payload", {})
                    peer_id = payload.get("peer_id") or from_peer
                    peer_addr = payload.get("addr", "")
                    if len(self._registry) >= self.max_peers:
                        await ws.close(1008, "relay at capacity")
                        return
                    self._registry.register(peer_id, peer_addr, ws)
                    log.debug("Peer registered: %s", peer_id)
                    # Send current peer list back
                    await ws.send(json.dumps({
                        "topic": "peer-list",
                        "version": 1,
                        "from_peer": "relay",
                        "payload": {"peers": self._registry.list_peers()},
                    }))
                    continue

                if topic == "relay-http-response":
                    payload = envelope.get("payload", {})
                    request_id = payload.get("request_id")
                    queue = self._pending_rpc.get(request_id)
                    if queue is not None:
                        queue.put_nowait(payload)
                    continue

                # Fan out to all other registered peers
                if peer_id:
                    self._registry.touch(peer_id)
                fan_out_peers = self._registry.all_except(peer_id or "")
                await _broadcast(raw, fan_out_peers)

        except Exception as exc:
            log.debug("Gossip connection error: %s", exc)
        finally:
            if peer_id:
                self._registry.unregister(peer_id)
                log.debug("Peer unregistered: %s", peer_id)

    async def _handle_circuit_relay(self, ws_requester, target_peer_id: str) -> None:
        """Proxy WebSocket traffic between ws_requester and target_peer_id's ws."""
        target = self._registry.get(target_peer_id)
        if target is None:
            try:
                await ws_requester.send(json.dumps({
                    "error": f"peer {target_peer_id!r} not connected to relay"
                }))
                await ws_requester.close(1011, "target peer not found")
            except Exception:
                pass
            return

        log.debug("Circuit relay: ??? → %s", target_peer_id)

        async def pipe(src, dst) -> None:
            try:
                async for msg in src:
                    await dst.send(msg)
            except Exception:
                pass

        await asyncio.gather(
            pipe(ws_requester, target.ws),
            pipe(target.ws, ws_requester),
            return_exceptions=True,
        )

    async def _handle_rpc(self, ws_requester, target_peer_id: str) -> None:
        """Relay sequential HTTP-shaped requests over one requester connection."""
        target = self._registry.get(target_peer_id)
        if target is None:
            await ws_requester.send(json.dumps({
                "status": 503,
                "headers": {"Content-Type": "application/json"},
                "body": json.dumps({"error": f"peer {target_peer_id!r} not connected to relay"}),
            }))
            await ws_requester.close()
            return

        while True:
            try:
                raw = await ws_requester.recv()
            except Exception:
                return

            request_id: str | None = None
            try:
                if isinstance(raw, (bytes, bytearray)):
                    header, body = decode_binary_frame(bytes(raw))
                    request_id = str(header.get("request_id") or uuid.uuid4())
                    header["request_id"] = request_id
                    header["target_peer"] = target_peer_id
                    outbound: str | bytes = encode_binary_frame(header, body)
                else:
                    payload = json.loads(raw)
                    request_id = str(payload.get("request_id") or uuid.uuid4())
                    payload["request_id"] = request_id
                    payload["target_peer"] = target_peer_id
                    outbound = json.dumps({
                        "topic": "relay-http-request",
                        "version": 1,
                        "from_peer": "relay",
                        "payload": payload,
                    })
            except Exception:
                await ws_requester.close(1003, "invalid relay rpc request")
                return

            queue: asyncio.Queue = asyncio.Queue()
            self._pending_rpc[request_id] = queue
            overall_timeout = 310.0
            idle_timeout = 120.0
            loop = asyncio.get_running_loop()
            deadline = loop.time() + overall_timeout
            target = self._registry.get(target_peer_id)
            try:
                if target is None:
                    await ws_requester.send(json.dumps({
                        "request_id": request_id,
                        "status": 503,
                        "headers": {"Content-Type": "application/json"},
                        "body": json.dumps({"error": f"peer {target_peer_id!r} disconnected"}),
                    }))
                    continue
                await target.ws.send(outbound)
                # Streamed responses end with done=true. Binary and legacy JSON
                # responses are complete in one frame.
                while True:
                    remaining = deadline - loop.time()
                    if remaining <= 0:
                        raise asyncio.TimeoutError
                    frame = await asyncio.wait_for(
                        queue.get(), timeout=min(idle_timeout, remaining)
                    )
                    if isinstance(frame, (bytes, bytearray)):
                        await ws_requester.send(frame)
                        break
                    await ws_requester.send(json.dumps(frame))
                    if not frame.get("stream") or frame.get("done"):
                        break
            except asyncio.TimeoutError:
                await ws_requester.send(json.dumps({
                    "request_id": request_id,
                    "status": 504,
                    "headers": {"Content-Type": "application/json"},
                    "body": json.dumps({"error": "relay rpc timed out"}),
                }))
            finally:
                self._pending_rpc.pop(request_id, None)


async def _broadcast(raw: str | bytes, peers: list) -> None:
    """Send raw message to all peers; ignore individual send failures."""
    if not peers:
        return
    import asyncio
    await asyncio.gather(
        *[_safe_send(p.ws, raw) for p in peers],
        return_exceptions=True,
    )


async def _safe_send(ws, msg) -> None:
    try:
        await ws.send(msg)
    except Exception:
        pass