pipeline/jobs.py

"""Async job manager for PDF extraction.

One job per uploaded file. Heavy CPU work (rendering + OCR) is serialized
through a SINGLE-worker ThreadPoolExecutor via loop.run_in_executor so OCR
never runs concurrently. Per-job progress is pushed to an asyncio.Queue that
is consumed by the SSE endpoint.

Implements the project contract exactly:
  Job dataclass + to_dict()
  JobManager.create_job / get / run / cancel / events / render_page_png
  manager = JobManager()  module-level singleton
"""

import asyncio
import threading
import time
import uuid
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
from typing import Optional

import fitz  # type: ignore

from pipeline.extractor import extract_page


# Single-worker executor: serializes all heavy CPU work (OCR is not
# thread-safe and we want deterministic, low-memory CPU usage).
_EXECUTOR = ThreadPoolExecutor(max_workers=1, thread_name_prefix="ocr-worker")


# --- Job retention / eviction policy (in-memory only; single-user local) ------
# A finished job keeps its full pdf_bytes in RAM. To bound memory across a long
# session we evict the oldest / expired *finished, unsubscribed* jobs. Jobs that
# are still pending/processing, or that have a live SSE subscriber, are NEVER
# eviction candidates regardless of these limits.
MAX_JOBS = 50                       # keep at most this many evictable jobs
JOB_TTL_SECONDS = 2 * 60 * 60       # evict finished+idle jobs older than this
_SWEEP_MIN_INTERVAL = 5.0           # throttle: don't sweep on get() more than
                                    # once per this many seconds (monotonic)

_TERMINAL_STATES = ("done", "error", "cancelled")


@dataclass
class PageResult:
    page: int  # 1-based
    source: Optional[str] = None  # "text" | "ocr" | None
    text: str = ""
    status: str = "pending"  # "pending" | "done" | "error"
    error: Optional[str] = None
    confidence: Optional[float] = None  # mean OCR confidence [0,1]; None for text-layer
    lines: Optional[list] = None  # per-line [{text, confidence}] for OCR pages (else None)

    def to_dict(self) -> dict:
        return {
            "page": self.page,
            "source": self.source,
            "text": self.text,
            "status": self.status,
            "error": self.error,
            "confidence": self.confidence,
            "lines": self.lines,
        }


@dataclass
class Job:
    job_id: str
    filename: str
    pdf_bytes: bytes
    mode: str
    lang: str
    preprocess: bool
    binarize: bool
    handwriting: bool
    online: bool
    online_key: str  # NOT serialized in to_dict() — never leaked to the API
    online_model: str
    force_ocr: bool
    remove_headers: bool
    total_pages: int
    pages: list = field(default_factory=list)  # list[PageResult]
    status: str = "pending"  # pending|processing|done|error|cancelled
    processed_pages: int = 0
    error: Optional[str] = None

    # Optional knowledge graph (opt-in, online). Built on demand from this job's
    # extracted pages and attached here, so it is reclaimed with the job by the
    # TTL/LRU eviction below and — because to_dict() is an explicit allow-list —
    # is NEVER serialized to the client or broadcast over SSE (same protection
    # the API key gets). Held loosely as ``object`` to keep jobs.py free of the
    # kg / numpy import at module load.
    knowledge_graph: object = field(default=None, repr=False)
    kg_status: str = field(default="none", repr=False)  # none|building|ready|error
    kg_error: Optional[str] = field(default=None, repr=False)

    # Internals (not serialized). One asyncio.Queue per live SSE subscriber so
    # progress is broadcast (fan-out), not consumed once — reconnects are safe.
    subscribers: list = field(default_factory=list, repr=False)
    cancel_event: "threading.Event" = field(default_factory=threading.Event, repr=False)
    # Eviction bookkeeping (monotonic clock; immune to wall-clock changes).
    created_at: float = field(default_factory=time.monotonic, repr=False)
    last_access: float = field(default_factory=time.monotonic, repr=False)

    def to_dict(self) -> dict:
        return {
            "job_id": self.job_id,
            "filename": self.filename,
            "status": self.status,
            "total_pages": self.total_pages,
            "processed_pages": self.processed_pages,
            "pages": [p.to_dict() for p in self.pages],
            "error": self.error,
        }


def _try_repair(pdf_bytes: bytes):
    """Attempt to repair a PDF fitz can't open, using pikepdf. Returns clean
    bytes on success, or None if repair is unavailable/failed."""
    try:
        import io as _io
        import pikepdf  # optional; only used as a fallback

        out = _io.BytesIO()
        with pikepdf.open(_io.BytesIO(pdf_bytes)) as pdf:
            pdf.save(out)
        return out.getvalue()
    except Exception:
        return None


class JobManager:
    def __init__(self) -> None:
        self._jobs: dict[str, Job] = {}
        self._lock = threading.Lock()
        # Throttle for get()-triggered sweeps (guarded by self._lock).
        self._last_sweep: float = 0.0

    # --- eviction internals --------------------------------------------------
    @staticmethod
    def _is_evictable(job: "Job") -> bool:
        """A job may be evicted only if it has reached a terminal state AND has
        no live SSE subscriber. pending/processing or subscribed jobs are never
        touched, so eviction can't break an in-flight run or an open stream.

        ``job.subscribers`` is mutated only on the event loop (events()); reading
        its length here under self._lock is a plain, atomic CPython list read.
        A momentarily-stale read is safe: worst case we *skip* evicting a job
        that just lost its last subscriber, deferring it to the next sweep.
        """
        return job.status in _TERMINAL_STATES and not job.subscribers

    def _evict_locked(self, now: float) -> None:
        """Evict expired, then over-cap, evictable jobs. MUST hold self._lock.

        1. TTL: drop evictable jobs whose last_access is older than the TTL.
        2. Cap: if more than MAX_JOBS evictable jobs remain, drop the
           least-recently-accessed ones until at most MAX_JOBS are left.
        Non-evictable (active/subscribed) jobs are never removed and do not
        count against the cap as targets — they're simply left in place.
        """
        self._last_sweep = now

        # Pass 1 — TTL on evictable jobs.
        if JOB_TTL_SECONDS > 0:
            expired = [
                jid
                for jid, job in self._jobs.items()
                if self._is_evictable(job)
                and (now - job.last_access) > JOB_TTL_SECONDS
            ]
            for jid in expired:
                del self._jobs[jid]

        # Pass 2 — count cap on remaining evictable jobs (LRU by last_access).
        if MAX_JOBS >= 0:
            evictable = [
                (job.last_access, jid)
                for jid, job in self._jobs.items()
                if self._is_evictable(job)
            ]
            overflow = len(evictable) - MAX_JOBS
            if overflow > 0:
                evictable.sort()  # oldest last_access first
                for _, jid in evictable[:overflow]:
                    job = self._jobs.get(jid)
                    if job is not None and self._is_evictable(job):
                        del self._jobs[jid]

    def create_job(
        self,
        filename: str,
        pdf_bytes: bytes,
        *,
        mode: str,
        lang: str,
        preprocess: bool,
        binarize: bool = False,
        handwriting: bool = False,
        online: bool = False,
        online_key: str = "",
        online_model: str = "",
        force_ocr: bool = False,
        remove_headers: bool = True,
    ) -> Job:
        """Open the PDF to validate it and count pages, then register a Job.

        Raises ValueError("encrypted") for password-protected PDFs and
        ValueError(<msg>) for corrupt / unreadable PDFs.
        """
        usable_bytes = pdf_bytes
        try:
            doc = fitz.open(stream=usable_bytes, filetype="pdf")
        except Exception as exc:  # corrupt / not a PDF — try to repair once
            repaired = _try_repair(pdf_bytes)
            if repaired is None:
                raise ValueError(f"corrupt: {exc}") from exc
            usable_bytes = repaired
            try:
                doc = fitz.open(stream=usable_bytes, filetype="pdf")
            except Exception as exc2:
                raise ValueError(f"corrupt: {exc2}") from exc2

        try:
            if doc.needs_pass:
                raise ValueError("encrypted")
            try:
                total_pages = doc.page_count
            except Exception as exc:
                raise ValueError(f"corrupt: {exc}") from exc
            if total_pages <= 0:
                raise ValueError("corrupt: no pages")
        finally:
            doc.close()

        job_id = uuid.uuid4().hex
        job = Job(
            job_id=job_id,
            filename=filename,
            pdf_bytes=usable_bytes,
            mode=mode,
            lang=lang,
            preprocess=preprocess,
            binarize=binarize,
            handwriting=handwriting,
            online=online,
            online_key=online_key,
            online_model=online_model,
            force_ocr=force_ocr,
            remove_headers=remove_headers,
            total_pages=total_pages,
            pages=[PageResult(page=i + 1) for i in range(total_pages)],
            status="pending",
        )

        with self._lock:
            self._evict_locked(time.monotonic())  # reclaim stale jobs first
            self._jobs[job_id] = job
        return job

    def get(self, job_id: str) -> Optional[Job]:
        now = time.monotonic()
        with self._lock:
            job = self._jobs.get(job_id)
            if job is not None:
                # Touch: keeps actively-used jobs (polling, image render, SSE
                # reconnect) from aging out under the TTL.
                job.last_access = now
            # Throttled TTL sweep so idle sessions still reclaim memory even
            # with no new uploads, without sweeping on every single get().
            if now - self._last_sweep >= _SWEEP_MIN_INTERVAL:
                self._evict_locked(now)
            return job

    @staticmethod
    def _broadcast(job: Job, event: dict) -> None:
        """Fan an event out to every live SSE subscriber's queue.

        Runs on the event loop; queues are unbounded so put_nowait never
        blocks. Each SSE connection drains its own queue independently.
        """
        for q in list(job.subscribers):
            try:
                q.put_nowait(event)
            except Exception:
                pass

    async def run(self, job: Job) -> None:
        """Process every page of the job sequentially.

        Each page's extraction runs in the single-worker executor. The cancel
        flag is checked between pages. Progress and terminal events are
        broadcast to all subscriber queues for SSE consumption. The whole body
        is guarded so the job always reaches a terminal state (no hung SSE).
        """
        loop = asyncio.get_running_loop()
        job.status = "processing"

        try:
            # Open one document for the lifetime of the run. Do it ON the single
            # OCR worker thread (not the event loop) so that every fitz call in
            # the app — open, load_page, get_pixmap, close — happens on that one
            # thread. PyMuPDF/MuPDF is not thread-safe even across separate
            # Document objects, and concurrent multi-file uploads would otherwise
            # let this run() open a doc on the event loop while another job
            # renders on the worker thread.
            try:
                doc = await loop.run_in_executor(
                    _EXECUTOR,
                    lambda: fitz.open(stream=job.pdf_bytes, filetype="pdf"),
                )
            except Exception as exc:
                job.status = "error"
                job.error = f"corrupt: {exc}"
                self._broadcast(job, {"type": "error", "error": job.error})
                return

            try:
                for idx in range(job.total_pages):
                    if job.cancel_event.is_set():
                        job.status = "cancelled"
                        self._broadcast(job, {"type": "cancelled"})
                        return

                    page_result = job.pages[idx]

                    def _work(page_index: int = idx) -> dict:
                        page = doc.load_page(page_index)
                        return extract_page(
                            page,
                            mode=job.mode,
                            lang=job.lang,
                            preprocess=job.preprocess,
                            binarize=job.binarize,
                            handwriting=job.handwriting,
                            online=job.online,
                            online_key=job.online_key,
                            online_model=job.online_model,
                            force_ocr=job.force_ocr,
                        )

                    try:
                        result = await loop.run_in_executor(_EXECUTOR, _work)
                        page_result.source = result.get("source")
                        page_result.text = result.get("text", "") or ""
                        page_result.confidence = result.get("confidence")
                        page_result.lines = result.get("lines")
                        page_result.status = "done"
                        page_result.error = None
                    except Exception as exc:  # per-page failure: keep going
                        page_result.source = None
                        page_result.text = ""
                        page_result.confidence = None
                        page_result.lines = None
                        page_result.status = "error"
                        page_result.error = str(exc)

                    job.processed_pages += 1

                    self._broadcast(
                        job,
                        {
                            "type": "progress",
                            "processed": job.processed_pages,
                            "total": job.total_pages,
                            "page": page_result.to_dict(),
                        },
                    )

                # Cross-page cleanup: strip running headers/footers + page
                # numbers now that every page's text is available. Best-effort;
                # never fail the job over it. Updated pages are re-broadcast so
                # the live UI (and reconnect snapshot/exports) show cleaned text.
                if job.remove_headers and not job.cancel_event.is_set():
                    try:
                        from pipeline.postprocess import (
                            strip_running_headers_footers,
                            realign_line_confidences,
                        )
                        originals = [p.text or "" for p in job.pages]
                        cleaned, kept_indices = strip_running_headers_footers(
                            originals, return_kept=True
                        )
                        for p, new_text, kept in zip(
                            job.pages, cleaned, kept_indices
                        ):
                            if p.status == "done" and new_text != (p.text or ""):
                                if p.lines:
                                    # An OCR page's per-line list is 1:1 with its
                                    # text lines, so map confidences to the EXACT
                                    # surviving lines by their original index
                                    # (correct even for identical-text duplicate
                                    # lines). Fall back to the text matcher only
                                    # if the indices don't line up with the
                                    # cleaned text (rare blank-detection edge).
                                    mapped = [
                                        p.lines[i]
                                        for i in kept
                                        if 0 <= i < len(p.lines)
                                    ]
                                    if [
                                        str(m.get("text", "")) for m in mapped
                                    ] == new_text.split("\n"):
                                        p.lines = mapped
                                    else:
                                        p.lines = realign_line_confidences(
                                            new_text, p.lines
                                        )
                                p.text = new_text
                                self._broadcast(
                                    job,
                                    {
                                        "type": "progress",
                                        "processed": job.processed_pages,
                                        "total": job.total_pages,
                                        "page": p.to_dict(),
                                    },
                                )
                    except Exception:
                        pass

                # Completed all pages (unless cancelled above).
                if job.cancel_event.is_set():
                    job.status = "cancelled"
                    self._broadcast(job, {"type": "cancelled"})
                else:
                    # If every page errored we still report done; the page
                    # records carry the per-page error state. A whole-job error
                    # is only for fatal failures handled above.
                    job.status = "done"
                    self._broadcast(job, {"type": "done"})
            finally:
                # Close on the worker thread too (same single-thread invariant),
                # and never let a close error overwrite the terminal status.
                try:
                    await loop.run_in_executor(_EXECUTOR, doc.close)
                except Exception:
                    pass
        except Exception as exc:  # never strand subscribers in a hung stream
            job.status = "error"
            job.error = str(exc)
            self._broadcast(job, {"type": "error", "error": job.error})
        finally:
            # The TTL is an IDLE timer, but last_access is otherwise only set at
            # creation/stream-start. A job that spent longer than JOB_TTL_SECONDS
            # actually PROCESSING (e.g. a 1000-page OCR, or a slow online-OCR run)
            # would otherwise be instantly evictable the moment it finishes — the
            # next sweep (triggered by any other upload) could drop a result the
            # user just received, 404-ing a page-image click or KG build. Start
            # the idle clock at completion instead. Covers every terminal path
            # (done / cancelled / error / corrupt-open early return).
            job.last_access = time.monotonic()

    def cancel(self, job_id: str) -> None:
        job = self.get(job_id)
        if job is not None:
            job.cancel_event.set()

    def delete(self, job_id: str) -> bool:
        """Drop a job from the registry and request cancellation.

        Removes the job from ``_jobs`` immediately (so it's no longer findable)
        and sets its cancel flag. The PDF bytes are reclaimed once any in-flight
        page finishes — cancellation is cooperative and checked between pages, so
        a long page already running (e.g. a 120 s online-OCR call) keeps the
        bytes alive until it returns. Idempotent — returns False if unknown.
        """
        with self._lock:
            job = self._jobs.pop(job_id, None)
        if job is not None:
            job.cancel_event.set()
            return True
        return False

    async def events(self, job_id: str):
        """Async generator yielding SSE event dicts for a job.

        On connect, replays a snapshot of every already-completed page so a
        fresh OR reconnecting client catches up without losing progress. Then
        streams live events from its own private subscriber queue until a
        terminal {"type":"done"|"error"|"cancelled"} event. Duplicate page
        events (snapshot + live) are harmless: the frontend keys pages by
        number and re-renders idempotently.
        """
        job = self.get(job_id)
        if job is None:
            yield {"type": "error", "error": "not found"}
            return

        terminal_types = {"done", "error", "cancelled"}

        # Subscribe BEFORE snapshotting so no event slips through the gap.
        q: "asyncio.Queue" = asyncio.Queue()
        job.subscribers.append(q)
        try:
            # Snapshot: replay completed pages for catch-up / reconnect.
            for p in job.pages:
                if p.status in ("done", "error"):
                    yield {
                        "type": "progress",
                        "processed": job.processed_pages,
                        "total": job.total_pages,
                        "page": p.to_dict(),
                    }
            if job.status in terminal_types:
                yield {"type": job.status, "error": job.error}
                return

            while True:
                try:
                    event = await asyncio.wait_for(q.get(), timeout=1.0)
                except asyncio.TimeoutError:
                    # If the job finished while we waited and our queue drained,
                    # synthesize the terminal event so the client closes.
                    if job.status in terminal_types and q.empty():
                        yield {"type": job.status, "error": job.error}
                        return
                    continue

                yield event
                if event.get("type") in terminal_types:
                    return
        finally:
            try:
                job.subscribers.remove(q)
            except ValueError:
                pass

    def render_page_png(self, job_id: str, page_no: int, dpi: int = 150) -> bytes:
        """Render a 1-based page number to PNG bytes on demand."""
        job = self.get(job_id)
        if job is None:
            raise KeyError("job not found")

        doc = fitz.open(stream=job.pdf_bytes, filetype="pdf")
        try:
            if page_no < 1 or page_no > doc.page_count:
                raise KeyError("page out of range")
            page = doc.load_page(page_no - 1)
            zoom = float(dpi) / 72.0
            matrix = fitz.Matrix(zoom, zoom)
            pix = page.get_pixmap(matrix=matrix, alpha=False)
            return pix.tobytes("png")
        finally:
            doc.close()


manager = JobManager()