"""WebSocket relay for live state changes."""

import asyncio
import json
import logging

from fastapi import APIRouter, WebSocket

from config import settings

from .bark import _handle_client_message
from .state import _get_current_state

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/api/world", tags=["world"])

_ws_clients: list[WebSocket] = []

_HEARTBEAT_INTERVAL = 15  # seconds — ping to detect dead iPad/Safari connections


async def _heartbeat(websocket: WebSocket) -> None:
    """Send periodic pings to detect dead connections (iPad resilience).

    Safari suspends background tabs, killing the TCP socket silently.
    A 15-second ping ensures we notice within one interval.

    Rescued from stale PR #399.
    """
    try:
        while True:
            await asyncio.sleep(_HEARTBEAT_INTERVAL)
            await websocket.send_text(json.dumps({"type": "ping"}))
    except Exception:
        logger.debug("Heartbeat stopped — connection gone")


async def _authenticate_ws(websocket: WebSocket) -> bool:
    """Authenticate WebSocket connection using matrix_ws_token.

    Checks for token in query param ?token=<token>. If no query param,
    accepts the connection and waits for first message with
    {"type": "auth", "token": "<token>"}.

    Returns True if authenticated (or if auth is disabled).
    Returns False and closes connection with code 4001 if invalid.
    """
    token_setting = settings.matrix_ws_token

    # Auth disabled in dev mode (empty/unset token)
    if not token_setting:
        return True

    # Check query param first (can validate before accept)
    query_token = websocket.query_params.get("token", "")
    if query_token:
        if query_token == token_setting:
            return True
        # Invalid token in query param - we need to accept to close properly
        await websocket.accept()
        await websocket.close(code=4001, reason="Invalid token")
        return False

    # No query token - accept and wait for auth message
    await websocket.accept()

    # Wait for auth message as first message
    try:
        raw = await websocket.receive_text()
        data = json.loads(raw)
        if data.get("type") == "auth" and data.get("token") == token_setting:
            return True
        # Invalid auth message
        await websocket.close(code=4001, reason="Invalid token")
        return False
    except (json.JSONDecodeError, TypeError):
        # Non-JSON first message without valid token
        await websocket.close(code=4001, reason="Authentication required")
        return False


@router.websocket("/ws")
async def world_ws(websocket: WebSocket) -> None:
    """Accept a Workshop client and keep it alive for state broadcasts.

    Sends a full ``world_state`` snapshot immediately on connect so the
    client never starts from a blank slate.  Incoming frames are parsed
    as JSON — ``visitor_message`` triggers a bark response.  A background
    heartbeat ping runs every 15 s to detect dead connections early.

    Authentication:
        - If matrix_ws_token is configured, clients must provide it via
          ?token=<token> param or in the first message as
          {"type": "auth", "token": "<token>"}.
        - Invalid token results in close code 4001.
        - Valid token receives a connection_ack message.
    """
    # Authenticate (may accept connection internally)
    is_authed = await _authenticate_ws(websocket)
    if not is_authed:
        logger.info("World WS connection rejected — invalid token")
        return

    # Auth passed - accept if not already accepted
    if websocket.client_state.name != "CONNECTED":
        await websocket.accept()

    # Send connection_ack if auth was required
    if settings.matrix_ws_token:
        await websocket.send_text(json.dumps({"type": "connection_ack"}))

    _ws_clients.append(websocket)
    logger.info("World WS connected — %d clients", len(_ws_clients))

    # Send full world-state snapshot so client bootstraps instantly
    try:
        snapshot = _get_current_state()
        await websocket.send_text(json.dumps({"type": "world_state", **snapshot}))
    except Exception as exc:
        logger.warning("Failed to send WS snapshot: %s", exc)

    ping_task = asyncio.create_task(_heartbeat(websocket))
    try:
        while True:
            raw = await websocket.receive_text()
            await _handle_client_message(raw)
    except Exception:
        logger.debug("WebSocket receive loop ended")
    finally:
        ping_task.cancel()
        if websocket in _ws_clients:
            _ws_clients.remove(websocket)
        logger.info("World WS disconnected — %d clients", len(_ws_clients))


async def _broadcast(message: str) -> None:
    """Send *message* to every connected Workshop client, pruning dead ones."""
    dead: list[WebSocket] = []
    for ws in _ws_clients:
        try:
            await ws.send_text(message)
        except Exception:
            logger.debug("Pruning dead WebSocket client")
            dead.append(ws)
    for ws in dead:
        if ws in _ws_clients:
            _ws_clients.remove(ws)


async def broadcast_world_state(presence: dict) -> None:
    """Broadcast a ``timmy_state`` message to all connected Workshop clients.

    Called by :class:`~timmy.workshop_state.WorkshopHeartbeat` via its
    ``on_change`` callback.
    """
    from .state import _build_world_state

    state = _build_world_state(presence)
    await _broadcast(json.dumps({"type": "timmy_state", **state["timmyState"]}))
