app.py

"""BudgetBuddy — gradio.Server backend (v2).

Custom dark frontend (frontend/) over JSON API endpoints that wrap the on-device
small models: MiniCPM-V-4.6 vision on the Space's ZeroGPU, MiniCPM4.1-8B text on
Modal (core.inference routes each). Flexible transaction model, an understanding
agent, per-item analytics, budgets, and a tool-using chat agent.
"""

from __future__ import annotations

import base64
import io
import json
from datetime import date
from pathlib import Path

import gradio as gr
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles

from core.extract import extract_receipt, extract_payment, reconcile
from core.categorize import CATEGORIES, understand
from core import storage, analytics, auth, agent

FRONTEND = Path(__file__).parent / "frontend"
app = gr.Server()

_PERSIST_KEYS = ["vendor", "date", "currency", "line_items", "charges", "total",
                 "category", "understanding", "note", "source"]


# --------------------------------------------------------------------------- #
# Helpers
# --------------------------------------------------------------------------- #
def _open_image(image_b64: str):
    from PIL import Image
    if not image_b64:
        return None
    if "," in image_b64 and image_b64.strip().lower().startswith("data:"):
        image_b64 = image_b64.split(",", 1)[1]
    return Image.open(io.BytesIO(base64.b64decode(image_b64))).convert("RGB")


def _loads(s, default):
    if isinstance(s, (dict, list)):
        return s
    try:
        return json.loads(s) if s else default
    except Exception:
        return default


def _clean(record: dict) -> dict:
    out = {k: record.get(k) for k in _PERSIST_KEYS if k in record}
    out.setdefault("line_items", record.get("line_items") or [])
    out.setdefault("charges", record.get("charges") or [])
    return out


# --------------------------------------------------------------------------- #
# Auth
# --------------------------------------------------------------------------- #
@app.api(name="signup")
def signup(username: str, pin: str) -> dict:
    ok, msg = auth.signup(username, pin)
    return {"ok": ok, "message": msg,
            "token": auth.issue_token(username) if ok else "",
            "user": auth.normalize_username(username) if ok else ""}


@app.api(name="login")
def login(username: str, pin: str) -> dict:
    ok, msg = auth.login(username, pin)
    return {"ok": ok, "message": msg,
            "token": auth.issue_token(username) if ok else "",
            "user": auth.normalize_username(username) if ok else ""}


# --------------------------------------------------------------------------- #
# Capture
# --------------------------------------------------------------------------- #
@app.api(name="extract_receipt")
def api_extract_receipt(image_b64: str) -> dict:
    rec = extract_receipt(_open_image(image_b64))
    rec["reconcile"] = reconcile(rec)
    return rec


@app.api(name="extract_payment")
def api_extract_payment(image_b64: str) -> dict:
    return extract_payment(_open_image(image_b64))


@app.api(name="understand")
def api_understand(record_json: str) -> dict:
    """The 8B reasons over the bill: per-item + overall categories + summary."""
    record = _loads(record_json, {})
    record.setdefault("source", "receipt")
    try:
        record = understand(record)
    except Exception as e:  # pragma: no cover
        print(f"[app] understand failed: {e}")
        record.setdefault("category", "Other")
    return {"record": record, "reconcile": reconcile(record),
            "uncertain": record.get("_uncertain") or [],
            "understanding": record.get("understanding", "")}


# --------------------------------------------------------------------------- #
# Data (token-scoped)
# --------------------------------------------------------------------------- #
@app.api(name="save")
def api_save(token: str, record_json: str) -> dict:
    user = auth.verify_token(token)
    if not user:
        return {"ok": False, "message": "Please sign in again."}
    record = _clean(_loads(record_json, {}))
    if analytics._num(record.get("total", 0)) <= 0 and not record.get("line_items"):
        return {"ok": False, "message": "Nothing to save — add an amount first."}
    try:
        storage.save(user, record)
        return {"ok": True, "message": "Saved.",
                "count": len(storage.load(user, force=True))}
    except Exception as e:
        return {"ok": False, "message": f"Save failed: {e}"}


def _txn_dict(r: dict) -> dict:
    return {"date": r.get("date", ""), "vendor": r.get("vendor", ""),
            "total": analytics._num(r.get("total")), "category": analytics._category(r),
            "source": r.get("source", "receipt"), "understanding": r.get("understanding", ""),
            "currency": r.get("currency", ""),
            "items": [{"name": str(it.get("name", "")), "amount": analytics._num(it.get("amount")),
                       "qty": analytics._num(it.get("qty", 1)) or 1,
                       "category": str(it.get("category") or analytics._category(r))}
                      for it in (r.get("line_items") or [])],
            "charges": [{"label": str(ch.get("label", "")), "amount": analytics._num(ch.get("amount"))}
                        for ch in (r.get("charges") or [])]}


def _sort_recent(records):
    return sorted(records, key=lambda r: (analytics.parse_date(r.get("date")) or date.min,
                                          str(r.get("saved_at", ""))), reverse=True)


@app.api(name="dashboard")
def api_dashboard(token: str, start: str = "", end: str = "", category: str = "All",
                  granularity: str = "Monthly", cal_month: str = "") -> dict:
    user = auth.verify_token(token)
    if not user:
        return {"ok": False, "error": "auth"}
    try:
        records = storage.load(user, force=True)
        summ = analytics.summary(records)
        cat = None if category in ("", "All", None) else category
        filt = analytics.filter_records(
            records, analytics.parse_date(start), analytics.parse_date(end), cat)
        today = date.today()
        try:
            cy, cm = (int(x) for x in cal_month.split("-")) if cal_month else (today.year, today.month)
        except Exception:
            cy, cm = today.year, today.month
        return {
            "ok": True,
            "summary": summ,
            "budget": storage.get_budget(user),
            "by_category": analytics.spend_by_category(filt),
            "over_time": analytics.spend_over_time(filt, granularity or "Monthly"),
            "comparison": analytics.category_comparison(records),
            "calendar": analytics.calendar_data(records, cy, cm),
            "calendar_month": f"{cy}-{cm:02d}",
            "transactions": [_txn_dict(r) for r in _sort_recent(filt)],
            "recent": [_txn_dict(r) for r in _sort_recent(records)[:8]],
            "has_any": len(records) > 0,
        }
    except Exception as e:
        import traceback
        print(f"[app] dashboard failed: {e}\n{traceback.format_exc()[-800:]}")
        return {"ok": False, "error": str(e)}


@app.api(name="set_budget")
def api_set_budget(token: str, amount: float) -> dict:
    user = auth.verify_token(token)
    if not user:
        return {"ok": False}
    try:
        return {"ok": True, "budget": storage.set_budget(user, amount)}
    except Exception as e:
        return {"ok": False, "message": str(e)}


@app.api(name="chat")
def api_chat(token: str, message: str, history_json: str = "") -> dict:
    user = auth.verify_token(token)
    if not user:
        return {"reply": "Please sign in again.", "trace": []}
    records = storage.load(user, force=True)
    try:
        budget = storage.get_budget(user)
    except Exception:
        budget = 0.0
    result = agent.run(message, records, budget=budget)
    return {"reply": result["reply"], "trace": result.get("trace", [])}


@app.api(name="categories")
def api_categories() -> dict:
    return {"categories": CATEGORIES}


# --------------------------------------------------------------------------- #
# Static custom frontend
# --------------------------------------------------------------------------- #
if (FRONTEND / "assets").exists():
    app.mount("/assets", StaticFiles(directory=str(FRONTEND / "assets")), name="assets")


@app.get("/", response_class=HTMLResponse)
async def home():
    html = (FRONTEND / "index.html").read_text(encoding="utf-8")
    try:
        v = int((FRONTEND / "assets" / "app.js").stat().st_mtime)
    except Exception:
        v = 1
    return html.replace("/assets/app.js", f"/assets/app.js?v={v}")


if __name__ == "__main__":
    app.launch(show_error=True)