"""WebSocket endpoint for real-time project event streaming.

Clients connect with a valid JWT access token (query param or first message).
All background worker events for a project are published to Redis and forwarded
here in real-time.

Connection:
    ws://<host>/ws/projects/<project_id>?token=<jwt_access_token>

Events are JSON objects. Agent stream lines use the canonical name ``run.progress`` (``stage`` may be
``generating``); legacy ``codegen.ai_output`` is not dual-published for those lines.

    {
        "event":      "run.progress",
        "project_id": "...",
        "timestamp":  "2026-03-31T12:00:00Z",
        "data":       {
            "run_id": "...",
            "text":   "line from the agent (paths redacted)",
            "type":   "output",
            "stream": "stdout" | "stderr",
            "stage":  "generating",
            "detail_format": "stream-json"   (optional; NDJSON event line for live progress / tool calls)
        }
    }

For scroll-back after disconnect, ``GET …/codegen/runs/{run}`` returns ``agent_transcript``
(capped) on the codegen run record.

See src/core/pubsub.py for the full event catalog.
"""
from __future__ import annotations

import asyncio
from uuid import UUID

from fastapi import APIRouter, Query, WebSocket, WebSocketDisconnect
from starlette.websockets import WebSocketState

import structlog

from src.core.pubsub import (
    get_pubsub_redis,
    subscribe_device_events,
    subscribe_project_events,
)

logger = structlog.get_logger()

router = APIRouter(tags=["websocket"])


async def _watch_client_disconnect(websocket: WebSocket) -> None:
    """Observe quiet browser closes even when no project events are outbound."""
    while True:
        message = await websocket.receive()
        if message.get("type") == "websocket.disconnect":
            return


@router.websocket("/ws/projects/{project_id}")
async def project_events_ws(
    websocket: WebSocket,
    project_id: UUID,
    token: str = Query(default=""),
):
    """Stream real-time project events over WebSocket.

    Authentication:
      - Pass a valid JWT access token as the ``?token=`` query parameter, OR
      - Send ``{"type": "auth", "token": "<jwt>"}`` as the first WebSocket message.

    The connection is closed with code 4001 if authentication fails.
    """
    from src.core.security.jwt import get_user_id_from_token

    # --- Auth via query param ---
    user_id = get_user_id_from_token(token) if token else None

    await websocket.accept()

    # --- Auth via first message (fallback for clients that cannot set query params) ---
    if not user_id:
        try:
            first = await asyncio.wait_for(websocket.receive_json(), timeout=10.0)
            if first.get("type") == "auth":
                user_id = get_user_id_from_token(first.get("token", ""))
        except Exception:
            pass

    if not user_id:
        await websocket.close(code=4001, reason="Unauthorized")
        return

    logger.info(
        "ws_connected",
        project_id=str(project_id),
        user_id=str(user_id),
    )

    redis = await get_pubsub_redis()
    try:
        async def forward_events() -> None:
            async for event in subscribe_project_events(redis, str(project_id)):
                await websocket.send_json(event)

        forward_task = asyncio.create_task(forward_events())
        disconnect_task = asyncio.create_task(_watch_client_disconnect(websocket))
        done, pending = await asyncio.wait(
            {forward_task, disconnect_task},
            return_when=asyncio.FIRST_COMPLETED,
        )
        for task in pending:
            task.cancel()
        if pending:
            await asyncio.gather(*pending, return_exceptions=True)
        for task in done:
            exc = task.exception()
            if exc and not isinstance(exc, (WebSocketDisconnect, RuntimeError, ConnectionError)):
                raise exc
    except asyncio.CancelledError:
        # Uvicorn/FastAPI cancels request tasks on shutdown — release Redis subscriber promptly.
        logger.info(
            "ws_cancelled",
            project_id=str(project_id),
            user_id=str(user_id),
        )
        try:
            if websocket.application_state != WebSocketState.DISCONNECTED:
                await websocket.close(code=1001, reason="Server shutting down")
        except Exception:
            pass
        raise
    except WebSocketDisconnect:
        pass
    except Exception as exc:
        logger.error(
            "ws_error",
            project_id=str(project_id),
            user_id=str(user_id),
            error=str(exc),
        )
    finally:
        await redis.aclose()
        logger.info(
            "ws_disconnected",
            project_id=str(project_id),
            user_id=str(user_id),
        )


@router.websocket("/ws/devices/{device_id}")
async def device_events_ws(
    websocket: WebSocket,
    device_id: UUID,
    token: str = Query(default=""),
):
    """Stream real-time run-assignment events to a paired desktop device.

    Authentication: prefer ``Authorization: Bearer <device_secret>`` during the WebSocket
    handshake. ``?token=`` remains as a backward-compatible fallback for existing builds. The
    connection is closed with code 4001 if the secret is missing, wrong, or the device has been
    revoked.

    Devices that aren't connected (or briefly drop) don't lose work: a dispatched run stays
    ``queued`` and is picked up via ``GET /api/devices/{device_id}/runs?status=queued`` on the
    next heartbeat/reconnect — see ``dispatch_codegen_local_run``.
    """
    from src.core.database import async_session_factory
    from src.core.security.hashing import verify_token_hash
    from src.domain.devices.repository import DeviceRepository

    auth_header = (websocket.headers.get("authorization") or "").strip()
    supplied = token
    if auth_header.lower().startswith("bearer "):
        supplied = auth_header[7:].strip()

    device_ok = False
    if supplied:
        async with async_session_factory() as db:
            device = await DeviceRepository(db).get_by_id(device_id)
            if device and device.status == "active" and device.device_secret_hash:
                device_ok = verify_token_hash(supplied, device.device_secret_hash)

    await websocket.accept()

    if not device_ok:
        await websocket.close(code=4001, reason="Unauthorized")
        return

    logger.info("device_ws_connected", device_id=str(device_id))

    redis = await get_pubsub_redis()
    try:
        async def forward_events() -> None:
            async for event in subscribe_device_events(redis, str(device_id)):
                await websocket.send_json(event)

        forward_task = asyncio.create_task(forward_events())
        disconnect_task = asyncio.create_task(_watch_client_disconnect(websocket))
        done, pending = await asyncio.wait(
            {forward_task, disconnect_task},
            return_when=asyncio.FIRST_COMPLETED,
        )
        for task in pending:
            task.cancel()
        if pending:
            await asyncio.gather(*pending, return_exceptions=True)
        for task in done:
            exc = task.exception()
            if exc and not isinstance(exc, (WebSocketDisconnect, RuntimeError, ConnectionError)):
                raise exc
    except asyncio.CancelledError:
        logger.info("device_ws_cancelled", device_id=str(device_id))
        try:
            if websocket.application_state != WebSocketState.DISCONNECTED:
                await websocket.close(code=1001, reason="Server shutting down")
        except Exception:
            pass
        raise
    except WebSocketDisconnect:
        pass
    except Exception as exc:
        logger.error("device_ws_error", device_id=str(device_id), error=str(exc))
    finally:
        await redis.aclose()
        logger.info("device_ws_disconnected", device_id=str(device_id))