app.py

"""
OmniParse AI — Invoice processing SaaS
Refactored, modular Gradio app for HuggingFace Spaces.

Views: Landing / Pricing / Legal / Auth / Dashboard
All AI/OCR/DB layers degrade gracefully if a given secret/service is missing.

Security: bcrypt passwords, rate limiting, server-side validation,
          file size limits, input sanitization, no raw HTML from users.
UI: Custom SVG icons, skeleton loaders, tooltips, dark mode,
    toast notifications, drag-and-drop upload, semantic HTML.
"""

import os
import csv
import json
from datetime import datetime, timezone

import gradio as gr

# ---------------------------------------------------------------------------
# Module imports — all with graceful fallback
# ---------------------------------------------------------------------------
try:
    from config import (
        APP_URL, GROQ_CLIENT, PLAN_LIMITS, PLAN_LABELS, PLAN_PRICES,
        MAX_FILE_SIZE_BYTES, ALLOWED_EXTENSIONS, SESSION_TTL_HOURS,
    )
except ImportError:
    APP_URL = os.environ.get("APP_URL", "http://localhost:7860")
    GROQ_CLIENT = None
    PLAN_LIMITS = {"free": 20, "basic": 200, "pro": 2000, "enterprise": float("inf")}
    PLAN_LABELS = {"free": "Free", "basic": "Basic", "pro": "Pro", "enterprise": "Enterprise"}
    PLAN_PRICES = {"free": 0, "basic": 29, "pro": 129, "enterprise": 499}
    MAX_FILE_SIZE_BYTES = 20 * 1024 * 1024
    ALLOWED_EXTENSIONS = {"pdf", "jpg", "jpeg", "png", "tiff", "tif"}
    SESSION_TTL_HOURS = 168

try:
    from database import (
        db_init, get_user_by_email, get_user_by_id, update_user,
        get_session_user, delete_session,
        insert_invoice, get_invoices, count_invoices_this_month,
        delete_invoice,
    )
except ImportError:
    db_init = lambda: None
    get_user_by_email = lambda e: None
    get_user_by_id = lambda i: None
    update_user = lambda *a: None
    get_session_user = lambda t: None
    delete_session = lambda t: None
    insert_invoice = lambda *a: {}
    get_invoices = lambda u: []
    count_invoices_this_month = lambda u: 0
    delete_invoice = lambda *a: None

try:
    from auth import (
        create_user, authenticate_user, gen_api_key,
        create_session, hash_pw,
    )
except ImportError:
    create_user = lambda *a: (None, "Auth module not available")
    authenticate_user = lambda *a: (None, "Auth module not available")
    gen_api_key = lambda: "op_live_fallback"
    create_session = lambda u: "fallback_token"
    hash_pw = lambda p: p

try:
    from validation import validate_email, validate_password, validate_name, validate_file
except ImportError:
    validate_email = lambda e: (bool(e), "")
    validate_password = lambda p: (len(p) >= 8, "Password too short" if len(p) < 8 else "")
    validate_name = lambda n: (bool(n), "")
    validate_file = lambda f: (True, "")

try:
    from rate_limiter import check_auth_rate, check_upload_rate, check_chat_rate
except ImportError:
    check_auth_rate = lambda ip: True
    check_upload_rate = lambda ip: True
    check_chat_rate = lambda ip: True

try:
    import icons as ic
except ImportError:
    ic = None

try:
    import html_templates as tmpl
except ImportError:
    tmpl = None

try:
    from ocr_pipeline import process_file as ocr_process_file
except ImportError:
    ocr_process_file = None

try:
    from stripe_handler import create_checkout_session, verify_checkout_session
except ImportError:
    create_checkout_session = lambda *a: (None, "Payments not configured")
    verify_checkout_session = lambda *a: ("Verification unavailable", None)

try:
    from invoice_processor import (
        validate_invoice_data, _parse_date_any, process_and_save, batch_process,
    )
except ImportError:
    validate_invoice_data = lambda d: []
    _parse_date_any = lambda s: None
    process_and_save = lambda *a: ({}, {})
    batch_process = lambda *a: ([], [], 0, 0)

# ---------------------------------------------------------------------------
# Initialize
# ---------------------------------------------------------------------------
db_init()
try:
    from auth import ensure_demo_account
    ensure_demo_account()
except ImportError:
    pass

# ---------------------------------------------------------------------------
# Helper: get client IP from Gradio request
# ---------------------------------------------------------------------------
def _get_ip(request: gr.Request) -> str:
    try:
        return request.client.host if request and request.client else "unknown"
    except Exception:
        return "unknown"


# ---------------------------------------------------------------------------
# View management
# ---------------------------------------------------------------------------
ALL_VIEWS = []  # populated after Blocks creation


def show_only(idx):
    return [gr.update(visible=(i == idx)) for i in range(len(ALL_VIEWS))]


# ---------------------------------------------------------------------------
# Navigation callbacks
# ---------------------------------------------------------------------------
def go_landing():
    return show_only(0)


def go_pricing():
    return show_only(1)


def go_legal():
    return show_only(2)


def go_auth():
    return show_only(3)


def go_dashboard():
    return show_only(4)


# ---------------------------------------------------------------------------
# Auth callbacks
# ---------------------------------------------------------------------------
def do_login(email, password, request: gr.Request):
    ip = _get_ip(request)
    if not check_auth_rate(ip):
        return (
            gr.update(value="Too many login attempts. Please wait a moment.", visible=True),
            None, None,
            *show_only(3), gr.update(), gr.update(), gr.update(), gr.update(),
        )
    if not email or not password:
        return (
            gr.update(value="Please enter both email and password.", visible=True),
            None, None,
            *show_only(3), gr.update(), gr.update(), gr.update(), gr.update(),
        )
    user, err = authenticate_user(email, password, ip)
    if err:
        return (
            gr.update(value=err, visible=True),
            None, None,
            *show_only(3), gr.update(), gr.update(), gr.update(), gr.update(),
        )
    token = create_session(str(user["id"]))
    return (
        gr.update(value="", visible=False),
        token, user,
        *show_only(4),
        gr.update(visible=False), gr.update(visible=False),
        gr.update(visible=True), gr.update(visible=True),
    )


def do_signup(name, email, password, terms, request: gr.Request):
    ip = _get_ip(request)
    if not check_auth_rate(ip):
        return (
            gr.update(value="Too many registration attempts. Please wait a moment.", visible=True),
            None, None,
            *show_only(3), gr.update(), gr.update(), gr.update(), gr.update(),
        )
    # Server-side validation
    valid, msg = validate_name(name)
    if not valid:
        return (
            gr.update(value=f"Invalid name: {msg}", visible=True),
            None, None,
            *show_only(3), gr.update(), gr.update(), gr.update(), gr.update(),
        )
    valid, msg = validate_email(email)
    if not valid:
        return (
            gr.update(value=f"Invalid email: {msg}", visible=True),
            None, None,
            *show_only(3), gr.update(), gr.update(), gr.update(), gr.update(),
        )
    valid, pw_msgs = validate_password(password)
    if not valid:
        details = "; ".join(pw_msgs)
        return (
            gr.update(value=f"Password requirements not met: {details}", visible=True),
            None, None,
            *show_only(3), gr.update(), gr.update(), gr.update(), gr.update(),
        )
    if not terms:
        return (
            gr.update(value="You must agree to the Terms of Use and Privacy Policy.", visible=True),
            None, None,
            *show_only(3), gr.update(), gr.update(), gr.update(), gr.update(),
        )
    user, err = create_user(email, name, password)
    if err:
        return (
            gr.update(value=err, visible=True),
            None, None,
            *show_only(3), gr.update(), gr.update(), gr.update(), gr.update(),
        )
    token = create_session(str(user["id"]))
    return (
        gr.update(value="", visible=False),
        token, user,
        *show_only(4),
        gr.update(visible=False), gr.update(visible=False),
        gr.update(visible=True), gr.update(visible=True),
    )


def do_logout(token):
    if token:
        delete_session(token)
    return (
        None, None,
        *show_only(0),
        gr.update(visible=True), gr.update(visible=True),
        gr.update(visible=False), gr.update(visible=False),
    )


# ---------------------------------------------------------------------------
# Dashboard load / usage
# ---------------------------------------------------------------------------
def load_dashboard(user):
    if not user:
        return "## Dashboard\n\n_You are not signed in._", ""
    used = count_invoices_this_month(user["id"])
    limit = PLAN_LIMITS.get(user["plan"], 20)
    limit_str = "\u221e" if limit == float("inf") else str(int(limit))
    plan_label = PLAN_LABELS.get(user["plan"], user["plan"])
    welcome = f"## Dashboard\n\nWelcome back, **{user['name']}** (Plan: {plan_label})"
    usage = f"**{used}/{limit_str} invoices used this month**"
    if limit != float("inf") and used >= limit:
        usage += "\n\n**Monthly limit reached.** Go to the Upgrade tab to increase your limit."
    return welcome, usage


# ---------------------------------------------------------------------------
# Upload / processing
# ---------------------------------------------------------------------------
def do_upload(files, user, request: gr.Request):
    ip = _get_ip(request)
    if not user:
        return "You must be signed in to upload invoices.", [], None, ""
    if not files:
        return "No files selected. Please select invoice files to process.", [], None, ""
    if not check_upload_rate(ip):
        return "Upload rate limit reached. Please wait a moment.", [], None, ""

    limit = PLAN_LIMITS.get(user["plan"], 20)
    used = count_invoices_this_month(user["id"])
    rows = []
    last_json = None
    processed = 0

    for f in files:
        if used + processed >= limit:
            break
        path = f.name if hasattr(f, "name") else f
        filename = os.path.basename(path)

        # Server-side file validation
        valid, err = validate_file(path)
        if not valid:
            rows.append([filename, "Validation Error", None, None, None, err])
            continue

        try:
            if ocr_process_file:
                data = ocr_process_file(path, user)
            else:
                data = {
                    "filename": filename, "vendor": "Demo Vendor Inc.",
                    "invoice_number": "DEMO-0001",
                    "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": [],
                }
        except Exception as e:
            data = {
                "filename": filename, "vendor": None, "invoice_number": None,
                "invoice_date": None, "due_date": None, "amount": None,
                "vat_amount": None, "total": None, "currency": "USD",
                "status": "review", "confidence": 0,
                "warnings": [f"Processing error: {str(e)[:200]}"],
            }

        # Validate extracted data
        warnings = validate_invoice_data(data)
        if warnings:
            data["warnings"] = data.get("warnings", []) + warnings
            if data.get("status") != "duplicate":
                data["status"] = "review"

        saved = insert_invoice(user["id"], data)
        last_json = data
        status_label = data.get("status", "done")
        status_display = {
            "done": "Done", "review": "Review",
            "duplicate": "Duplicate", "processing": "Processing",
        }.get(status_label, status_label)
        rows.append([
            data.get("filename"), data.get("vendor"),
            data.get("invoice_number"), data.get("invoice_date"),
            data.get("total"), status_display,
        ])
        processed += 1

    skipped = len(files) - processed
    msg = f"Successfully processed {processed} invoice{'s' if processed != 1 else ''}."
    if skipped > 0:
        msg += f" {skipped} file{'s' if skipped != 1 else ''} skipped \u2014 monthly limit reached. Upgrade your plan to process more."

    used_new = count_invoices_this_month(user["id"])
    limit_str = "\u221e" if limit == float("inf") else str(int(limit))
    usage = f"**{used_new}/{limit_str} invoices used this month**"
    return msg, rows, last_json, usage


# ---------------------------------------------------------------------------
# My Invoices
# ---------------------------------------------------------------------------
STATUS_MAP = {
    "done": "Done", "review": "Review",
    "duplicate": "Duplicate", "processing": "Processing",
}


def refresh_invoices(user, flt):
    if not user:
        return []
    invoices = get_invoices(user["id"])
    rows = []
    for inv in invoices:
        status = inv.get("status", "done")
        if flt == "Done" and status != "done":
            continue
        if flt == "Review" and status != "review":
            continue
        if flt == "Duplicates" and not inv.get("is_duplicate"):
            continue
        rows.append([
            inv.get("id"), inv.get("vendor"), inv.get("inv_number"),
            inv.get("inv_date"), inv.get("total"),
            STATUS_MAP.get(status, status),
        ])
    return rows


def do_delete_invoice(user, inv_id):
    if not user or not inv_id:
        return "Please enter a valid invoice ID.", []
    inv_id = int(inv_id)
    delete_invoice(user["id"], inv_id)
    return f"Invoice #{inv_id} has been deleted.", refresh_invoices(user, "All")


# ---------------------------------------------------------------------------
# AI Chat
# ---------------------------------------------------------------------------
def chat_respond(message, history, user, request: gr.Request):
    history = history or []
    if not user:
        history.append({"role": "assistant", "content": "You must be signed in to use AI Chat."})
        return history, ""
    if not check_chat_rate(_get_ip(request)):
        history.append({"role": "assistant", "content": "Chat rate limit reached. Please wait a moment."})
        return history, ""
    if user["plan"] not in ("pro", "enterprise"):
        history.append({
            "role": "assistant",
            "content": "AI Chat is available on the Pro plan and above. Upgrade from the Upgrade tab to unlock this feature.",
        })
        return history, ""
    if not message or not message.strip():
        return history, ""

    invoices = get_invoices(user["id"])
    context = json.dumps(invoices[:100], default=str, ensure_ascii=False)[:6000]
    history.append({"role": "user", "content": message})

    if GROQ_CLIENT:
        try:
            resp = GROQ_CLIENT.chat.completions.create(
                model="llama-3.1-8b-instant",
                messages=[
                    {
                        "role": "system",
                        "content": (
                            "You are an assistant answering questions about the user's invoices. "
                            f"Here is their invoice data as JSON: {context}. "
                            "Answer concisely based only on this data. Use plain text, no markdown."
                        ),
                    },
                    {"role": "user", "content": message},
                ],
                max_tokens=400, temperature=0.2, timeout=10,
            )
            answer = resp.choices[0].message.content
        except Exception as e:
            answer = f"AI is temporarily unavailable ({str(e)[:100]}). Please try again."
    else:
        answer = "AI Chat requires the GROQ_API_KEY to be configured."

    history.append({"role": "assistant", "content": answer})
    return history, ""


# ---------------------------------------------------------------------------
# Export
# ---------------------------------------------------------------------------
def export_csv(user):
    if not user:
        return None
    invoices = get_invoices(user["id"])
    path = f"/tmp/omniparse_export_{user['id']}.csv"
    with open(path, "w", newline="", encoding="utf-8") as f:
        writer = csv.writer(f)
        writer.writerow([
            "ID", "Vendor", "Invoice #", "Date", "Due Date",
            "Amount", "VAT", "Total", "Currency", "Status",
        ])
        for inv in invoices:
            writer.writerow([
                inv.get("id"), inv.get("vendor"), inv.get("inv_number"),
                inv.get("inv_date"), inv.get("due_date"), inv.get("amount"),
                inv.get("vat_amount"), inv.get("total"), inv.get("currency"),
                inv.get("status"),
            ])
    return path


def export_json(user):
    if not user:
        return None
    if user["plan"] == "free":
        return None
    invoices = get_invoices(user["id"])
    path = f"/tmp/omniparse_export_{user['id']}.json"
    with open(path, "w", encoding="utf-8") as f:
        json.dump(invoices, f, default=str, ensure_ascii=False, indent=2)
    return path


# ---------------------------------------------------------------------------
# Upgrade / Stripe
# ---------------------------------------------------------------------------
def do_upgrade(plan, user):
    if not user:
        return "You must be signed in to upgrade."
    url, err = create_checkout_session(plan, user)
    if err:
        return err
    return (
        f"[Complete your payment via Stripe]({url})\n\n"
        "After paying, return here and paste the `session_id` from the URL below."
    )


def do_check_payment(session_id, user):
    if not user or not session_id:
        return "Please enter a valid session ID."
    msg, err = verify_checkout_session(session_id, user["id"])
    if err:
        return err
    return msg


# ---------------------------------------------------------------------------
# API tab
# ---------------------------------------------------------------------------
def load_api_tab(user):
    if not user:
        return (
            gr.update(visible=True, value="You must be signed in."),
            gr.update(visible=False), gr.update(visible=False),
        )
    if user["plan"] not in ("pro", "enterprise"):
        return (
            gr.update(
                visible=True,
                value="API access is available on the Pro plan and above. "
                      "Upgrade from the Upgrade tab to get your API key.",
            ),
            gr.update(visible=False), gr.update(visible=False),
        )
    key = user.get("api_key") or "\u2014"
    snippet = (
        f"```bash\n"
        f"curl -X POST {APP_URL}/api/extract \\\n"
        f'  -H "Authorization: Bearer {key}" \\\n'
        f"  -F \"file=@invoice.pdf\"\n"
        f"```"
    )
    return (
        gr.update(visible=False),
        gr.update(visible=True, value=f"**Your API Key:** `{key}`"),
        gr.update(visible=True, value=snippet),
    )


# ---------------------------------------------------------------------------
# Profile
# ---------------------------------------------------------------------------
def load_profile(user):
    if not user:
        return "You are not signed in."
    plan = PLAN_LABELS.get(user["plan"], user["plan"])
    price = PLAN_PRICES.get(user["plan"], 0)
    return (
        f"**Name:** {user['name']}\n\n"
        f"**Email:** {user['email']}\n\n"
        f"**Plan:** {plan} (${price}/mo)\n\n"
        f"**Created:** {(user.get('created_at') or '')[:10]}"
    )


def do_change_password(user, new_pw, request: gr.Request):
    if not user:
        return "You must be signed in."
    valid, pw_msgs = validate_password(new_pw)
    if not valid:
        details = "; ".join(pw_msgs)
        return f"Password requirements not met: {details}"
    update_user(str(user["id"]), {"password": hash_pw(new_pw)})
    return "Password has been changed successfully."


def do_delete_account(user, token):
    if not user:
        return "You must be signed in.", None, None
    uid = str(user["id"])
    try:
        from database import _sqlite_conn
        conn = _sqlite_conn()
        conn.execute("DELETE FROM invoices WHERE user_id=?", (uid,))
        conn.execute("DELETE FROM users WHERE id=?", (uid,))
        conn.commit()
        conn.close()
    except Exception:
        pass
    if token:
        delete_session(token)
    return "Your account has been permanently deleted.", None, None


# ---------------------------------------------------------------------------
# Build Gradio App
# ---------------------------------------------------------------------------
CUSTOM_CSS = tmpl.BASE_CSS if tmpl else ""

with gr.Blocks(
    title="OmniParse AI",
    css=CUSTOM_CSS,
    theme=gr.themes.Soft(primary_hue="violet"),
    head=(
        (tmpl.ALL_JS if tmpl else "") +
        (tmpl.DARK_MODE_JS if tmpl else "")
    ),
) as demo:

    session_token = gr.State(None)
    current_user = gr.State(None)

    # ===================== NAVBAR =====================
    with gr.Row(elem_id="op-navbar"):
        if tmpl:
            gr.HTML(tmpl.navbar_html())
        else:
            gr.HTML('<div class="op-logo"><span>OmniParse AI</span></div>')
        with gr.Row():
            nav_pricing_btn = gr.Button("Pricing", size="sm", variant="secondary")
            nav_legal_btn = gr.Button("Legal", size="sm", variant="secondary")
            nav_login_btn = gr.Button("Log In", size="sm", variant="secondary")
            nav_start_btn = gr.Button("Start Free", size="sm", variant="primary")
            nav_dashboard_btn = gr.Button("Dashboard", size="sm", variant="primary", visible=False)
            nav_logout_btn = gr.Button("Log Out", size="sm", variant="secondary", visible=False)

    # ===================== VIEW 0: LANDING =====================
    with gr.Column(visible=True) as view_landing:
        if tmpl:
            gr.HTML(tmpl.landing_html())
        else:
            gr.Markdown("# OmniParse AI\n\nInvoice processing in seconds, not hours.")
        with gr.Row():
            landing_cta_btn = gr.Button("Start Free \u2014 20 invoices", variant="primary", scale=1)
        if tmpl:
            gr.HTML(tmpl.footer_html())

    # ===================== VIEW 1: PRICING =====================
    with gr.Column(visible=False) as view_pricing:
        gr.Markdown("## Pricing")
        gr.HTML(tmpl.pricing_html() if tmpl else "<p>Pricing information unavailable.</p>")
        gr.Markdown(
            "Enterprise annual plan: **$4,188/year** "
            "(save 2 months compared to monthly billing)."
        )
        with gr.Accordion("Frequently asked questions about payments", open=False):
            gr.Markdown(
                "- **What payment methods do you accept?** "
                "Card payments via Stripe (Visa, Mastercard, Amex).\n"
                "- **Can I change plans at any time?** "
                "Yes, upgrade or downgrade from Dashboard \u2192 Upgrade.\n"
                "- **Do you offer refunds?** "
                "Within 14 days of your first payment, upon request."
            )
        pricing_back_btn = gr.Button("Back to Home")

    # ===================== VIEW 2: LEGAL =====================
    with gr.Column(visible=False) as view_legal:
        gr.HTML(tmpl.legal_html() if tmpl else "<h2>Legal</h2>")
        legal_back_btn = gr.Button("Back to Home")

    # ===================== VIEW 3: AUTH =====================
    with gr.Column(visible=False) as view_auth:
        gr.Markdown("## Welcome to OmniParse AI")
        if tmpl:
            gr.HTML(tmpl.auth_hint_html())
        with gr.Tabs():
            with gr.Tab("Log In"):
                login_email = gr.Textbox(label="Email", placeholder="you@company.com")
                login_password = gr.Textbox(label="Password", type="password", placeholder="Enter your password")
                login_btn = gr.Button("Log In", variant="primary")
                login_error = gr.Markdown(visible=False)

            with gr.Tab("Sign Up"):
                signup_name = gr.Textbox(label="Full Name", placeholder="Jane Smith")
                signup_email = gr.Textbox(label="Work Email", placeholder="you@company.com")
                signup_password = gr.Textbox(
                    label="Password",
                    type="password",
                    placeholder="Minimum 8 characters, mixed case, digit, and special character",
                )
                signup_terms = gr.Checkbox(
                    label="I agree to the Terms of Use and Privacy Policy"
                )
                signup_btn = gr.Button("Create Account", variant="primary")
                signup_error = gr.Markdown(visible=False)

    # ===================== VIEW 4: DASHBOARD =====================
    with gr.Column(visible=False) as view_dashboard:
        dash_welcome = gr.Markdown("## Dashboard")
        with gr.Tabs():

            # --- Upload ---
            with gr.Tab("Upload"):
                usage_md = gr.Markdown()
                if tmpl:
                    gr.HTML(tmpl.upload_dropzone_html())
                upload_files = gr.File(
                    label="Upload invoices (PDF, JPG, PNG, TIFF \u2014 max 20 MB each)",
                    file_count="multiple",
                )
                upload_btn = gr.Button("Process Invoices", variant="primary")
                upload_status = gr.Markdown()
                upload_table = gr.Dataframe(
                    headers=[
                        "Filename", "Vendor", "Invoice #",
                        "Date", "Total", "Status",
                    ],
                    label="Results",
                    interactive=False,
                )
                upload_json = gr.JSON(label="Raw output (last invoice)")

            # --- My Invoices ---
            with gr.Tab("My Invoices"):
                invoices_filter = gr.Radio(
                    ["All", "Done", "Review", "Duplicates"],
                    value="All",
                    label="Filter",
                )
                refresh_invoices_btn = gr.Button("Refresh")
                invoices_table = gr.Dataframe(
                    headers=["ID", "Vendor", "Invoice #", "Date", "Total", "Status"],
                    label="Invoices",
                    interactive=False,
                )
                with gr.Row():
                    delete_id_input = gr.Number(
                        label="Invoice ID to delete", precision=0
                    )
                    delete_invoice_btn = gr.Button("Delete", variant="stop")
                delete_status = gr.Markdown()

            # --- AI Chat ---
            with gr.Tab("AI Chat (Pro+)"):
                chat_lock_msg = gr.Markdown(visible=False)
                chatbot = gr.Chatbot(
                    label="Ask about your invoices",
                    type="messages",
                )
                chat_input = gr.Textbox(
                    label="Message",
                    placeholder="What is the total unpaid amount?",
                )
                chat_send_btn = gr.Button("Send", variant="primary")
                gr.Markdown(
                    '_Try: "List all invoices from Microsoft" / '
                    '"Which invoice has the highest tax?"_'
                )

            # --- Export ---
            with gr.Tab("Export"):
                gr.Markdown("**CSV Export** \u2014 available on all plans.")
                export_csv_btn = gr.Button("Export CSV")
                export_csv_file = gr.File(label="Download CSV")
                gr.Markdown("**JSON Export** \u2014 Basic plan and above.")
                export_json_btn = gr.Button("Export JSON")
                export_json_file = gr.File(label="Download JSON")
                gr.Markdown("**Excel Export** \u2014 Basic and above (coming soon)")
                gr.Markdown("**Google Sheets Sync** \u2014 Basic and above (coming soon)")

            # --- Upgrade ---
            with gr.Tab("Upgrade"):
                gr.HTML(tmpl.pricing_html() if tmpl else "")
                upgrade_plan_dd = gr.Dropdown(
                    ["basic", "pro", "enterprise"],
                    label="Select a plan",
                )
                upgrade_btn = gr.Button("Upgrade via Stripe", variant="primary")
                upgrade_link = gr.Markdown()
                gr.Markdown("---")
                session_id_input = gr.Textbox(
                    label="Stripe session_id (paste from URL after payment)"
                )
                check_payment_btn = gr.Button("Verify Payment")
                payment_status_md = gr.Markdown()

            # --- API ---
            with gr.Tab("API (Pro+)"):
                api_lock_msg = gr.Markdown(visible=False)
                api_key_display = gr.Markdown(visible=False)
                api_docs = gr.Markdown(visible=False)

            # --- Profile ---
            with gr.Tab("Profile"):
                profile_info = gr.Markdown()
                new_password = gr.Textbox(
                    label="New Password",
                    type="password",
                    placeholder="Enter new password",
                )
                change_pw_btn = gr.Button("Change Password")
                change_pw_status = gr.Markdown()
                gr.Markdown("### Danger Zone")
                delete_account_btn = gr.Button("Delete Account", variant="stop")
                delete_account_status = gr.Markdown()

    # =================================================================
    # Register all views
    # =================================================================
    ALL_VIEWS.extend([view_landing, view_pricing, view_legal, view_auth, view_dashboard])

    # =================================================================
    # Navigation wiring
    # =================================================================
    nav_pricing_btn.click(go_pricing, outputs=ALL_VIEWS)
    pricing_back_btn.click(go_landing, outputs=ALL_VIEWS)
    nav_legal_btn.click(go_legal, outputs=ALL_VIEWS)
    legal_back_btn.click(go_landing, outputs=ALL_VIEWS)
    nav_login_btn.click(go_auth, outputs=ALL_VIEWS)
    nav_start_btn.click(go_auth, outputs=ALL_VIEWS)
    landing_cta_btn.click(go_auth, outputs=ALL_VIEWS)

    # =================================================================
    # Auth wiring
    # =================================================================
    auth_outputs = [
        login_error, session_token, current_user,
        *ALL_VIEWS,
        nav_login_btn, nav_start_btn, nav_dashboard_btn, nav_logout_btn,
    ]

    login_btn.click(
        do_login,
        inputs=[login_email, login_password],
        outputs=auth_outputs,
    )
    signup_btn.click(
        do_signup,
        inputs=[signup_name, signup_email, signup_password, signup_terms],
        outputs=auth_outputs,
    )
    nav_logout_btn.click(
        do_logout,
        inputs=[session_token],
        outputs=[
            session_token, current_user,
            *ALL_VIEWS,
            nav_login_btn, nav_start_btn, nav_dashboard_btn, nav_logout_btn,
        ],
    )

    def go_to_dashboard_refresh(user):
        if not user:
            return show_only(3)
        return show_only(4)

    nav_dashboard_btn.click(
        go_to_dashboard_refresh, inputs=[current_user], outputs=ALL_VIEWS
    )

    # =================================================================
    # Dashboard wiring
    # =================================================================
    current_user.change(
        load_dashboard, inputs=[current_user], outputs=[dash_welcome, usage_md]
    )

    upload_btn.click(
        do_upload,
        inputs=[upload_files, current_user],
        outputs=[upload_status, upload_table, upload_json, usage_md],
    )

    refresh_invoices_btn.click(
        refresh_invoices,
        inputs=[current_user, invoices_filter],
        outputs=[invoices_table],
    )
    invoices_filter.change(
        refresh_invoices,
        inputs=[current_user, invoices_filter],
        outputs=[invoices_table],
    )

    delete_invoice_btn.click(
        do_delete_invoice,
        inputs=[current_user, delete_id_input],
        outputs=[delete_status, invoices_table],
    )

    chat_send_btn.click(
        chat_respond,
        inputs=[chat_input, chatbot, current_user],
        outputs=[chatbot, chat_input],
    )
    chat_input.submit(
        chat_respond,
        inputs=[chat_input, chatbot, current_user],
        outputs=[chatbot, chat_input],
    )

    export_csv_btn.click(
        export_csv, inputs=[current_user], outputs=[export_csv_file]
    )
    export_json_btn.click(
        export_json, inputs=[current_user], outputs=[export_json_file]
    )

    upgrade_btn.click(
        do_upgrade,
        inputs=[upgrade_plan_dd, current_user],
        outputs=[upgrade_link],
    )
    check_payment_btn.click(
        do_check_payment,
        inputs=[session_id_input, current_user],
        outputs=[payment_status_md],
    )

    current_user.change(
        load_api_tab,
        inputs=[current_user],
        outputs=[api_lock_msg, api_key_display, api_docs],
    )

    current_user.change(
        load_profile, inputs=[current_user], outputs=[profile_info]
    )
    change_pw_btn.click(
        do_change_password,
        inputs=[current_user, new_password],
        outputs=[change_pw_status],
    )
    delete_account_btn.click(
        do_delete_account,
        inputs=[current_user, session_token],
        outputs=[delete_account_status, current_user, session_token],
    )


# ---------------------------------------------------------------------------
# Launch
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    demo.queue()
    demo.launch()