"""Runs router: list/read runs and steps, and stream a run live over WebSocket."""

from __future__ import annotations

import uuid

from fastapi import APIRouter, Depends, HTTPException, Query, WebSocket, status
from sqlalchemy.orm import Session, sessionmaker

from legate.agent.enums import TERMINAL_STATUSES, RunStatus
from legate.agent.events import run_event_bus
from legate.api.deps import Principal, require
from legate.api.schemas import RunOut, StepOut
from legate.db.session import get_db, get_session_factory
from legate.policy.rbac import Permission
from legate.repositories import runs as runs_repo
from legate.repositories import steps as steps_repo
from legate.repositories import users as users_repo
from legate.repositories import workspaces as workspaces_repo
from legate.security.tokens import TokenError, decode_token

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

_TERMINAL_VALUES = {s.value for s in TERMINAL_STATUSES}


@router.get("", response_model=list[RunOut])
def list_runs(
    limit: int = Query(default=100, ge=1, le=500),
    offset: int = Query(default=0, ge=0),
    principal: Principal = Depends(require(Permission.VIEW)),
    db: Session = Depends(get_db),
) -> list[RunOut]:
    runs = runs_repo.list_for_workspace(
        db, workspace_id=principal.workspace_id, limit=limit, offset=offset
    )
    return [RunOut.model_validate(r) for r in runs]


@router.get("/{run_id}", response_model=RunOut)
def get_run(
    run_id: uuid.UUID,
    principal: Principal = Depends(require(Permission.VIEW)),
    db: Session = Depends(get_db),
) -> RunOut:
    run = runs_repo.get(db, workspace_id=principal.workspace_id, run_id=run_id)
    if run is None:
        raise HTTPException(status.HTTP_404_NOT_FOUND, detail="run not found")
    return RunOut.model_validate(run)


@router.get("/{run_id}/steps", response_model=list[StepOut])
def list_steps(
    run_id: uuid.UUID,
    principal: Principal = Depends(require(Permission.VIEW)),
    db: Session = Depends(get_db),
) -> list[StepOut]:
    run = runs_repo.get(db, workspace_id=principal.workspace_id, run_id=run_id)
    if run is None:
        raise HTTPException(status.HTTP_404_NOT_FOUND, detail="run not found")
    steps = steps_repo.list_for_run(db, workspace_id=principal.workspace_id, run_id=run_id)
    return [StepOut.model_validate(s) for s in steps]


def _authenticate_ws(token: str | None, db: Session) -> Principal | None:
    """Resolve a JWT access token to a workspace-bound principal for WebSocket use."""
    if not token:
        return None
    try:
        payload = decode_token(token, expected_type="access")
    except TokenError:
        return None
    if payload.workspace_id is None:
        return None
    user = users_repo.get_by_id(db, payload.user_id)
    if user is None or not user.is_active:
        return None
    membership = workspaces_repo.get_membership(
        db, user_id=user.id, workspace_id=payload.workspace_id
    )
    if membership is None:
        return None
    return Principal(
        auth_type="jwt", workspace_id=payload.workspace_id, user=user, role=membership.role
    )


@router.websocket("/{run_id}/stream")
async def stream_run(
    websocket: WebSocket,
    run_id: uuid.UUID,
    session_factory: sessionmaker = Depends(get_session_factory),
) -> None:
    """Stream a run's steps live.

    Auth uses a JWT access token via the ``token`` query parameter or an
    ``Authorization: Bearer`` header. The client first receives any already
    persisted steps (replay), then live events until the run is terminal.
    """
    await websocket.accept()
    token = websocket.query_params.get("token")
    if token is None:
        header = websocket.headers.get("authorization", "")
        if header.lower().startswith("bearer "):
            token = header[7:]

    db = session_factory()
    try:
        principal = _authenticate_ws(token, db)
        if principal is None or not principal.has_permission(Permission.VIEW):
            await websocket.close(code=4401)
            return

        run = runs_repo.get(db, workspace_id=principal.workspace_id, run_id=run_id)
        if run is None:
            await websocket.close(code=4404)
            return

        # Subscribe before replay so no live event is missed in the gap.
        queue = run_event_bus.subscribe(run_id)
        try:
            steps = steps_repo.list_for_run(db, workspace_id=principal.workspace_id, run_id=run_id)
            for step in steps:
                await websocket.send_json(
                    {"type": "step", "step": StepOut.model_validate(step).model_dump(mode="json")}
                )
            await websocket.send_json({"type": "run.status", "status": run.status.value})

            if run.status in TERMINAL_STATUSES or run.status == RunStatus.WAITING_APPROVAL:
                await websocket.close()
                return

            while True:
                event = await queue.get()
                await websocket.send_json(event)
                if event.get("type") == "run.status" and event.get("status") in _TERMINAL_VALUES:
                    break
        finally:
            run_event_bus.unsubscribe(run_id, queue)
        await websocket.close()
    finally:
        db.close()
