"""File-based OCR job queue for BadNote.""" import json import os import shutil from datetime import datetime, timezone from uuid import uuid4 from ..config import settings def _queue_dir(subdir: str) -> str: path = os.path.join(settings.queue_path, subdir) os.makedirs(path, exist_ok=True) return path def _job_path(job_id: str, subdir: str) -> str: return os.path.join(_queue_dir(subdir), f"{job_id}.json") def _read_job(path: str) -> dict | None: try: with open(path, "r") as f: return json.load(f) except (FileNotFoundError, json.JSONDecodeError): return None def _write_job(path: str, data: dict) -> None: with open(path, "w") as f: json.dump(data, f, indent=2) def enqueue(job_data: dict) -> str: """Add a job to the pending queue. Returns job_id.""" job_id = job_data.get("id", str(uuid4())) job_data["id"] = job_id job_data["status"] = "pending" job_data["created_at"] = datetime.now(timezone.utc).isoformat() _write_job(_job_path(job_id, "pending"), job_data) return job_id def dequeue() -> dict | None: """Move the first pending job to processing. Returns job dict or None.""" pending_dir = _queue_dir("pending") try: files = sorted(os.listdir(pending_dir)) except FileNotFoundError: return None for fname in files: if not fname.endswith(".json"): continue src = os.path.join(pending_dir, fname) job = _read_job(src) if job is None: continue job["status"] = "processing" dst = _job_path(job["id"], "processing") shutil.move(src, dst) return job return None def complete(job_id: str, result: str) -> None: """Mark a job as done with result text.""" src = _job_path(job_id, "processing") job = _read_job(src) if job is None: return job["status"] = "done" job["result_text"] = result job["completed_at"] = datetime.now(timezone.utc).isoformat() dst = _job_path(job_id, "done") if os.path.exists(src): os.remove(src) _write_job(dst, job) def fail(job_id: str, error: str) -> None: """Mark a job as failed with error message.""" src = _job_path(job_id, "processing") job = _read_job(src) if job is None: return job["status"] = "failed" job["error_message"] = error job["completed_at"] = datetime.now(timezone.utc).isoformat() dst = _job_path(job_id, "failed") if os.path.exists(src): os.remove(src) _write_job(dst, job) def get_status(job_id: str) -> dict | None: """Check all directories for a job and return its data.""" for subdir in ("pending", "processing", "done", "failed"): job = _read_job(_job_path(job_id, subdir)) if job is not None: return job return None def get_jobs_for_note(note_id: str) -> list[dict]: """Return all completed OCR jobs for a given note_id.""" results = [] for subdir in ("done", "pending", "processing", "failed"): dir_path = _queue_dir(subdir) try: for fname in os.listdir(dir_path): if not fname.endswith(".json"): continue job = _read_job(os.path.join(dir_path, fname)) if job and job.get("note_id") == note_id: results.append(job) except FileNotFoundError: pass return results def get_jobs_for_document(document_id: str) -> list[dict]: """Return all OCR jobs for a given document_id.""" results = [] for subdir in ("done", "pending", "processing", "failed"): dir_path = _queue_dir(subdir) try: for fname in os.listdir(dir_path): if not fname.endswith(".json"): continue job = _read_job(os.path.join(dir_path, fname)) if job and job.get("document_id") == document_id: results.append(job) except FileNotFoundError: pass return results