BadNote/server/badnote_server/routers/ocr_router.py

"""OCR router for BadNote."""

from fastapi import APIRouter, Depends, HTTPException, status

from ..auth import get_current_user
from ..database import get_db
from ..models import OcrJobRequest, OcrJobStatus, OcrResult
from ..ocr import queue as job_queue

router = APIRouter()


@router.post("/process", status_code=status.HTTP_201_CREATED)
async def submit_ocr_job(
    body: OcrJobRequest,
    user_id: str = Depends(get_current_user),
) -> dict:
    """Enqueue an OCR job."""
    job_data: dict = {
        "user_id": user_id,
        "note_id": body.note_id,
        "document_id": body.document_id,
        "page_number": body.page_number,
    }
    job_id = job_queue.enqueue(job_data)
    return {"job_id": job_id}


@router.get("/status/{job_id}", response_model=OcrJobStatus)
async def get_job_status(
    job_id: str,
    user_id: str = Depends(get_current_user),
) -> OcrJobStatus:
    """Get OCR job status and result."""
    job = job_queue.get_status(job_id)
    # Treat jobs owned by another user as not found to avoid leaking their data.
    if job is None or job.get("user_id") != user_id:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Job not found")

    return OcrJobStatus(
        id=job["id"],
        status=job["status"],
        result_text=job.get("result_text"),
        error_message=job.get("error_message"),
        created_at=job["created_at"],
        completed_at=job.get("completed_at"),
    )


@router.get("/results/{note_id}", response_model=list[OcrResult])
async def get_ocr_results(
    note_id: str,
    user_id: str = Depends(get_current_user),
) -> list[OcrResult]:
    """Get all OCR results for a note."""
    jobs = [
        j for j in job_queue.get_jobs_for_note(note_id) if j.get("user_id") == user_id
    ]
    return [
        OcrResult(
            id=j["id"],
            note_id=j.get("note_id"),
            document_id=j.get("document_id"),
            page_number=j.get("page_number"),
            status=j["status"],
            result_text=j.get("result_text"),
            error_message=j.get("error_message"),
            created_at=j["created_at"],
            completed_at=j.get("completed_at"),
        )
        for j in jobs
    ]


@router.get("/results/document/{document_id}", response_model=list[OcrResult])
async def get_document_ocr_results(
    document_id: str,
    user_id: str = Depends(get_current_user),
) -> list[OcrResult]:
    """Get all OCR results for a document."""
    jobs = [
        j
        for j in job_queue.get_jobs_for_document(document_id)
        if j.get("user_id") == user_id
    ]
    return [
        OcrResult(
            id=j["id"],
            note_id=j.get("note_id"),
            document_id=j.get("document_id"),
            page_number=j.get("page_number"),
            status=j["status"],
            result_text=j.get("result_text"),
            error_message=j.get("error_message"),
            created_at=j["created_at"],
            completed_at=j.get("completed_at"),
        )
        for j in jobs
    ]