from __future__ import annotations

from fastapi import APIRouter, WebSocket, status
from fastapi.exceptions import WebSocketException
from starlette.websockets import WebSocketDisconnect

from app.ws.manager import ConnectionManager

router = APIRouter()
manager = ConnectionManager()

DEFAULT_HZ = 10
MIN_HZ = 1
MAX_HZ = 60


def _parse_hz(ws: WebSocket) -> int:
    raw = ws.query_params.get("hz")
    if raw is None:
        return DEFAULT_HZ
    try:
        hz = int(raw)
    except ValueError:
        raise WebSocketException(
            code=status.WS_1008_POLICY_VIOLATION, reason="Invalid 'hz' (int expected)"
        )
    if hz < MIN_HZ or hz > MAX_HZ:
        raise WebSocketException(
            code=status.WS_1008_POLICY_VIOLATION,
            reason=f"Invalid 'hz' (allowed {MIN_HZ}..{MAX_HZ})",
        )
    return hz


@router.websocket("/ws/live")
async def ws_live(ws: WebSocket) -> None:
    hz = _parse_hz(ws)

    await manager.connect(ws, hz=hz)
    try:
        # Не обязательно принимать сообщения от клиента
        # Но чтобы корректно ловить disconnect в некоторых клиентах - держим receive loop
        while True:
            await ws.receive_text()
    except WebSocketDisconnect:
        await manager.disconnect(ws)
    except WebSocketException:
        # если прилетит exception после accept — корректно удалим
        await manager.disconnect(ws)
        raise
    except Exception:
        await manager.disconnect(ws)