"""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 threading
import uuid
from pathlib import Path

from .peer_registry import PeerRegistry

log = logging.getLogger(__name__)


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,
        )
        # 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:
                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:
        """Send one HTTP-shaped request to a connected peer and relay its response."""
        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

        try:
            raw = await asyncio.wait_for(ws_requester.recv(), timeout=30.0)
            payload = json.loads(raw)
        except Exception:
            await ws_requester.close(1003, "invalid relay rpc request")
            return

        request_id = str(payload.get("request_id") or uuid.uuid4())
        payload["request_id"] = request_id
        payload["target_peer"] = target_peer_id
        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
        try:
            await target.ws.send(json.dumps({
                "topic": "relay-http-request",
                "version": 1,
                "from_peer": "relay",
                "payload": payload,
            }))
            # Forward frames until a terminal one: streamed responses (US-036)
            # end with {"stream": true, "done": true}; a frame without "stream"
            # is a complete legacy single response.
            while True:
                remaining = deadline - loop.time()
                if remaining <= 0:
                    raise asyncio.TimeoutError
                frame = await asyncio.wait_for(
                    queue.get(), timeout=min(idle_timeout, remaining)
                )
                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)
            await ws_requester.close()


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