from fastapi import APIRouter, WebSocket, WebSocketDisconnect

from led_platform.core.clock import epoch_ms
from led_platform.domain import CommandEnvelope, CommandType, PerformanceSample
from led_platform.runtime import hub, scenes, settings, sync, tiles

router = APIRouter()


def initial_state_payload() -> dict:
    apply_at_ms = epoch_ms() + settings.initial_state_delay_ms
    state = scenes.state().model_copy(update={"apply_at_ms": apply_at_ms})
    return {
        "state": state.model_dump(mode="json"),
        "server_time_ms": epoch_ms(),
        "apply_at_ms": apply_at_ms,
    }


@router.websocket("/ws/admin")
async def admin_ws(websocket: WebSocket) -> None:
    await hub.connect("admin", websocket)
    try:
        await hub.send(websocket, CommandEnvelope(type=CommandType.STATE, payload=initial_state_payload()))
        while True:
            raw = await websocket.receive_json()
            if raw.get("type") == CommandType.CLOCK_PING:
                await hub.send(
                    websocket,
                    CommandEnvelope(
                        type=CommandType.CLOCK_PONG,
                        payload={
                            "client_send_ms": raw.get("client_send_ms"),
                            "client_send_perf_ms": raw.get("client_send_perf_ms"),
                            "server_time_ms": epoch_ms(),
                        },
                    ),
                )
    except WebSocketDisconnect:
        await hub.disconnect("admin", websocket)


@router.websocket("/ws/output/{tile_id}")
async def output_ws(websocket: WebSocket, tile_id: str) -> None:
    try:
        tiles.get(tile_id)
    except KeyError:
        await websocket.close(code=1008)
        return

    node_id: str | None = None
    await hub.connect("outputs", websocket)
    try:
        await hub.send(websocket, CommandEnvelope(type=CommandType.STATE, payload=initial_state_payload()))
        while True:
            raw = await websocket.receive_json()
            message_type = raw.get("type")

            if message_type == CommandType.HELLO:
                node_id = str(raw.get("node_id") or f"{tile_id}-anonymous")
                node = sync.connect_node(
                    node_id=node_id,
                    tile_id=tile_id,
                    app_version=raw.get("app_version"),
                    user_agent=raw.get("user_agent"),
                )
                await hub.broadcast(
                    CommandEnvelope(
                        type=CommandType.HEARTBEAT,
                        payload=node.model_dump(mode="json"),
                    ),
                    "admin",
                )

            elif message_type == CommandType.CLOCK_PING:
                await hub.send(
                    websocket,
                    CommandEnvelope(
                        type=CommandType.CLOCK_PONG,
                        payload={
                            "node_id": raw.get("node_id") or node_id,
                            "client_send_ms": raw.get("client_send_ms"),
                            "client_send_perf_ms": raw.get("client_send_perf_ms"),
                            "server_time_ms": epoch_ms(),
                        },
                    ),
                )

            elif message_type == CommandType.ACK:
                node_id = str(raw.get("node_id") or node_id or "")
                sync.clock_quality(node_id, raw.get("clock_offset_ms"), raw.get("rtt_ms"))
                record = sync.ack(raw)
                await hub.broadcast(
                    CommandEnvelope(
                        type=CommandType.ACK,
                        payload={**raw, "record": record.model_dump(mode="json") if record else None},
                    ),
                    "admin",
                )

            elif message_type == CommandType.TELEMETRY:
                if raw.get("node_id") and raw.get("tile_id"):
                    sync.telemetry(
                        PerformanceSample(
                            node_id=raw["node_id"],
                            tile_id=raw["tile_id"],
                            fps=float(raw.get("fps") or 0),
                            frame_time_ms=float(raw.get("frame_time_ms") or 0),
                            dropped_frames=int(raw.get("dropped_frames") or 0),
                        )
                    )

    except WebSocketDisconnect:
        sync.disconnect_node(node_id)
        await hub.disconnect("outputs", websocket)