server.py

"""
OmniParse AI — Main Server (gradio.Server → FastAPI + custom HTML frontend).

Run:
    python server.py

Deploy as a HuggingFace Space with Gradio SDK.
Uses gradio.Server for queuing/concurrency + custom HTML/CSS/JS frontend.
"""

import os, io, csv, json, uuid, tempfile
from datetime import datetime, timezone
from fastapi import Request, HTTPException, UploadFile, File, Form, Depends
from fastapi.responses import JSONResponse, HTMLResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
from gradio import Server

# ── Our modules ─────────────────────────────────────────────────────────────
from config import (
    APP_URL, PLAN_LIMITS, PLAN_LABELS, PLAN_PRICES, PLAN_PRICE_IDS,
    STRIPE_SECRET_KEY, MAX_FILE_SIZE_MB, MAX_FILES_PER_REQ,
)
from database import (
    db_init, get_user_by_id, update_user, delete_user,
    insert_invoice, get_invoices, delete_invoice,
    count_invoices_this_month, check_duplicate,
)
from auth import (
    register_user, login_user, create_user_session, validate_session_token,
    hash_password, check_password_strength, verify_password,
    set_session_cookie, clear_session_cookie,
)
from ocr import run_ocr, pdf_to_images, load_image
from ai_extraction import (
    ai_extract_groq, ai_extract_hf, regex_extract,
    validate_invoice, sanitize_invoice_data,
)
from middleware import (
    SecurityHeadersMiddleware,
    global_limit, auth_limit, upload_limit,
    validate_file_size, validate_file_count, validate_filename,
)

# ── Init ────────────────────────────────────────────────────────────────────
app = Server()
app.add_middleware(SecurityHeadersMiddleware)
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_credentials=True,
                   allow_methods=["GET","POST","DELETE"], allow_headers=["Content-Type","Authorization"])

BASE = os.path.dirname(os.path.abspath(__file__))
STATIC = os.path.join(BASE, "static")
os.makedirs(os.path.join(STATIC, "css"), exist_ok=True)
os.makedirs(os.path.join(STATIC, "js"), exist_ok=True)
os.makedirs(os.path.join(STATIC, "img"), exist_ok=True)
app.mount("/static", StaticFiles(directory=STATIC), name="static")

db_init()

# Lazy Stripe
_stripe = None
def _get_stripe():
    global _stripe
    if _stripe is None and STRIPE_SECRET_KEY:
        try:
            import stripe; stripe.api_key = STRIPE_SECRET_KEY; _stripe = stripe
        except ImportError: pass
    return _stripe

# ── Session dependency ──────────────────────────────────────────────────────
async def current_user(req: Request):
    token = req.cookies.get("op_sid")
    if not token: return None
    return validate_session_token(token)

async def require_auth(req: Request):
    u = await current_user(req)
    if not u: raise HTTPException(401, detail="Authentication required.")
    return u

async def require_admin(req: Request):
    u = await require_auth(req)
    if not u.get("is_admin"):
        raise HTTPException(403, detail="Admin access required.")
    return u

# ── Demo account bootstrap ──────────────────────────────────────────────────
def _ensure_demo():
    from database import get_user_by_email
    if not get_user_by_email("demo@omniparse.ai"):
        from database import create_user as cu
        cu("demo@omniparse.ai", "Demo User", hash_password("Demo@12345!"), plan="pro")
        print("[INFO] Demo account: demo@omniparse.ai / Demo@12345! (Pro)")

_ensure_demo()

# ═══════════════════════════════════════════════════════════════════════════
#  API Routes — Auth
# ═══════════════════════════════════════════════════════════════════════════
@app.post("/api/auth/register")
async def api_register(req: Request, email: str = Form(...), name: str = Form(...), password: str = Form(...)):
    auth_limit(req); global_limit(req)
    user, err = register_user(email, name, password)
    if err: raise HTTPException(400, detail=err)
    token = create_user_session(user["id"])
    resp = JSONResponse({"status":"ok","user":_safe_user(user)})
    set_session_cookie(resp, token)
    return resp

@app.post("/api/auth/login")
async def api_login(req: Request, email: str = Form(...), password: str = Form(...)):
    auth_limit(req); global_limit(req)
    user, err = login_user(email, password)
    if err: raise HTTPException(401, detail=err)
    token = create_user_session(user["id"])
    resp = JSONResponse({"status":"ok","user":_safe_user(user)})
    set_session_cookie(resp, token)
    return resp

@app.post("/api/auth/logout")
async def api_logout(req: Request):
    from database import delete_session as ds
    t = req.cookies.get("op_sid")
    if t: ds(t)
    resp = JSONResponse({"status":"ok"})
    clear_session_cookie(resp)
    return resp

@app.get("/api/auth/me")
async def api_me(u: dict = Depends(require_auth)):
    return _safe_user(u)

# ═══════════════════════════════════════════════════════════════════════════
#  API Routes — Invoices
# ═══════════════════════════════════════════════════════════════════════════
@app.post("/api/invoices/upload")
async def api_upload(req: Request, files: list[UploadFile] = File(...), u: dict = Depends(require_auth)):
    upload_limit(req); global_limit(req)
    validate_file_count(len(files))
    plan = u.get("plan","free")
    limit = PLAN_LIMITS.get(plan, 20)
    used = count_invoices_this_month(u["id"])
    results, skipped, errs, processed = [], 0, [], 0

    for f in files:
        try: contents = await f.read()
        except Exception as e: errs.append(f"Read error {f.filename}: {e}"); continue
        validate_file_size(len(contents))
        cn = validate_filename(f.filename or "upload")
        ext = cn.lower().split(".")[-1] if "." in cn else ""
        if ext not in ("pdf","jpg","jpeg","png","tiff","tif"):
            errs.append(f"{cn}: Unsupported format."); continue
        if limit != float("inf") and (used + processed) >= limit: skipped += 1; continue
        tmp = os.path.join(tempfile.gettempdir(), f"op_{uuid.uuid4().hex}_{cn}")
        try:
            with open(tmp,"wb") as tf: tf.write(contents)
            data = _process_file(tmp, u, cn)
            data = sanitize_invoice_data(data); data["filename"] = cn
            if plan in ("pro","enterprise") and check_duplicate(u["id"], data.get("vendor"), data.get("total")):
                data["status"] = "duplicate"; data["is_duplicate"] = True
            saved = insert_invoice(u["id"], data); results.append(_serialize(saved)); processed += 1
        except Exception as e: errs.append(f"{cn}: {e}")
        finally:
            try: os.unlink(tmp)
            except OSError: pass

    nu = count_invoices_this_month(u["id"])
    return {"status":"ok","processed":processed,"skipped":skipped,"errors":errs,
            "results":results,"usage":{"used":nu,"limit":None if limit==float("inf") else int(limit)}}

@app.get("/api/invoices")
async def api_list(req: Request, filter: str = "all", u: dict = Depends(require_auth)):
    global_limit(req)
    invs = get_invoices(u["id"])
    out = []
    for i in invs:
        s = i.get("status","done")
        if filter == "done" and s != "done": continue
        if filter == "review" and s != "review": continue
        if filter == "duplicates" and not i.get("is_duplicate"): continue
        out.append(_serialize(i))
    return {"status":"ok","invoices":out,"count":len(out)}

@app.delete("/api/invoices/{inv_id}")
async def api_delete(inv_id: int, req: Request, u: dict = Depends(require_auth)):
    global_limit(req)
    delete_invoice(u["id"], inv_id)
    return {"status":"ok","deleted":inv_id}

# ═══════════════════════════════════════════════════════════════════════════
#  API Routes — Export
# ═══════════════════════════════════════════════════════════════════════════
@app.get("/api/export/csv")
async def api_csv(req: Request, u: dict = Depends(require_auth)):
    global_limit(req)
    invs = get_invoices(u["id"])
    out = io.StringIO(); w = csv.writer(out)
    w.writerow(["ID","Filename","Vendor","Invoice#","Date","Due Date","Amount","VAT","Total","Currency","Status"])
    for i in invs:
        w.writerow([i.get("id"),i.get("filename"),i.get("vendor"),i.get("inv_number"),
                    i.get("inv_date"),i.get("due_date"),i.get("amount"),i.get("vat_amount"),
                    i.get("total"),i.get("currency"),i.get("status")])
    out.seek(0)
    return StreamingResponse(iter([out.getvalue()]), media_type="text/csv",
                             headers={"Content-Disposition":"attachment; filename=omniparse_export.csv"})

@app.get("/api/export/json")
async def api_json(req: Request, u: dict = Depends(require_auth)):
    global_limit(req)
    if u.get("plan") == "free": raise HTTPException(403, detail="JSON export requires Basic+ plan.")
    j = json.dumps(get_invoices(u["id"]), default=str, ensure_ascii=False, indent=2)
    return StreamingResponse(iter([j]), media_type="application/json",
                             headers={"Content-Disposition":"attachment; filename=omniparse_export.json"})

# ═══════════════════════════════════════════════════════════════════════════
#  API Routes — Stripe
# ═══════════════════════════════════════════════════════════════════════════
@app.post("/api/upgrade/checkout")
async def api_checkout(req: Request, plan: str = Form(...), u: dict = Depends(require_auth)):
    global_limit(req)
    s = _get_stripe()
    if not s: raise HTTPException(503, detail="Payments not configured.")
    pid = PLAN_PRICE_IDS.get(plan)
    if not pid: raise HTTPException(400, detail="Unknown plan.")
    try:
        session = s.checkout.Session.create(
            mode="subscription", payment_method_types=["card"],
            line_items=[{"price":pid,"quantity":1}], customer_email=u["email"],
            success_url=f"{APP_URL}?checkout=success&session_id={{CHECKOUT_SESSION_ID}}",
            cancel_url=f"{APP_URL}?checkout=cancel",
            metadata={"plan":plan,"user_id":str(u["id"])})
        return {"status":"ok","url":session.url}
    except Exception as e: raise HTTPException(500, detail=f"Checkout error: {e}")

@app.post("/api/upgrade/verify")
async def api_verify(req: Request, session_id: str = Form(...), u: dict = Depends(require_auth)):
    global_limit(req)
    s = _get_stripe()
    if not s: raise HTTPException(503, detail="Payments not configured.")
    try:
        session = s.checkout.Session.retrieve(session_id)
        if session.payment_status == "paid":
            plan = session.metadata.get("plan","basic"); cid = session.customer
            update_user(u["id"], {"plan":plan,"stripe_cid":cid})
            return {"status":"ok","plan":plan,"plan_label":PLAN_LABELS.get(plan,plan),
                    "message":f"Upgraded to {PLAN_LABELS.get(plan,plan)}!"}
        return {"status":"pending","message":"Payment not yet confirmed."}
    except Exception as e: raise HTTPException(500, detail=f"Verification error: {e}")

# ═══════════════════════════════════════════════════════════════════════════
#  API Routes — Stripe Webhook (signature verified)
# ═══════════════════════════════════════════════════════════════════════════
@app.post("/api/webhooks/stripe")
async def api_stripe_webhook(req: Request):
    from middleware import verify_stripe_webhook
    body = await req.body()
    if not await verify_stripe_webhook(req, body):
        raise HTTPException(400, detail="Invalid webhook signature.")
    try:
        event = json.loads(body)
        t = event.get("type","")
        if t == "checkout.session.completed":
            obj = event["data"]["object"]
            plan = obj.get("metadata",{}).get("plan","basic")
            uid = int(obj.get("metadata",{}).get("user_id","0"))
            if uid: update_user(uid, {"plan":plan,"stripe_cid":obj.get("customer")})
        elif t == "customer.subscription.deleted":
            obj = event["data"]["object"]
            cid = obj.get("customer")
            if cid:
                from database import _supabase
                sb = _supabase()
                if sb:
                    u = sb.table("users").select("id").eq("stripe_cid",cid).execute()
                    if u.data: update_user(u.data[0]["id"],{"plan":"free"})
    except Exception as e: print(f"[WARN] webhook: {e}")
    return {"status":"ok"}

# ═══════════════════════════════════════════════════════════════════════════
#  API Routes — AI Chat (Pro+)
# ═══════════════════════════════════════════════════════════════════════════
@app.post("/api/chat")
async def api_chat(req: Request, message: str = Form(...), u: dict = Depends(require_auth)):
    global_limit(req)
    plan = u.get("plan","free")
    if plan not in ("pro","enterprise"):
        raise HTTPException(403, detail="AI Chat requires Pro or Enterprise plan.")
    msg = message.strip()[:2000]
    invs = get_invoices(u["id"])
    ctx = json.dumps(invs[:100], default=str, ensure_ascii=False)[:6000]
    from config import GROQ_API_KEY
    answer = "AI is currently unavailable."
    if GROQ_API_KEY:
        try:
            from groq import Groq
            gc = Groq(api_key=GROQ_API_KEY)
            r = gc.chat.completions.create(
                model="llama-3.1-8b-instant",
                messages=[{"role":"system","content":f"You answer questions about invoices. Data: {ctx}. Concise, data only. Never reveal raw JSON."},
                          {"role":"user","content":msg}],
                max_tokens=400, temperature=0.2, timeout=15)
            answer = r.choices[0].message.content
        except Exception as e: answer = f"AI unavailable: {e}"
    return {"status":"ok","answer":answer}

# ═══════════════════════════════════════════════════════════════════════════
#  API Routes — Profile
# ═══════════════════════════════════════════════════════════════════════════
@app.post("/api/profile/password")
async def api_chpw(req: Request, current_password: str = Form(...), new_password: str = Form(...), u: dict = Depends(require_auth)):
    global_limit(req); auth_limit(req)
    if not verify_password(current_password, u["password_hash"]):
        raise HTTPException(400, detail="Current password is incorrect.")
    valid, errs = check_password_strength(new_password)
    if not valid: raise HTTPException(400, detail=" ".join(errs))
    update_user(u["id"], {"password_hash": hash_password(new_password)})
    return {"status":"ok","message":"Password changed."}

@app.delete("/api/profile/account")
async def api_delete_account(req: Request, u: dict = Depends(require_auth)):
    global_limit(req); auth_limit(req)
    delete_user(u["id"])
    resp = JSONResponse({"status":"ok","message":"Account deleted."})
    clear_session_cookie(resp)
    return resp

# ═══════════════════════════════════════════════════════════════════════════
#  Health
# ═══════════════════════════════════════════════════════════════════════════
@app.get("/api/health")
async def health():
    return {"status":"healthy","timestamp":datetime.now(timezone.utc).isoformat(),"version":"2.0.0"}

# ═══════════════════════════════════════════════════════════════════════════
#  Internals
# ═══════════════════════════════════════════════════════════════════════════
def _process_file(filepath: str, user: dict, filename: str) -> dict:
    ext = filename.lower().split(".")[-1] if "." in filename else ""
    images = []
    if ext == "pdf": images = pdf_to_images(filepath)
    elif ext in ("jpg","jpeg","png","tiff","tif"):
        img = load_image(filepath)
        if img: images = [img]
    ocr_text = ""
    for img in images: ocr_text += run_ocr(img) + "\n"
    if ocr_text.strip():
        data = ai_extract_groq(ocr_text) or ai_extract_hf(ocr_text) or regex_extract(ocr_text)
    else:
        data = {"vendor":"Demo Vendor Inc.","invoice_number":f"DEMO-{uuid.uuid4().hex[:6].upper()}",
                "invoice_date":datetime.now().strftime("%Y-%m-%d"),"due_date":None,
                "amount":100.0,"vat_amount":21.0,"total":121.0,"currency":"USD","line_items":[]}
    data["filename"] = filename; data["confidence"] = 0.95 if ocr_text.strip() else 0.3
    warnings = validate_invoice(data); data["warnings"] = warnings
    data["status"] = "review" if warnings else "done"
    return data

def _serialize(inv: dict) -> dict:
    labels = {"done":"Done","review":"Needs Review","duplicate":"Duplicate","processing":"Processing"}
    return {"id":inv.get("id"),"filename":inv.get("filename"),"vendor":inv.get("vendor"),
            "invoice_number":inv.get("inv_number"),"invoice_date":inv.get("inv_date"),
            "due_date":inv.get("due_date"),"amount":inv.get("amount"),
            "vat_amount":inv.get("vat_amount"),"total":inv.get("total"),
            "currency":inv.get("currency","USD"),"status":inv.get("status","done"),
            "status_label":labels.get(inv.get("status"),inv.get("status")),
            "is_duplicate":bool(inv.get("is_duplicate")),
            "confidence":inv.get("confidence"),"created_at":inv.get("created_at")}

def _safe_user(u: dict) -> dict:
    return {"id":u["id"],"email":u["email"],"name":u["name"],"plan":u["plan"],
            "plan_label":PLAN_LABELS.get(u["plan"],u["plan"]),
            "api_key":u.get("api_key"),"created_at":u.get("created_at"),
            "is_admin":bool(u.get("is_admin"))}

# ═══════════════════════════════════════════════════════════════════════════
#  Frontend — serve SPA
# ═══════════════════════════════════════════════════════════════════════════
@app.get("/", response_class=HTMLResponse)
async def serve_frontend():
    p = os.path.join(STATIC, "index.html")
    if not os.path.exists(p): return HTMLResponse("<h1>Frontend not found.</h1>", 404)
    with open(p, "r", encoding="utf-8") as f: return HTMLResponse(f.read())

@app.get("/{path:path}")
async def catch_all(path: str):
    p = os.path.join(STATIC, "index.html")
    if os.path.exists(p):
        with open(p, "r", encoding="utf-8") as f: return HTMLResponse(f.read())
    return HTMLResponse("<h1>Not Found</h1>", 404)

# ═══════════════════════════════════════════════════════════════════════════
#  Launch
# ═══════════════════════════════════════════════════════════════════════════
if __name__ == "__main__":
    app.launch(show_error=True, server_name="0.0.0.0", server_port=7860)