"""
WebSocket routes for real-time job status updates

Provides WebSocket endpoints for:
1. Individual job status updates: /ws/jobs/{job_id}
2. Job list updates: /ws/jobs (all jobs for authenticated user)
"""
import asyncio
import logging

from fastapi import (
    APIRouter,
    Depends,
    Query,
    WebSocket,
    WebSocketDisconnect,
)
from fastapi.security import HTTPBearer

from ...core.authz import PLATFORM_ADMIN_ROLES, _cached_memberships
from ...core.database import get_database
from ...models.user import UserRole
from ...services.websocket import (
    ConnectionManager,
    authenticate_websocket,
    connection_manager,
    get_connection_manager,
)

logger = logging.getLogger(__name__)

router = APIRouter(tags=["WebSocket"])
security = HTTPBearer()

# Close codes that indicate a permanent auth/permission failure — frontend must NOT retry
_TERMINAL_CLOSE_CODES = {4001, 4003, 4004, 4403}

# Seconds between server-side keepalive frames.
# Must be < Apache mod_proxy_wstunnel idle timeout.
# Mod Comms incident 2026-03-18: 25s was insufficient; 20s is safe.
_KEEPALIVE_INTERVAL_S = 20


async def _resolve_user_and_org(websocket: WebSocket, user_id: str, db):
    """
    Fetch user document and resolve org memberships from cache.
    Returns (user_doc, memberships_dict) or closes the socket and returns (None, None).
    """
    user = await db["users"].find_one({"_id": user_id})
    if not user:
        try:
            from bson import ObjectId
            user = await db["users"].find_one({"_id": ObjectId(user_id)})
        except Exception:
            pass

    if not user:
        await websocket.close(code=4001, reason="User not found")
        return None, None

    is_platform_admin = UserRole(user.get("role", "")) in PLATFORM_ADMIN_ROLES
    if is_platform_admin:
        return user, None  # None memberships = unrestricted

    memberships = await _cached_memberships(user_id, db)
    return user, memberships


def _can_access_org(org_id: str | None, memberships: dict | None) -> bool:
    """Return True if user (with these memberships) may access the given org_id."""
    if memberships is None:
        return True  # platform admin
    if not org_id:
        return True  # legacy job without org: allow (further checks done below if needed)
    return org_id in memberships


@router.websocket("/ws/jobs/{job_id}")
async def websocket_job_status(
    websocket: WebSocket,
    job_id: str,
    token: str | None = Query(None),
    manager: ConnectionManager = Depends(get_connection_manager)
):
    """
    WebSocket endpoint for real-time job status updates.

    Usage:
    - Connect: ws://localhost:8000/api/v1/ws/jobs/{job_id}?token={jwt_token}
    - Receives: Real-time status updates for the specific job

    Close codes:
      4001 — user not found
      4003 — role-based access denied
      4004 — job not found
      4403 — org membership access denied (do not retry)
    """
    user_id = await authenticate_websocket(websocket, token)
    if not user_id:
        return

    try:
        db = await get_database()

        job = await db["jobs"].find_one({"_id": job_id})
        if not job:
            await websocket.close(code=4004, reason="Job not found")
            return

        user, memberships = await _resolve_user_and_org(websocket, user_id, db)
        if user is None:
            return  # socket already closed inside helper

        # Role-based client restriction
        if user["role"] == "client" and job.get("created_by") != user_id:
            await websocket.close(code=4003, reason="Access denied")
            return

        # Org membership check
        job_org = job.get("organization_id")
        if not _can_access_org(job_org, memberships):
            await websocket.close(code=4403, reason="Org access denied")
            return

        await manager.connect_job_status(websocket, user_id, job_id)

        while True:
            try:
                # Wait up to _KEEPALIVE_INTERVAL_S for a client message.
                # On timeout send a keepalive frame so the proxy idle timer resets.
                message = await asyncio.wait_for(
                    websocket.receive_text(),
                    timeout=_KEEPALIVE_INTERVAL_S,
                )
                logger.debug(f"Received WebSocket message from user {user_id}: {message}")
                if message == "ping":
                    await websocket.send_text("pong")

            except TimeoutError:
                await websocket.send_text("keepalive")

            except WebSocketDisconnect:
                break
            except Exception as e:
                logger.error(f"Error in WebSocket message handling: {e}")
                break

    except WebSocketDisconnect:
        pass
    except Exception as e:
        logger.error(f"WebSocket job status error: {e}")
    finally:
        await manager.disconnect(websocket, user_id)


@router.websocket("/ws/jobs")
async def websocket_job_list(
    websocket: WebSocket,
    token: str | None = Query(None),
    manager: ConnectionManager = Depends(get_connection_manager)
):
    """
    WebSocket endpoint for real-time job list updates.

    Usage:
    - Connect: ws://localhost:8000/api/v1/ws/jobs?token={jwt_token}
    - Receives: Real-time status updates for all jobs the user can access

    Only events for jobs in the user's accessible orgs are delivered.
    """
    user_id = await authenticate_websocket(websocket, token)
    if not user_id:
        return

    try:
        logger.info(f"WebSocket: Looking up user {user_id} in database")
        db = await get_database()

        user, memberships = await _resolve_user_and_org(websocket, user_id, db)
        if user is None:
            return  # socket already closed inside helper

        logger.info(f"WebSocket: User {user_id} found, role: {user.get('role', 'unknown')}")

        accessible_org_ids = None if memberships is None else list(memberships.keys())
        await manager.connect_job_list(websocket, user_id, accessible_org_ids=accessible_org_ids)

        while True:
            try:
                message = await asyncio.wait_for(
                    websocket.receive_text(),
                    timeout=_KEEPALIVE_INTERVAL_S,
                )
                logger.debug(f"Received WebSocket message from user {user_id}: {message}")
                if message == "ping":
                    await websocket.send_text("pong")

            except TimeoutError:
                await websocket.send_text("keepalive")

            except WebSocketDisconnect:
                break
            except Exception as e:
                logger.error(f"Error in WebSocket message handling: {e}")
                break

    except WebSocketDisconnect:
        pass
    except Exception as e:
        logger.error(f"WebSocket job list error: {e}")
    finally:
        await manager.disconnect(websocket, user_id)


@router.get("/ws/status")
async def websocket_status():
    """Get WebSocket connection status and statistics (debug/monitoring)."""
    stats = {
        "active_connections": len(connection_manager.active_connections),
        "job_subscriptions": len(connection_manager.job_subscriptions),
        "global_subscriptions": len(connection_manager.global_subscriptions),
        "redis_connected": connection_manager.redis_client is not None,
        "subscriber_running": (
            connection_manager.subscriber_task is not None and
            not connection_manager.subscriber_task.done()
        )
    }
    return stats
