app/pipeline.py

from __future__ import annotations

import hashlib
import shutil
import threading
from pathlib import Path

from app.config import Settings
from app.db import (
    get_by_sha,
    insert_receipt,
    open_db,
    set_line_match,
    update_receipt_status,
    upsert_vector,
)
from app.embed import embed_texts
from app.extract import extract_receipt
from app.match import match_receipt
from app.media import to_jpeg_bytes
from app.ocr import maybe_ocr
from app.schemas import ProcessResult, ReceiptStatus
from backends import build_embed, build_llm, build_ocr

_LOCK = threading.Lock()


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def _safe_move(src: Path, dest_dir: Path) -> Path:
    dest_dir.mkdir(parents=True, exist_ok=True)
    dest = dest_dir / src.name
    stem, suffix = dest.stem, dest.suffix
    n = 1
    while dest.exists():
        dest = dest_dir / f"{stem}-{n}{suffix}"
        n += 1
    shutil.move(str(src), str(dest))
    return dest


def process_file(path: Path, settings: Settings) -> ProcessResult:
    with _LOCK:
        return _process_file_locked(path, settings)


def _process_file_locked(path: Path, settings: Settings) -> ProcessResult:
    settings.ensure_dirs()
    src = Path(path)
    digest = sha256_file(src)
    con = open_db(settings)
    try:
        existing = get_by_sha(con, digest)
        if existing is not None:
            if src.parent.resolve() == settings.inbox_dir.resolve():
                _safe_move(src, settings.processed_dir)
            return ProcessResult(
                receipt_id=int(existing["id"]),
                status=ReceiptStatus.duplicate,
                source_path=str(src),
                error="duplicate sha256",
            )
        working = src
        if src.parent.resolve() == settings.inbox_dir.resolve():
            working = _safe_move(src, settings.processing_dir)

        ocr_backend = build_ocr(settings)
        llm = build_llm(settings)
        embed = build_embed(settings)
        ocr = maybe_ocr(ocr_backend, working)
        ocr_text = None if ocr is None else ocr.text
        try:
            image = to_jpeg_bytes(working, settings)
        except Exception as exc:
            failed = _safe_move(working, settings.failed_dir)
            rid = insert_receipt(
                con,
                source_path=str(failed),
                sha256=digest,
                status=ReceiptStatus.failed,
                ocr_text=ocr_text,
                error=str(exc),
            )
            return ProcessResult(
                receipt_id=rid,
                status=ReceiptStatus.failed,
                source_path=str(failed),
                error=str(exc),
            )

        if image is None and not ocr_text:
            rid = insert_receipt(
                con,
                source_path=str(working),
                sha256=digest,
                status=ReceiptStatus.needs_ocr,
            )
            return ProcessResult(
                receipt_id=rid,
                status=ReceiptStatus.needs_ocr,
                source_path=str(working),
                error="no image/text for extract",
            )

        try:
            extract = extract_receipt(
                llm, settings=settings, image_jpeg=image, ocr_text=ocr_text
            )
        except Exception as exc:
            final = _safe_move(working, settings.processed_dir)
            rid = insert_receipt(
                con,
                source_path=str(final),
                sha256=digest,
                status=ReceiptStatus.needs_extract,
                ocr_text=ocr_text,
                error=str(exc),
            )
            return ProcessResult(
                receipt_id=rid,
                status=ReceiptStatus.needs_extract,
                source_path=str(final),
                error=str(exc),
            )

        final = _safe_move(working, settings.processed_dir)
        rid = insert_receipt(
            con,
            source_path=str(final),
            sha256=digest,
            status=ReceiptStatus.needs_review,
            extract=extract,
            ocr_text=ocr_text,
        )
        try:
            if extract.vendor or extract.line_items:
                blob = " | ".join(
                    [
                        extract.doc_kind,
                        extract.category,
                        extract.vendor or "",
                        extract.date.isoformat() if extract.date else "",
                        *(item.description for item in extract.line_items[:12]),
                    ]
                )
                vec = embed_texts(embed, [blob], input_type="passage", settings=settings)[0]
                upsert_vector(con, "receipt_vec", "receipt_id", rid, vec)
            matches = match_receipt(con, settings, embed, extract)
            line_rows = con.execute(
                "SELECT id FROM line_items WHERE receipt_id = ? ORDER BY id", (rid,)
            ).fetchall()
            for row, hit in zip(line_rows, matches, strict=False):
                if hit.reason != "vendor knn":
                    set_line_match(con, int(row["id"]), hit)
        except Exception:
            matches = []
        return ProcessResult(
            receipt_id=rid,
            status=ReceiptStatus.needs_review,
            source_path=str(final),
            extract=extract,
            matches=matches,
        )
    finally:
        con.close()