"""
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.

    If the relay connection drops mid-session, RelayTransport automatically
    reconnects with exponential backoff (2 s → 4 s → … → 30 s max).
    """

    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.create_task(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 (with reconnect) ────────────────────────────

    async def _run(self) -> None:
        first_connect = True
        reconnect_delay = 2.0

        while True:
            self._ws = None
            try:
                async with websockets.connect(self._host_url) as ws:
                    self._ws = ws
                    reconnect_delay = 2.0  # reset backoff on successful connect
                    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"]
                                if first_connect:
                                    first_connect = False
                                    self._ready.set()
                                else:
                                    logger.warning(
                                        f"Relay reconnected — new game code: {self._game_code}"
                                    )
                                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.warning("Relay connection lost, reconnecting...")
            except Exception as e:
                logger.error(f"Relay connection error: {e}", exc_info=True)
                if first_connect and self._ready and not self._ready.is_set():
                    # Initial connection failed — surface the error and stop
                    self._start_error = e
                    self._ready.set()
                    return
            finally:
                self._ws = None

            # Fire disconnect for all tracked connections so the server cleans up
            disconnected = list(self._connections.keys())
            self._connections.clear()
            for cid in disconnected:
                if self._on_disconnect:
                    try:
                        await self._on_disconnect(cid)
                    except Exception:
                        pass

            logger.info(f"Reconnecting to relay in {reconnect_delay:.1f}s...")
            await asyncio.sleep(reconnect_delay)
            reconnect_delay = min(reconnect_delay * 2, 30.0)

    # ── 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 close(self, cid: str) -> None:
        """Ask the relay to close a specific client connection (kick)."""
        if not self._ws:
            return
        try:
            await self._ws.send(json.dumps({
                "event": "__kick__",
                "data": {"cid": cid},
            }))
        except Exception as e:
            logger.error(f"Relay kick error: {e}")

    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())
