import asyncio

import websockets
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Query

from ..config import settings

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


@router.websocket("/ws")
async def console_ws(
    websocket: WebSocket,
    node: str = Query(...),
    vmid: int = Query(...),
    guest_type: str = Query(...),
    port: int = Query(...),
    ticket: str = Query(...),
):
    """
    Проксирует бинарный VNC-поток между браузером клиента и websocket-эндпоинтом
    Proxmox (nodes/{node}/(qemu|lxc)/{vmid}/vncwebsocket). Тикет и порт берутся
    из ответа /instances/{id}/console (см. instances.py -> pve.get_vnc_ticket).

    Примечание: в зависимости от версии Proxmox и способа аутентификации
    (API-токен vs cookie-тикет) может понадобиться донастройка заголовков —
    см. README, раздел "Консоль VNC".
    """
    await websocket.accept()

    guest_path = "qemu" if guest_type == "vm" else "lxc"
    pve_host = settings.pve_host.replace("http://", "").replace("https://", "")
    upstream_url = (
        f"wss://{pve_host}/api2/json/nodes/{node}/{guest_path}/{vmid}/vncwebsocket"
        f"?port={port}&vncticket={ticket}"
    )

    headers = {
        "Authorization": f"PVEAPIToken={settings.pve_token_name}={settings.pve_token_value}"
    }

    try:
        async with websockets.connect(
            upstream_url,
            extra_headers=headers,
            subprotocols=["binary"],
            ssl=None if settings.pve_verify_ssl else False,
        ) as upstream:

            async def client_to_upstream():
                try:
                    while True:
                        data = await websocket.receive_bytes()
                        await upstream.send(data)
                except WebSocketDisconnect:
                    pass

            async def upstream_to_client():
                try:
                    async for message in upstream:
                        if isinstance(message, str):
                            message = message.encode()
                        await websocket.send_bytes(message)
                except websockets.ConnectionClosed:
                    pass

            await asyncio.gather(client_to_upstream(), upstream_to_client())
    except Exception as exc:  # noqa: BLE001
        await websocket.close(code=1011, reason=str(exc)[:120])