main.py

"""
Receipt OCR Service — FastAPI entry point.

Endpoints:
  POST /ocr          Upload a receipt image, get structured JSON back.
  GET  /health       Liveness check.

Usage:
  uvicorn main:app --host 0.0.0.0 --port 8000 --reload
"""

import asyncio
import hmac
import logging
import os
import time
from collections import defaultdict
from contextlib import asynccontextmanager

from fastapi import FastAPI, File, Header, HTTPException, Query, UploadFile
from fastapi.responses import JSONResponse, RedirectResponse

from ocr.reader import PaddleOCRReader
from ocr.parser import parse_blocks, is_complete_parse
from utils.debug import timing
from utils.image_prep import preprocess_image

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Global reader instance — loaded once at startup
_reader: PaddleOCRReader | None = None

# Shared secret for /ocr. When the env var is unset the endpoint stays open
# (local dev / CLI); when set, requests must carry it in X-OCR-Key. This is
# abuse deterrence for a public URL, not real auth — the key ships inside
# the mobile app.
_OCR_SHARED_SECRET = os.getenv("OCR_SHARED_SECRET", "")

# Reject uploads larger than this before decoding. Phone receipt photos are
# 1-6 MB; anything bigger is a mistake or abuse.
_MAX_UPLOAD_BYTES = 15 * 1024 * 1024


def _check_ocr_key(provided: str | None) -> None:
    if not _OCR_SHARED_SECRET:
        return
    if provided is None or not hmac.compare_digest(provided, _OCR_SHARED_SECRET):
        raise HTTPException(status_code=401, detail="Missing or invalid X-OCR-Key")


def _load_and_warm_reader():
    """Construct the reader (the slow model load) and JIT-warm it.

    Runs in a worker thread so the event loop — and therefore /health —
    responds from the moment the server binds. /ocr and /warmup return 503
    until _reader is set.
    """
    global _reader
    import numpy as np
    logger.info("Model load starting...")
    reader = PaddleOCRReader()
    _reader = reader
    logger.info("Model loaded; warm-up starting...")
    try:
        reader.extract(np.ones((200, 600, 3), dtype=np.uint8) * 255)
        logger.info("Warm-up complete.")
    except Exception as exc:
        logger.warning("Warm-up failed (non-fatal): %s", exc)


def _log_model_load_failure(future: asyncio.Future) -> None:
    if not future.cancelled() and future.exception() is not None:
        logger.error("Model load failed — /ocr will keep returning 503: %s",
                     future.exception())


@asynccontextmanager
async def lifespan(app: FastAPI):
    global _reader
    # Keep a reference on app.state so the future isn't garbage-collected
    # and startup failures surface in the log.
    future = asyncio.get_running_loop().run_in_executor(None, _load_and_warm_reader)
    future.add_done_callback(_log_model_load_failure)
    app.state.model_load_future = future
    yield
    _reader = None


app = FastAPI(
    title="Receipt OCR API",
    description="Extract structured data from receipt images using PaddleOCR.",
    version="1.0.0",
    lifespan=lifespan,
)


@app.get("/")
async def root():
    return RedirectResponse(url="/docs")


@app.get("/health")
async def health():
    return {"status": "ok", "model_loaded": _reader is not None}


@app.get("/warmup")
async def warmup():
    """Run OCR on a blank image to exercise the inference engine. Used by keepalive cron."""
    if _reader is None:
        raise HTTPException(status_code=503, detail="OCR model not loaded yet")
    import numpy as np
    _reader.extract(np.ones((200, 600, 3), dtype=np.uint8) * 255)
    return {"status": "warm"}


def _print_parse_result(result: dict, elapsed: float) -> None:
    """
    Print the structured parse output and flag anything that looks wrong.

    Uses print(flush=True) rather than the logging module — plain logger
    calls don't reliably surface in the HF Spaces log viewer. This line
    always prints (with total elapsed time); per-stage TIMING lines are
    opt-in via OCR_DEBUG_TIMING (utils/debug.py).
    """
    items = result.get("line_items", [])
    print(
        f"PARSED {elapsed:.3f}s | {result.get('merchant_name')} / {result.get('store_location')} | "
        f"{result.get('date')} | items={len(items)} subtotal={result.get('subtotal')} "
        f"savings={result.get('savings')} total={result.get('total')}",
        flush=True,
    )
    for n, it in enumerate(items, 1):
        discount = f" discount={it['discount']}" if it.get("discount") else ""
        desc = (it.get("description") or "<no description>")[:40]
        print(f"  {n:02d} qty{it.get('quantity')} {desc:<40} total={it.get('total_price')}{discount}", flush=True)

    # Anomaly flags come from parse_blocks() itself (result["anomalies"]) so the
    # API response and these prints always agree on what looks wrong.
    by_type: dict[str, list[int | None]] = defaultdict(list)
    for a in result.get("anomalies", []):
        by_type[a["type"]].append(a["item_index"])

    if by_type.get("no_description"):
        idxs = by_type["no_description"]
        print(f"ANOMALY: {len(idxs)} item(s) with no description at {idxs}", flush=True)
    if by_type.get("no_price"):
        idxs = by_type["no_price"]
        print(f"ANOMALY: {len(idxs)} item(s) with no price at {idxs}", flush=True)
    if by_type.get("no_items"):
        print("ANOMALY: no line items extracted", flush=True)
    if by_type.get("no_total"):
        print("ANOMALY: no total extracted", flush=True)


@app.post("/ocr")
async def ocr_receipt(
    file: UploadFile = File(..., description="Receipt image (JPEG, PNG, etc.)"),
    debug: bool = Query(False, description="Include raw OCR text blocks in the response"),
    x_ocr_key: str | None = Header(None, description="Shared secret when OCR_SHARED_SECRET is configured"),
):
    """
    Process a receipt image and return structured JSON.

    - **file**: multipart image upload
    - **debug**: if true, includes raw OCR text blocks in the response
    """
    _check_ocr_key(x_ocr_key)

    if _reader is None:
        raise HTTPException(status_code=503, detail="OCR model not loaded yet")

    # Validate content type loosely
    content_type = file.content_type or ""
    if content_type and not content_type.startswith("image/"):
        raise HTTPException(
            status_code=415,
            detail=f"Unsupported media type: {content_type}. Upload an image file.",
        )

    t0 = time.perf_counter()

    raw_bytes = await file.read(_MAX_UPLOAD_BYTES + 1)
    if not raw_bytes:
        raise HTTPException(status_code=400, detail="Empty file uploaded")
    if len(raw_bytes) > _MAX_UPLOAD_BYTES:
        raise HTTPException(
            status_code=413,
            detail=f"Upload too large — limit is {_MAX_UPLOAD_BYTES // (1024 * 1024)}MB",
        )

    try:
        image = preprocess_image(raw_bytes)
    except (ValueError, FileNotFoundError) as exc:
        raise HTTPException(status_code=400, detail=f"Image preprocessing failed: {exc}")
    t1 = time.perf_counter()
    timing(f"preprocess: {t1-t0:.2f}s | image size: {image.shape[1]}x{image.shape[0]} | upload: {len(raw_bytes)//1024}KB")

    try:
        # Skip the second (CLAHE) OCR pass when the first pass already
        # parses into a complete receipt whose items sum to the total —
        # a pass that missed rows can't satisfy that arithmetic.
        blocks = _reader.extract(
            image,
            second_pass_needed=lambda first: not is_complete_parse(parse_blocks(first)),
        )
    except Exception as exc:
        logger.exception("OCR extraction failed")
        raise HTTPException(status_code=500, detail=f"OCR failed: {exc}")
    t2 = time.perf_counter()
    timing(f"ocr extract: {t2-t1:.2f}s | blocks found: {len(blocks)}")

    result = parse_blocks(blocks)
    t3 = time.perf_counter()
    timing(f"parse: {t3-t2:.2f}s | items: {len(result.get('line_items', []))} | total: {t3-t0:.2f}s")
    _print_parse_result(result, t3 - t0)

    if debug:
        result["_raw_blocks"] = blocks

    return JSONResponse(content=result)