"""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 ]