ppt-tool/backend/api/v1/ppt/endpoints/jobs.py

"""Job status polling and SSE streaming endpoints."""
import asyncio
import json
import uuid

from fastapi import APIRouter, Depends, HTTPException
from fastapi.responses import StreamingResponse
from sqlalchemy.ext.asyncio import AsyncSession

from models.sse_response import SSEResponse, SSEErrorResponse
from models.sql.job import JobModel
from models.sql.user import UserModel
from services.database import get_async_session
from services.redis_service import get_arq_pool
from utils.auth_dependencies import get_current_user

JOBS_ROUTER = APIRouter(prefix="/api/v1/ppt", tags=["Jobs"])


@JOBS_ROUTER.get("/jobs/{job_id}")
async def get_job_status(
    job_id: uuid.UUID,
    _current_user: UserModel = Depends(get_current_user),
    session: AsyncSession = Depends(get_async_session),
):
    """Poll job status."""
    job = await session.get(JobModel, job_id)
    if not job:
        raise HTTPException(status_code=404, detail="Job not found")

    return {
        "id": str(job.id),
        "job_type": job.job_type,
        "status": job.status,
        "progress": job.progress,
        "progress_message": job.progress_message,
        "error_message": job.error_message,
        "presentation_id": str(job.presentation_id) if job.presentation_id else None,
        "created_at": job.created_at.isoformat() if job.created_at else None,
        "started_at": job.started_at.isoformat() if job.started_at else None,
        "completed_at": job.completed_at.isoformat() if job.completed_at else None,
    }


@JOBS_ROUTER.get("/jobs/{job_id}/stream")
async def stream_job_progress(
    job_id: uuid.UUID,
    _current_user: UserModel = Depends(get_current_user),
    session: AsyncSession = Depends(get_async_session),
):
    """SSE stream of job progress events via Redis pub/sub."""
    job = await session.get(JobModel, job_id)
    if not job:
        raise HTTPException(status_code=404, detail="Job not found")

    # If already completed/failed, return final event immediately
    if job.status in ("completed", "failed"):
        async def done_stream():
            data = {
                "type": "progress",
                "job_id": str(job.id),
                "progress": job.progress,
                "message": job.progress_message or job.status,
                "status": job.status,
            }
            yield SSEResponse(
                event="response", data=json.dumps(data)
            ).to_string()

        return StreamingResponse(done_stream(), media_type="text/event-stream")

    async def progress_stream():
        try:
            pool = await get_arq_pool()
            pubsub = pool.pubsub()
            channel = f"job:{job_id}:progress"
            await pubsub.subscribe(channel)

            # Send initial status
            yield SSEResponse(
                event="response",
                data=json.dumps({
                    "type": "progress",
                    "job_id": str(job_id),
                    "progress": job.progress,
                    "message": job.progress_message or "Waiting",
                    "status": job.status,
                }),
            ).to_string()

            # Listen for updates with timeout
            timeout = 600  # 10 minutes max
            elapsed = 0
            while elapsed < timeout:
                message = await pubsub.get_message(
                    ignore_subscribe_messages=True, timeout=2.0
                )
                if message and message["type"] == "message":
                    raw = message["data"]
                    if isinstance(raw, bytes):
                        raw = raw.decode()
                    payload = json.loads(raw)
                    yield SSEResponse(
                        event="response",
                        data=json.dumps({"type": "progress", **payload}),
                    ).to_string()

                    # Stop streaming on terminal status
                    if payload.get("status") in ("completed", "failed"):
                        break
                else:
                    elapsed += 2

                await asyncio.sleep(0)

            await pubsub.unsubscribe(channel)
            await pubsub.aclose()

        except Exception as e:
            yield SSEErrorResponse(detail=str(e)[:200]).to_string()

    return StreamingResponse(progress_stream(), media_type="text/event-stream")


@JOBS_ROUTER.delete("/jobs/{job_id}")
async def cancel_job(
    job_id: uuid.UUID,
    _current_user: UserModel = Depends(get_current_user),
    session: AsyncSession = Depends(get_async_session),
):
    """Cancel a queued or processing job."""
    job = await session.get(JobModel, job_id)
    if not job:
        raise HTTPException(status_code=404, detail="Job not found")

    if job.status in ("completed", "failed"):
        raise HTTPException(status_code=409, detail="Job already finished")

    # Try to abort via ARQ
    try:
        pool = await get_arq_pool()
        await pool.abort_job(str(job_id))
    except Exception:
        pass  # Best effort

    job.status = "failed"
    job.error_message = "Cancelled by user"
    job.progress_message = "Cancelled"
    await session.commit()

    return {"ok": True, "status": "cancelled"}