"""
nodelink.transport.relay_transport
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Relay transport — connects outward to a nodelink relay server instead of
accepting inbound connections directly.  Enables public hosting with zero
port forwarding or ngrok setup.
"""

import asyncio
import json
import logging
from typing import Any, Callable, Dict, Optional, Set

import websockets
from websockets.exceptions import ConnectionClosed

logger = logging.getLogger("nodelink.transport.relay")


class RelayTransport:
    """
    Transport that connects to a hosted nodelink relay server.

    Instead of binding a local port, the game server connects *outward*
    to the relay.  Players then connect to the relay with the game code and
    the relay forwards all traffic transparently.
    """

    def __init__(self, relay_url: str) -> None:
        # Always connect to the /host endpoint
        base = relay_url.rstrip("/")
        self._host_url = f"{base}/host"
        self._base_url = base

        self._ws = None
        self._connections: Dict[str, bool] = {}   # cid → connected

        self._on_message:    Optional[Callable] = None
        self._on_connect:    Optional[Callable] = None
        self._on_disconnect: Optional[Callable] = None

        self._game_code:   Optional[str]       = None
        self._ready:       Optional[asyncio.Event] = None   # created in start()
        self._start_error: Optional[Exception] = None
        self._task:        Optional[asyncio.Task]  = None

    # ── Registration hooks (mirrors WebSocketTransport API) ───────────────

    def on_message(self, fn: Callable) -> None:
        self._on_message = fn

    def on_connect(self, fn: Callable) -> None:
        self._on_connect = fn

    def on_disconnect(self, fn: Callable) -> None:
        self._on_disconnect = fn

    # ── Public info ───────────────────────────────────────────────────────

    @property
    def game_code(self) -> Optional[str]:
        """The game code assigned by the relay (available after start())."""
        return self._game_code

    @property
    def join_url(self) -> Optional[str]:
        """Full relay URL players connect to, e.g. ws://host:9000/join/ABC123."""
        if self._game_code:
            return f"{self._base_url}/join/{self._game_code}"
        return None

    # ── Lifecycle ─────────────────────────────────────────────────────────

    async def start(self) -> None:
        """Connect to relay and wait until the game code is received."""
        # Create the Event here so it always belongs to the running loop
        # (avoids Python 3.8/3.9 wrong-loop issues when created in __init__)
        self._ready = asyncio.Event()
        self._start_error = None
        self._task = asyncio.ensure_future(self._run())
        try:
            await asyncio.wait_for(asyncio.shield(self._ready.wait()), timeout=10.0)
        except asyncio.TimeoutError:
            self._task.cancel()
            raise RuntimeError(
                f"Could not connect to relay at {self._host_url}. "
                "Is the relay server running?"
            )
        if self._start_error:
            raise RuntimeError(
                f"Relay connection failed: {self._start_error}"
            )

    async def stop(self) -> None:
        if self._ws:
            await self._ws.close()
        if self._task:
            self._task.cancel()

    # ── Internal receive loop ─────────────────────────────────────────────

    async def _run(self) -> None:
        try:
            async with websockets.connect(self._host_url) as ws:
                self._ws = ws
                logger.info(f"Connected to relay at {self._host_url}")

                async for raw in ws:
                    try:
                        msg      = json.loads(raw)
                        event    = msg.get("event")
                        raw_data = msg.get("data")
                        data     = raw_data if raw_data is not None else {}

                        if event == "registered":
                            self._game_code = data["code"]
                            self._ready.set()
                            logger.info(f"Relay assigned game code: {self._game_code}")

                        elif event == "__client_connect__":
                            cid = data["cid"]
                            self._connections[cid] = True
                            if self._on_connect:
                                await self._on_connect(cid, data.get("address", "relay"))

                        elif event == "__client_message__":
                            cid = data["cid"]
                            if self._on_message:
                                await self._on_message(cid, data["payload"])

                        elif event == "__client_disconnect__":
                            cid = data["cid"]
                            self._connections.pop(cid, None)
                            if self._on_disconnect:
                                await self._on_disconnect(cid)

                    except Exception as e:
                        logger.error(f"Relay message error: {e}", exc_info=True)

        except ConnectionClosed:
            logger.info("Relay connection closed")
        except Exception as e:
            logger.error(f"Relay connection error: {e}", exc_info=True)
            if self._ready and not self._ready.is_set():
                self._start_error = e
                self._ready.set()   # unblock start() so it can re-raise the error
        finally:
            self._ws = None

    # ── Send / broadcast (mirrors WebSocketTransport API) ─────────────────

    async def send(self, cid: str, event: str, data: Any) -> bool:
        if not self._ws:
            return False
        try:
            await self._ws.send(json.dumps({
                "event": "__send__",
                "data": {
                    "cid":     cid,
                    "payload": {"event": event, "data": data},
                },
            }))
            return True
        except ConnectionClosed:
            return False
        except Exception as e:
            logger.error(f"Relay send error: {e}")
            return False

    async def broadcast(
        self,
        event: str,
        data: Any,
        exclude: Optional[Set[str]] = None,
    ) -> None:
        if not self._ws:
            return
        try:
            await self._ws.send(json.dumps({
                "event": "__broadcast__",
                "data": {
                    "payload": {"event": event, "data": data},
                    "exclude": list(exclude or []),
                },
            }))
        except ConnectionClosed:
            pass
        except Exception as e:
            logger.error(f"Relay broadcast error: {e}")

    # ── Info ──────────────────────────────────────────────────────────────

    @property
    def connection_count(self) -> int:
        return len(self._connections)

    def get_connection_ids(self) -> Set[str]:
        return set(self._connections.keys())
