"""WebSocket-прокси для VNC-консоли.

Поток данных:
  1. Backend получает от Proxmox одноразовый тикет через `vncproxy` (REST API).
  2. Открывает WebSocket к Proxmox с этим тикетом и портом. Аутентификация:
     - API-токен (заголовок `Authorization: PVEAPIToken=...`) — новые PVE ≥ 7.x
     - Cookie PVEAuthCookie + CSRFPreventionToken — если заданы PVE_USERNAME + PVE_PASSWORD
  3. Проксирует бинарный VNC-поток между браузером клиента и Proxmox.

Особенности Proxmox 8.x+:
  - WebSocket-эндпоинт vncwebsocket возвращает 302 Redirect с Location со
    схемой https:// вместо wss://. Обрабатываем редирект вручную.
  - На WebSocket Proxmox может требовать одновременно Cookie и CSRF-токен.
  - Proxmox в Location редиректа часто возвращает свой self-reported
    origin (например, pve1.input.netcraze.pro — тот, что в сертификате).
    Если backend работает внутри сети и должен стучаться на внутренний
    IP (192.168.31.4), нужно заменить hostname в Location обратно.
"""

import asyncio
import logging
from typing import Optional, Tuple
from urllib.parse import quote, urlparse, urlunparse

import aiohttp
import httpx
from aiohttp import ClientSession, WSMsgType
from fastapi import APIRouter, Query, WebSocket, WebSocketDisconnect

from ..config import settings

logger = logging.getLogger(__name__)

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


def _rewrite_redirect(location: str, fallback_host: str) -> str:
    """Переписывает URL редиректа: заменяет https→wss и нежелательные хосты.

    Args:
        location: значение заголовка Location из Proxmox.
        fallback_host: хост, на который нужно стучаться по сети
            (например, 192.168.31.4:8006). Если в location хост
            отличается — он будет заменён.
    """
    if not location:
        return location

    # 1. ws/wss вместо http/https.
    if location.startswith("https://"):
        location = "wss://" + location[len("https://"):]
    elif location.startswith("http://"):
        location = "ws://" + location[len("http://"):]

    # 2. Если хост в редиректе — не тот, через который мы работаем
    # (например, Proxmox отдаёт свой публичный DNS, а backend сидит
    # внутри сети), заменяем host:port на fallback_host.
    parsed = urlparse(location)
    if parsed.hostname:
        if ":" in fallback_host:
            fb_hostname, _, fb_port = fallback_host.partition(":")
        else:
            fb_hostname, fb_port = fallback_host, ""

        # Если хост в location не совпадает с нашим (например, Proxmox
        # отдал свой публичный DNS) — заменяем на наш внутренний.
        if parsed.hostname != fb_hostname:
            new_netloc = fb_hostname
            if fb_port:
                new_netloc = f"{fb_hostname}:{fb_port}"
            location = urlunparse(parsed._replace(netloc=new_netloc))
    return location


async def _get_pve_auth_cookie() -> Optional[Tuple[str, str]]:
    """Аутентифицируется в Proxmox по логину/паролю через REST API.

    Возвращает кортеж (cookie_header, csrf_token) или None,
    если PVE_USERNAME/PVE_PASSWORD не заданы.
    """
    if not settings.pve_username or not settings.pve_password:
        return None

    url = f"{settings.pve_host}/api2/json/access/ticket"
    payload = {"username": settings.pve_username, "password": settings.pve_password}
    verify = settings.pve_verify_ssl

    try:
        async with httpx.AsyncClient(verify=verify, timeout=10.0) as client:
            resp = await client.post(url, data=payload)
            resp.raise_for_status()
            data = resp.json().get("data", {})
            cookie = data.get("ticket")
            csrf = data.get("CSRFPreventionToken")
            if not cookie:
                logger.error("auth: Proxmox не вернул PVEAuthCookie")
                return None
            logger.info(
                "auth: PVEAuthCookie получен (len=%d, csrf_prefix=%s)",
                len(cookie), (csrf or "")[:8],
            )
            return f"PVEAuthCookie={cookie}", csrf or ""
    except Exception as exc:
        logger.exception("auth: ошибка аутентификации в Proxmox")
        return None


@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-поток между браузером и Proxmox."""
    await websocket.accept()

    guest_path = "qemu" if guest_type == "vm" else "lxc"
    # Хост, через который backend реально ходит к Proxmox внутри сети.
    pve_backend_host = settings.pve_host.replace("http://", "").replace("https://", "").split(":")[0]
    pve_backend_port = settings.pve_host.replace("http://", "").replace("https://", "").split(":")[-1]
    if not pve_backend_port.isdigit():
        pve_backend_port = "8006"
    fallback_host = f"{pve_backend_host}:{pve_backend_port}"

    safe_ticket = quote(ticket, safe="")
    safe_port = quote(str(port), safe="")
    upstream_url = (
        f"wss://{fallback_host}/api2/json/nodes/{node}/{guest_path}/{vmid}/vncwebsocket"
        f"?port={safe_port}&vncticket={safe_ticket}"
    )

    # Определяем способ аутентификации.
    # Cookie-вариант приоритетнее — Proxmox 8.x/9.x часто отказывает API-токену
    # на WebSocket-эндпоинте vncwebsocket.
    use_cookie = bool(settings.pve_username and settings.pve_password)
    headers: dict = {}
    csrf_token: Optional[str] = None
    auth_kind = "api-token"
    if use_cookie:
        auth = await _get_pve_auth_cookie()
        if auth:
            cookie_header, csrf_token = auth
            headers["Cookie"] = cookie_header
            if csrf_token:
                headers["CSRFPreventionToken"] = csrf_token
            auth_kind = "cookie"
        else:
            logger.warning("console_ws: cookie не получен, fallback на API-токен")
    if auth_kind == "api-token":
        headers["Authorization"] = (
            f"PVEAPIToken={settings.pve_token_name}={settings.pve_token_value}"
        )

    logger.info(
        "console_ws: connecting node=%s vmid=%s guest=%s port=%s auth=%s ticket_prefix=%s",
        node, vmid, guest_type, port, auth_kind, safe_ticket[:24],
    )

    upstream: Optional[aiohttp.ClientWebSocketResponse] = None
    session: Optional[ClientSession] = None
    try:
        session = ClientSession()
        resp = await session.get(
            upstream_url,
            headers=headers,
            allow_redirects=False,
            ssl=False if not settings.pve_verify_ssl else None,
        )
        logger.info(
            "console_ws: initial GET status=%s, location=%s, body=%s",
            resp.status,
            resp.headers.get("Location", "<none>"),
            (await resp.text())[:200],
        )
        if resp.status in (301, 302, 303, 307, 308):
            location = resp.headers.get("Location", "")
            upstream_url = _rewrite_redirect(location, fallback_host)
            logger.info("console_ws: rewritten URL to %s", upstream_url)
            await resp.release()
        elif resp.status == 101:
            await resp.release()
        else:
            await resp.release()
            raise aiohttp.ClientResponseError(
                request_info=resp.request_info,
                history=resp.history,
                status=resp.status,
                message=f"unexpected status {resp.status}",
                headers=resp.headers,
            )

        upstream = await session.ws_connect(
            upstream_url,
            headers=headers,
            ssl=False if not settings.pve_verify_ssl else None,
            autoclose=False,
            autoping=True,
            max_msg_size=8 * 1024 * 1024,
        )
        logger.info("console_ws: upstream connected")

        async def client_to_upstream():
            try:
                while True:
                    data = await websocket.receive_bytes()
                    await upstream.send_bytes(data)
            except WebSocketDisconnect:
                logger.info("console_ws: client disconnected")
            except Exception as exc:
                logger.warning("console_ws: client_to_upstream error: %s", exc)

        async def upstream_to_client():
            try:
                async for msg in upstream:
                    if msg.type == WSMsgType.BINARY:
                        await websocket.send_bytes(msg.data)
                    elif msg.type == WSMsgType.TEXT:
                        await websocket.send_bytes(msg.data.encode())
                    elif msg.type == WSMsgType.CLOSE:
                        logger.info("console_ws: upstream close: %s", msg.data)
                        break
                    elif msg.type == WSMsgType.ERROR:
                        logger.warning("console_ws: upstream error: %s", msg.data)
                        break
            except Exception as exc:
                logger.warning("console_ws: upstream_to_client error: %s", exc)

        await asyncio.gather(client_to_upstream(), upstream_to_client())
    except aiohttp.ClientResponseError as exc:
        logger.error(
            "console_ws: Proxmox ответил status=%s, auth=%s, message=%s",
            exc.status, auth_kind, exc.message,
        )
        await websocket.close(code=1011, reason=f"Proxmox {exc.status}")
    except Exception as exc:  # noqa: BLE001
        logger.exception("console_ws: ошибка подключения к Proxmox")
        await websocket.close(code=1011, reason=str(exc)[:120])
    finally:
        if upstream is not None:
            try:
                await upstream.close()
            except Exception:
                pass
        if session is not None:
            try:
                await session.close()
            except Exception:
                pass