BadNote/server/badnote_server/ocr/queue.py

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