File size: 23,796 Bytes

from __future__ import annotations

import json
import shutil
import socket
import threading
import uuid
from pathlib import Path

import gradio as gr
from fastapi import FastAPI, File, UploadFile
from fastapi.responses import HTMLResponse, JSONResponse

from app.config import CATEGORIES, DOC_KINDS, Settings, load_settings
from app.db import (
    add_catalog_item,
    delete_receipt,
    get_receipt,
    list_catalog,
    list_line_items,
    list_receipts,
    open_db,
    update_extract,
    update_receipt_status,
)
from app.pipeline import process_file
from app.schemas import ReceiptExtract, ReceiptStatus
from app.watcher import start_inbox_watcher
from backends import build_embed, build_llm

CSS = """
.gradio-container {max-width: 1200px !important;}
"""

PHONE_HTML = """<!doctype html>
<html lang="en"><head>
<meta charset="utf-8"/>
<meta name="viewport" content="width=device-width,initial-scale=1,viewport-fit=cover"/>
<meta name="apple-mobile-web-app-capable" content="yes"/>
<title>Scan a receipt</title>
<style>
 :root { color-scheme: dark; }
 * { -webkit-tap-highlight-color: transparent; }
 body{font-family:-apple-system,system-ui,sans-serif;background:#111814;color:#e8f0e8;margin:0;padding:24px}
 .card{max-width:480px;margin:0 auto;background:#1b2420;border-radius:16px;padding:24px}
 h1{font-size:1.5rem;margin:0 0 8px}
 p{line-height:1.4;color:#c5d4c8}
 /* iOS ignores taps on display:none file inputs. Cover the whole button with an opacity-0 input. */
 .hit{position:relative;display:block;box-sizing:border-box;width:100%;margin:14px 0 0;min-height:56px;
      border-radius:12px;font:inherit;font-weight:600;text-align:center;line-height:56px;
      overflow:hidden;touch-action:manipulation}
 .cam{background:#c4a35a;color:#1b1408}
 .lib{background:#2b3a32;color:#e8f0e8;border:1px solid #4d6556}
 .hit input[type=file]{
   position:absolute;left:0;top:0;width:100%;height:100%;margin:0;padding:0;
   opacity:0;font-size:24px;cursor:pointer;z-index:2}
 .ok{color:#9fdfb2;margin-top:12px;word-break:break-word}
 .err{color:#ff8a80;margin-top:12px;word-break:break-word}
</style></head>
<body><div class="card">
<h1>Scan a receipt</h1>
<p>Tap <b>Take photo</b> for the camera, or <b>Choose file</b> for Photos / Files. Use Safari.</p>
<div class="hit cam"><span>Take photo</span>
  <input id="cam" type="file" accept="image/*" capture="environment" tabindex="0"/>
</div>
<div class="hit lib"><span>Choose file</span>
  <input id="lib" type="file" accept="image/*,image/heic,image/heif,application/pdf,.heic,.heif,.pdf,.jpg,.jpeg,.png" tabindex="0"/>
</div>
<div id="out"></div>
<script>
const out = document.getElementById('out');
function add(cls, msg) {
  const d = document.createElement('div');
  d.className = cls;
  d.textContent = msg;
  out.prepend(d);
}
async function waitDone(jobId) {
  add('ok', 'Uploaded. Reading the receipt on the GPU box…');
  const deadline = Date.now() + 180000;
  while (Date.now() < deadline) {
    await new Promise(r => setTimeout(r, 2000));
    const r = await fetch('/api/jobs/' + jobId);
    const j = await r.json();
    if (j.state === 'processing') continue;
    if (j.state === 'done') {
      const bits = [];
      if (j.vendor) bits.push(j.vendor);
      if (j.total) bits.push('$' + j.total);
      if (j.category) bits.push(j.category);
      if (j.receipt_id != null) bits.push('#' + j.receipt_id);
      const line = bits.join(' · ') || j.status || 'done';
      if (j.status === 'failed' || j.status === 'needs_extract') {
        add('err', 'Finished with issues: ' + line + (j.error ? ' — ' + j.error : ''));
      } else {
        add('ok', 'Completed: ' + line + '. Confirm it on the desktop Review tab.');
      }
      return;
    }
    add('err', 'Failed: ' + (j.error || j.state));
    return;
  }
  add('err', 'Still processing after 3 minutes — check Review on the desktop.');
}
async function upload(file) {
  if (!file) { add('err', 'No file selected'); return; }
  add('ok', 'Uploading ' + (file.name || 'photo') + '…');
  const body = new FormData();
  const name = file.name && file.name !== 'image.jpg' ? file.name : ('iphone-' + Date.now() + '.jpg');
  body.append('file', file, name);
  try {
    const r = await fetch('/api/inbox', { method: 'POST', body });
    const text = await r.text();
    let j;
    try { j = JSON.parse(text); } catch (e) { throw new Error('Server returned ' + r.status + ': ' + text.slice(0, 180)); }
    if (!r.ok || !j.ok) throw new Error(text.slice(0, 180));
    await waitDone(j.job_id);
  } catch (err) {
    add('err', String(err));
  }
}
function bind(id) {
  const el = document.getElementById(id);
  el.addEventListener('change', () => {
    const files = el.files;
    if (!files || !files.length) { add('err', 'Picker closed with no file'); return; }
    for (const f of files) upload(f);
    el.value = '';
  });
}
bind('cam');
bind('lib');
</script></div></body></html>
"""


def _lan_ip() -> str:
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    try:
        sock.connect(("192.0.2.1", 1))
        return sock.getsockname()[0]
    except OSError:
        return "127.0.0.1"
    finally:
        sock.close()


def _as_text(raw: object) -> str:
    if raw is None:
        return ""
    if isinstance(raw, list):
        return "\n".join(str(part) for part in raw)
    return str(raw)


def pretty_json(raw: object) -> str:
    text = _as_text(raw).strip() or "{}"
    return json.dumps(json.loads(text), indent=2)


def parse_rid(value: object) -> int | None:
    if value is None or value == "":
        return None
    try:
        return int(float(str(value)))
    except (TypeError, ValueError):
        return None


def _table_rows(data: object) -> list[list[object]]:
    if data is None:
        return []
    if hasattr(data, "values"):
        return [list(row) for row in data.values]
    return [list(row) for row in data]


def _image_preview(path: str) -> str | None:
    if not path:
        return None
    suffix = Path(path).suffix.lower()
    if suffix in {".jpg", ".jpeg", ".png", ".webp", ".tif", ".tiff"}:
        return path
    return None


def _save_upload(settings: Settings, name: str, data: bytes) -> Path:
    safe = Path(name).name or "upload.bin"
    dest = settings.inbox_dir / safe
    n = 1
    while dest.exists():
        dest = settings.inbox_dir / f"{dest.stem}-{n}{dest.suffix}"
        n += 1
    dest.write_bytes(data)
    return dest


_JOBS: dict[str, dict[str, object]] = {}
_JOBS_LOCK = threading.Lock()


def build_app(settings: Settings) -> FastAPI:
    settings.ensure_dirs()
    api = FastAPI(title="keys-automatic-receipt-doc-scanner")

    @api.post("/api/inbox")
    async def api_inbox(file: UploadFile = File(...)) -> JSONResponse:
        data = await file.read()
        dest = _save_upload(settings, file.filename or "iphone.jpg", data)
        job_id = uuid.uuid4().hex
        with _JOBS_LOCK:
            _JOBS[job_id] = {"state": "processing"}

        def _run() -> None:
            try:
                result = process_file(dest, settings)
                extract = result.extract
                payload: dict[str, object] = {
                    "state": "done",
                    "status": result.status.value,
                    "receipt_id": result.receipt_id,
                    "vendor": extract.vendor if extract else None,
                    "total": str(extract.total) if extract and extract.total is not None else None,
                    "category": extract.category if extract else None,
                    "error": result.error,
                }
            except Exception as exc:
                payload = {"state": "error", "error": str(exc)}
            with _JOBS_LOCK:
                _JOBS[job_id] = payload

        threading.Thread(target=_run, daemon=True, name="inbox-scan").start()
        return JSONResponse(
            {
                "ok": True,
                "job_id": job_id,
                "path": str(dest),
                "bytes": len(data),
                "processing": True,
            }
        )

    @api.get("/api/health")
    async def api_health() -> JSONResponse:
        from app.route import resolve_brain

        try:
            brain = resolve_brain(settings)
            payload = {
                "ok": True,
                "route": settings.llm_route,
                "kind": brain.kind,
                "base_url": brain.base_url,
                "model": brain.model,
                "source": brain.source,
            }
        except Exception as exc:
            payload = {"ok": False, "error": str(exc), "route": settings.llm_route}
        return JSONResponse(payload)

    @api.get("/api/jobs/{job_id}")
    async def job_status(job_id: str) -> JSONResponse:
        with _JOBS_LOCK:
            job = _JOBS.get(job_id)
        if job is None:
            return JSONResponse({"state": "error", "error": "unknown job"}, status_code=404)
        return JSONResponse(job)

    @api.get("/phone", response_class=HTMLResponse)
    async def phone() -> str:
        return PHONE_HTML

    def health() -> str:
        llm = build_llm(settings)
        embed = build_embed(settings)
        return (
            f"LLM {settings.llm_backend}/{settings.llm_model} @ {settings.llm_base_url} "
            f"vision={llm.accepts_images} health={llm.health()}\n"
            f"Embed {settings.embed_backend}/{settings.embed_model} dim={settings.embed_dim} "
            f"health={embed.health()}\n"
            f"Camera {settings.camera_url}  idle={settings.idle_seconds}s"
        )

    def inbox_list() -> str:
        files = sorted(p.name for p in settings.inbox_dir.iterdir() if p.is_file())
        return "\n".join(files) or "(empty inbox)"

    def ingest(files: list) -> str:
        if not files:
            return "no files"
        names = []
        for item in files:
            src = Path(item if isinstance(item, str) else item.name)
            dest = settings.inbox_dir / src.name
            shutil.copy2(src, dest)
            names.append(dest.name)
        return "queued: " + ", ".join(names)

    def review_table() -> list[list[str]]:
        con = open_db(settings)
        try:
            rows = list_receipts(con, limit=40)
            return [
                [
                    str(r["id"]),
                    r["status"],
                    r["doc_kind"] or "",
                    r["category"] or "",
                    r["vendor"] or "",
                    r["receipt_date"] or "",
                    "" if r["total_cents"] is None else f"{r['total_cents']/100:.2f}",
                ]
                for r in rows
            ]
        finally:
            con.close()

    def _empty_load() -> tuple:
        return (
            None,
            "",
            "receipt",
            "other",
            "",
            "",
            "",
            "",
            "",
            "{}",
            [],
            "not found",
        )

    def load_one(receipt_id: object) -> tuple:
        rid = parse_rid(receipt_id)
        if rid is None:
            return _empty_load()
        con = open_db(settings)
        try:
            row = get_receipt(con, rid)
            if row is None:
                return _empty_load()
            try:
                payload = pretty_json(row["extract_json"] or "{}")
            except json.JSONDecodeError:
                payload = row["extract_json"] or "{}"
            extract = None
            try:
                extract = ReceiptExtract.model_validate_json(payload)
            except Exception:
                extract = None
            lines = [
                [
                    item.description,
                    "" if item.qty is None else item.qty,
                    "" if item.unit_price is None else str(item.unit_price),
                    "" if item.amount is None else str(item.amount),
                    item.sku or "",
                ]
                for item in (extract.line_items if extract else [])
            ]
            if not lines:
                lines = [
                    [
                        r["description"],
                        r["qty"] if r["qty"] is not None else "",
                        "" if r["unit_price_cents"] is None else f"{r['unit_price_cents']/100:.2f}",
                        "" if r["amount_cents"] is None else f"{r['amount_cents']/100:.2f}",
                        r["sku"] or "",
                    ]
                    for r in list_line_items(con, rid)
                ]
            return (
                _image_preview(row["source_path"] or ""),
                str(rid),
                (extract.doc_kind if extract else row["doc_kind"]) or "receipt",
                (extract.category if extract else row["category"]) or "other",
                (extract.vendor if extract else row["vendor"]) or "",
                (
                    extract.date.isoformat()
                    if extract and extract.date
                    else (row["receipt_date"] or "")
                ),
                "" if extract is None or extract.tax is None else str(extract.tax),
                "" if extract is None or extract.total is None else str(extract.total),
                (extract.currency if extract else row["currency"]) or "",
                payload,
                lines,
                f"loaded #{rid}",
            )
        finally:
            con.close()

    def load_from_table(data: object, evt: gr.SelectData) -> tuple:
        rows = _table_rows(data)
        index = evt.index
        row_i = index[0] if isinstance(index, (list, tuple)) else index
        if row_i is None or row_i < 0 or row_i >= len(rows) or not rows[row_i]:
            return _empty_load()
        return load_one(rows[row_i][0])

    def save_extract(receipt_id: object, raw: object) -> tuple[str, list[list[str]]]:
        rid = parse_rid(receipt_id)
        if rid is None:
            return "pick a receipt (click a row or enter id)", review_table()
        try:
            extract = ReceiptExtract.model_validate_json(pretty_json(raw))
        except Exception as exc:
            return f"invalid JSON: {exc}", review_table()
        con = open_db(settings)
        try:
            update_extract(con, rid, extract, ReceiptStatus.needs_review)
        finally:
            con.close()
        return f"saved JSON for #{rid}", review_table()

    def save_fields(
        receipt_id: object,
        doc_kind: str,
        category: str,
        vendor: str,
        receipt_date: str,
        tax: str,
        total: str,
        currency: str,
        lines: object,
    ) -> tuple[str, str, list[list[str]]]:
        rid = parse_rid(receipt_id)
        if rid is None:
            return "pick a receipt first", "{}", review_table()
        def _num(val: object) -> str | None:
            text = str(val).strip()
            return None if text in {"", "None", "null"} else text

        items = []
        for row in _table_rows(lines):
            if not row or not str(row[0]).strip():
                continue
            items.append(
                {
                    "description": str(row[0]).strip(),
                    "qty": _num(row[1] if len(row) > 1 else None),
                    "unit_price": _num(row[2] if len(row) > 2 else None),
                    "amount": _num(row[3] if len(row) > 3 else None),
                    "sku": (str(row[4]).strip() or None) if len(row) > 4 else None,
                }
            )
        payload = {
            "doc_kind": doc_kind or "receipt",
            "category": category or "other",
            "vendor": vendor.strip() or None,
            "date": receipt_date.strip() or None,
            "tax": tax.strip() or None,
            "total": total.strip() or None,
            "currency": currency.strip() or None,
            "line_items": items,
        }
        raw = json.dumps(payload, indent=2)
        try:
            extract = ReceiptExtract.model_validate(payload)
        except Exception as exc:
            return f"invalid fields: {exc}", raw, review_table()
        con = open_db(settings)
        try:
            update_extract(con, rid, extract, ReceiptStatus.needs_review)
        finally:
            con.close()
        return f"saved receipt #{rid}", pretty_json(extract.model_dump_json()), review_table()

    def confirm(receipt_id: object) -> tuple[str, list[list[str]]]:
        rid = parse_rid(receipt_id)
        if rid is None:
            return "pick a receipt first", review_table()
        con = open_db(settings)
        try:
            update_receipt_status(con, rid, ReceiptStatus.confirmed)
        finally:
            con.close()
        return f"confirmed #{rid}", review_table()

    def delete_one(receipt_id: object) -> tuple:
        rid = parse_rid(receipt_id)
        empty = _empty_load()
        if rid is None:
            return (*empty[:-1], "pick a receipt first", review_table())
        con = open_db(settings)
        try:
            found = delete_receipt(con, rid)
        finally:
            con.close()
        if not found:
            return (*empty[:-1], f"not found #{rid}", review_table())
        return (*empty[:-1], f"deleted #{rid}", review_table())

    def catalog_table() -> list[list[str]]:
        con = open_db(settings)
        try:
            return [
                [str(r["id"]), r["sku"] or "", r["vendor"] or "", r["description"]]
                for r in list_catalog(con)
            ]
        finally:
            con.close()

    def add_sku(sku: str, vendor: str, description: str) -> str:
        if not description.strip():
            return "description required"
        con = open_db(settings)
        try:
            add_catalog_item(con, sku=sku or None, vendor=vendor or None, description=description)
        finally:
            con.close()
        return "added"

    def process_now(path: str) -> str:
        if not path:
            return "no path"
        result = process_file(Path(path), settings)
        return result.model_dump_json(indent=2)

    with gr.Blocks(title="Receipt Studio", theme=gr.themes.Soft(primary_hue="amber"), css=CSS) as demo:
        gr.Markdown(
            "# Receipt Studio\n"
            "Lamp camera · phone upload · Gemma 4 12B Unified on the GPU box (not on the Lamp)."
        )
        with gr.Tab("Inbox"):
            gr.Markdown(
                f"Drop files here or on your phone: `http://{_lan_ip()}:{settings.ui_port}/phone` "
                f"(bind `0.0.0.0` / `RECEIPT_UI_SHARE_LAN=true`). Syncthing can also land in `inbox/`."
            )
            files = gr.File(label="Photos / PDFs", file_count="multiple", type="filepath")
            ingest_btn = gr.Button("Queue in inbox")
            ingest_out = gr.Textbox(label="Queued")
            listing = gr.Textbox(label="Inbox", lines=8)
            refresh = gr.Button("Refresh inbox")
            ingest_btn.click(ingest, inputs=[files], outputs=[ingest_out]).then(
                inbox_list, outputs=[listing]
            )
            refresh.click(inbox_list, outputs=[listing])
            demo.load(inbox_list, outputs=[listing])
        with gr.Tab("Review"):
            gr.Markdown(
                "Click a row to load. Edit **Kind / Category / Vendor / Date / Tax / Total** "
                "(or the JSON), then **Save fields + lines**. "
                "**Delete** removes a bad scan from the database and disk."
            )
            table = gr.Dataframe(
                headers=["id", "status", "kind", "category", "vendor", "date", "total"],
                datatype=["str"] * 7,
                interactive=False,
                wrap=True,
            )
            refresh_r = gr.Button("Refresh queue")
            with gr.Row():
                img = gr.Image(label="Scan", type="filepath")
                with gr.Column():
                    rid = gr.Textbox(label="Receipt id")
                    kind = gr.Dropdown(choices=list(DOC_KINDS), label="Kind", value="receipt")
                    category = gr.Dropdown(choices=list(CATEGORIES), label="Category", value="other")
                    vendor = gr.Textbox(label="Vendor")
                    receipt_date = gr.Textbox(label="Date (YYYY-MM-DD)")
                    tax = gr.Textbox(label="Tax")
                    total = gr.Textbox(label="Total")
                    currency = gr.Textbox(label="Currency")
            raw = gr.Textbox(
                label="Extract JSON (editable)",
                lines=18,
                max_lines=40,
                interactive=True,
            )
            lines = gr.Dataframe(
                headers=["description", "qty", "unit_price", "amount", "sku"],
                datatype=["str", "str", "str", "str", "str"],
                label="Line items (editable)",
                interactive=True,
                wrap=True,
            )
            with gr.Row():
                load_btn = gr.Button("Load id")
                save_fields_btn = gr.Button("Save fields + lines")
                save_json_btn = gr.Button("Save JSON")
                ok_btn = gr.Button("Confirm")
                del_btn = gr.Button("Delete", variant="stop")
            msg = gr.Textbox(label="Status")
            load_outputs = [
                img,
                rid,
                kind,
                category,
                vendor,
                receipt_date,
                tax,
                total,
                currency,
                raw,
                lines,
                msg,
            ]
            refresh_r.click(review_table, outputs=[table])
            table.select(load_from_table, inputs=[table], outputs=load_outputs)
            load_btn.click(load_one, inputs=[rid], outputs=load_outputs)
            save_fields_btn.click(
                save_fields,
                inputs=[rid, kind, category, vendor, receipt_date, tax, total, currency, lines],
                outputs=[msg, raw, table],
            )
            save_json_btn.click(save_extract, inputs=[rid, raw], outputs=[msg, table])
            ok_btn.click(confirm, inputs=[rid], outputs=[msg, table])
            del_btn.click(delete_one, inputs=[rid], outputs=load_outputs + [table])
            demo.load(review_table, outputs=[table])
        with gr.Tab("Catalog"):
            cat = gr.Dataframe(headers=["id", "sku", "vendor", "description"])
            sku = gr.Textbox(label="SKU")
            vendor = gr.Textbox(label="Vendor")
            desc = gr.Textbox(label="Description")
            add_btn = gr.Button("Add SKU")
            add_msg = gr.Textbox()
            add_btn.click(add_sku, inputs=[sku, vendor, desc], outputs=[add_msg]).then(
                catalog_table, outputs=[cat]
            )
            demo.load(catalog_table, outputs=[cat])
        with gr.Tab("Settings"):
            gr.Markdown(
                "Gemma 4 12B Unified does **not** load on the Lamp (6 GB). "
                "This UI talks to the GPU box URLs in `.env`."
            )
            gr.Textbox(value=health, label="Backends", every=15)
            path = gr.Textbox(label="Process this path now")
            run = gr.Button("Process")
            run_out = gr.Textbox(lines=16)
            run.click(process_now, inputs=[path], outputs=[run_out])

    return gr.mount_gradio_app(api, demo, path="/")


def main() -> None:
    settings = load_settings()
    if settings.ui_share_lan:
        settings.ui_host = "0.0.0.0"
    start_inbox_watcher(settings)
    import uvicorn

    uvicorn.run(
        build_app(settings),
        host=settings.ui_host,
        port=settings.ui_port,
        log_level="info",
    )


if __name__ == "__main__":
    main()