"""PocketAccountant — Gradio Space entry point.

A custom 'ledger book' accountant dashboard over the deterministic engine, the
double-entry ledger, the regulation retriever and the (fine-tuned) classifier. The
tax math is always deterministic and local; the conversational agent uses whichever
LLM client is configured (local llama.cpp · Modal endpoint · deterministic router).

UI is English by default; the agent answers in English or Spanish (the language a
Mexican freelancer actually files in) via the toggle on the Ask tab.
"""

from __future__ import annotations

import os

import gradio as gr

from src import config
from src.agent import Agent, ToolContext, build_tools
from src.agent.serving import RouterClient, get_client
from src.engine import (
    isr_provisional_monthly,
    iva_monthly,
    resico_isr_monthly,
    taxable_base,
    us_annual_estimate,
)
from src.finetune import RemoteClassifier, RuleClassifier
from src.ledger import Ledger
from src.retrieval import Retriever
from src.ui.theme import CSS, MASTHEAD, theme


def make_classifier():
    """Use the Modal-served fine-tuned model when configured; else the rule classifier.

    PA_CLASSIFY_ENDPOINT, when set, points at the Modal classifier endpoint. The
    RemoteClassifier falls back to RuleClassifier on any timeout/error, so the app
    never blocks even on a cold endpoint.
    """
    endpoint = os.environ.get("PA_CLASSIFY_ENDPOINT", "").strip()
    if endpoint:
        return RemoteClassifier(endpoint, timeout=float(os.environ.get("PA_CLASSIFY_TIMEOUT", "60")))
    return RuleClassifier()


# --- shared, read-only singletons ----------------------------------------
RETRIEVER = Retriever.from_corpus_dir(config.REGULATION_DIR)
CLASSIFIER = make_classifier()

MONTHS = [("01 Jan", 1), ("02 Feb", 2), ("03 Mar", 3), ("04 Apr", 4),
          ("05 May", 5), ("06 Jun", 6), ("07 Jul", 7), ("08 Aug", 8),
          ("09 Sep", 9), ("10 Oct", 10), ("11 Nov", 11), ("12 Dec", 12)]
YEARS = [2024, 2025, 2026]


def new_ledger() -> Ledger:
    """A fresh per-session ledger, seeded with a demo month so nothing is empty."""
    lg = Ledger(":memory:")
    lg.record_income("2024-05-03", "Branding — Café Luna", 18000, iva_rate="0.16")
    lg.record_income("2024-05-17", "Website — Dental MX", 27000, iva_rate="0.16",
                     isr_retenido="2700", iva_retenido="2880")
    lg.record_expense("2024-05-05", "Adobe Creative Cloud", 1300, iva_rate="0.16")
    lg.record_expense("2024-05-12", "Office rent (share)", 4000, iva_rate="0.16")
    lg.record_expense("2024-05-22", "Non-deductible meal", 600, iva_rate="0.16",
                      deductible=False)
    return lg


# Per-session ledgers live here, keyed by a session id. We CANNOT put the Ledger
# (which holds a sqlite3.Connection) in gr.State, because Gradio deep-copies State
# values per session and a Connection can't be pickled/deep-copied. So gr.State holds
# only the id string, and the live object stays server-side.
_LEDGERS: dict = {}


def _new_session_id() -> str:
    import uuid
    sid = uuid.uuid4().hex
    _LEDGERS[sid] = new_ledger()
    return sid


def _ledger(sid):
    lg = _LEDGERS.get(sid)
    if lg is None:                      # session expired / server restarted
        lg = _LEDGERS[sid] = new_ledger()
    return lg


# --- HTML helpers ---------------------------------------------------------
def _m(v) -> str:
    try:
        f = float(v)
        cls = "pa-amount pa-neg" if f < 0 else "pa-amount"
        return f'<span class="{cls}">${f:,.2f}</span>'
    except (TypeError, ValueError):
        return f'<span class="pa-amount">{v}</span>'


def cards(items) -> str:
    cells = ""
    for k, v, src in items:
        cells += (f'<div class="pa-card"><div class="k">{k}</div>'
                  f'<div class="v">{v}</div>'
                  f'<div class="src">{src or ""}</div></div>')
    return f'<div class="pa-cards">{cells}</div>'


def breakdown_html(result) -> str:
    rows = "".join(f"<tr><td>{d}</td><td class='pa-num' style='text-align:right'>{v}</td></tr>"
                   for d, v in result.breakdown)
    src = f"<div class='src'>{result.source or ''}" + \
          (f" ({result.effective_year})" if result.effective_year else "") + "</div>"
    notes = "".join(f"<div class='pa-note'>{n}</div>" for n in result.notes)
    return (f"<table style='width:100%'>{rows}</table>{src}{notes}")


# --- callbacks ------------------------------------------------------------
def _cc(country) -> str:
    return "us" if country and "US" in str(country).upper() else "mx"


def book_expense_auto(sid, date, desc, amount, country):
    if not (date and desc and amount):
        return "<div class='pa-note'>Fill in date, description and amount.</div>"
    lg = _ledger(sid)
    if _cc(country) == "us":
        # USA: no SAT codes, no VAT — a deductible business expense (Schedule C).
        lg.record_expense(date, desc, float(amount), iva_rate="0", deductible=True)
        return (f"<div class='pa-note pa-cite'>✓ Booked <b>${float(amount):,.2f}</b> as a "
                f"deductible business expense (Schedule C). Enter business expenses only — "
                f"personal expenses aren't deductible.</div>")
    # Mexico: auto-classify to the SAT chart of accounts (+ IVA treatment).
    ctx = ToolContext(lg, classifier=CLASSIFIER)
    tools = build_tools(ctx)
    res = tools["record_expense_auto"].handler(date=date, description=desc, amount=float(amount))
    c = res["classification"]
    via = "fine-tuned model" if c.get("method") == "model" else "rule classifier"
    return (f"<div class='pa-note pa-cite'>🏷️ <b>{c['cuenta']}</b> ({c['sat_code']}) · "
            f"{'deductible' if c['deducible'] else 'non-deductible'} · IVA {c['iva_tasa']} · "
            f"{c['kind']}<br>{res['note']}<br><span class='src'>classified via {via}</span></div>")


def book_income(sid, date, desc, amount, iva_rate, country):
    if not (date and desc and amount):
        return "<div class='pa-note'>Fill in date, description and amount.</div>"
    # USA income has no VAT; ignore the (Mexico-only) IVA rate.
    rate = "0" if _cc(country) == "us" else str(iva_rate)
    _ledger(sid).record_income(date, desc, float(amount), iva_rate=rate)
    tax_note = ("no sales tax on this" if _cc(country) == "us"
                else f"IVA {float(rate) * 100:.0f}%")
    return (f"<div class='pa-note pa-cite'>✓ Income of ${float(amount):,.2f} "
            f"booked on {date} ({tax_note}).</div>")


def classify_only(desc, country):
    if not desc:
        return ""
    if _cc(country) == "us":
        return ("<div class='pa-note'>SAT account classification is Mexico-specific. "
                "For the 🇺🇸 USA, business purchases are simply <b>deductible business "
                "expenses</b> on Schedule C — just use “Book expense”.</div>")
    c = CLASSIFIER.classify(desc)
    ded = "deductible" if c.deducible else "non-deductible"
    ratio = f" (at {c.deducible_ratio*100:.1f}%)" if c.deducible_ratio < 1 else ""
    via = "fine-tuned model" if c.method == "model" else "rule classifier"
    return (f"<div class='pa-note'>🏷️ <b>{c.cuenta}</b> ({c.sat_code}) · {ded}{ratio} · "
            f"IVA {c.iva_tasa} · {c.kind}<br><span class='src'>via {via}</span></div>")


def refresh_ledger(sid, year, month, country):
    ledger = _ledger(sid)
    m = ledger.month_totals(int(year), int(month))
    if _cc(country) == "us":
        total_exp = m.deductible_expenses + m.nondeductible_expenses
        head = cards([
            ("Income", _m(m.income), f"{int(year)}-{int(month):02d}"),
            ("Deductible expenses", _m(m.deductible_expenses), "Schedule C"),
            ("Non-deductible", _m(m.nondeductible_expenses), ""),
            ("Net", _m(m.income - total_exp), "income − expenses"),
        ])
    else:
        head = cards([
            ("Income", _m(m.income), "CFDI"),
            ("Deductible expenses", _m(m.deductible_expenses), ""),
            ("IVA collected", _m(m.iva_trasladado), "trasladado"),
            ("IVA paid", _m(m.iva_acreditable), "acreditable"),
        ])
    rows = ledger.list_transactions(int(year), int(month))
    data = [[r["date"], r["description"], r["kind"], r["deductible"], r["amount"]] for r in rows]
    return head, data


def compute_taxes(sid, year, month, country):
    ledger = _ledger(sid)
    y, mo = int(year), int(month)
    if _cc(country) == "us":
        s = ledger.income_statement(y)  # annual
        est = us_annual_estimate(s.revenue, s.expenses)
        head = cards([
            ("Net profit (Schedule C)", _m(est["net_profit"].amount), f"FY {y}"),
            ("Self-employment tax", _m(est["self_employment_tax"].amount), "Sch. SE"),
            ("Federal income tax", _m(est["federal_income_tax"].amount), "after std. ded."),
            ("Total tax / year", _m(est["total_annual_tax"]), "SE + federal"),
        ])
        detail = (f"<div class='pa-note'>Taxable income = net profit − ½ SE tax "
                  f"({_m(est['half_se_deduction'])}) − standard deduction "
                  f"({_m(est['standard_deduction'])}) = <b>{_m(est['taxable_income'])}</b></div>"
                  f"<h4>Schedule C — net profit</h4>{breakdown_html(est['net_profit'])}"
                  f"<h4>Self-employment tax</h4>{breakdown_html(est['self_employment_tax'])}"
                  f"<h4>Federal income tax</h4>{breakdown_html(est['federal_income_tax'])}"
                  f"<h4>Quarterly estimate (1040-ES)</h4>{breakdown_html(est['quarterly_estimated_tax'])}"
                  "<div class='pa-note'>Simplified single-filer estimate (ignores QBI, state "
                  "tax, credits). A ballpark — confirm with a CPA.</div>")
        return head, detail

    m = ledger.month_totals(y, mo)
    resico = resico_isr_monthly(m.income)
    base = taxable_base(m.income, m.deductible_expenses)
    general = isr_provisional_monthly(base)
    iva = iva_monthly(m.iva_trasladado, m.iva_acreditable, m.iva_retenido)
    cheaper = "RESICO" if resico.amount <= general.amount else "GENERAL"
    head = cards([
        ("Income tax — RESICO", _m(resico.amount), "Art. 113-E"),
        ("Income tax — General", _m(general.amount), "Art. 96"),
        ("VAT for the month", _m(iva.amount), "Art. 5-D"),
        ("Suggested regime", f"<span style='color:#2f7d4f'>{cheaper}</span>", "lower ISR"),
    ])
    detail = (f"<h4>Income tax — RESICO regime</h4>{breakdown_html(resico)}"
              f"<h4>Income tax — General regime</h4>{breakdown_html(general)}"
              f"<h4>VAT (IVA)</h4>{breakdown_html(iva)}"
              "<div class='pa-note'>Regime choice is annual and has eligibility rules — "
              "confirm with your accountant.</div>")
    return head, detail


def show_statements(sid, year, month):
    ledger = _ledger(sid)
    y, mo = int(year), int(month)
    inc = ledger.income_statement(y, mo)
    last = f"{y:04d}-{mo:02d}-28"
    bs = ledger.balance_sheet(last)
    pl = cards([
        ("Revenue", _m(inc.revenue), inc.period),
        ("Expenses", _m(inc.expenses), ""),
        ("Net profit", _m(inc.net_profit), ""),
    ])
    balance = cards([
        ("Assets", _m(bs.assets), bs.as_of),
        ("Liabilities", _m(bs.liabilities), ""),
        ("Equity", _m(bs.equity), "assets − liabilities"),
    ])
    return f"<h4>Income Statement (P&amp;L)</h4>{pl}<h4>Balance Sheet</h4>{balance}"


def ask(question, sid, year, month, lang, country):
    if not question:
        yield "Type a question.", ""
        return
    # Immediate feedback — the reasoning model takes ~20–60s (longer on the first
    # question while it wakes up), so show progress instead of a frozen button.
    yield ("⏳ **The accountant is thinking…** reading the books, choosing the right "
           "calculation, and checking the regulation. The reasoning model can take "
           "20–60 seconds (longer on the first question while it warms up)."), ""

    ctx = ToolContext(_ledger(sid), retriever=RETRIEVER, classifier=CLASSIFIER, country=_cc(country).upper())
    tools = build_tools(ctx)
    code = "es" if lang and lang.lower().startswith(("es", "espa")) else "en"
    tagged = f"[{_cc(country)}][{code}][{int(year)}-{int(month):02d}] {question}"
    note = ""
    try:
        trace = Agent(get_client(), tools).run(tagged)
    except Exception:
        # Model endpoint cold/unavailable → deterministic fallback so the user
        # always gets a grounded answer (numbers still come from the engine).
        trace = Agent(RouterClient(), tools).run(tagged)
        note = "\n\n_(the reasoning model was unavailable — answered with the built-in fallback)_"
    used = " · ".join(s.tool for s in trace.steps) or "—"
    yield (trace.final_answer or "") + note, f"tools: {used}"


# --- UI -------------------------------------------------------------------
def build_demo() -> gr.Blocks:
    with gr.Blocks(theme=theme, css=CSS, title="PocketAccountant") as demo:
        gr.HTML(MASTHEAD)
        # Holds only a session-id string (never the un-pickleable Ledger object).
        session = gr.State(value=_new_session_id)
        with gr.Row():
            # Use (label, value) tuples with PLAIN ascii values ("MX"/"US"). Flag-emoji
            # values were breaking Gradio's option matching, so the selection never
            # actually switched to USA. The emoji stays in the visible label only.
            country = gr.Dropdown([("🇲🇽 Mexico", "MX"), ("🇺🇸 USA", "US")], value="MX",
                                  label="Country / tax system", filterable=False,
                                  info="Mexico: RESICO · IVA · SAT. USA: Schedule C · SE tax · federal.")

        with gr.Tabs():
            # ---- Capture ----
            with gr.Tab("🧾 Capture"):
                gr.Markdown("Record income and expenses (set your **country above**). "
                            "🇲🇽 Mexico: expenses are auto-classified to a SAT account and IVA "
                            "is tracked. 🇺🇸 USA: booked as deductible business expenses (no VAT).")
                with gr.Row():
                    with gr.Column():
                        gr.Markdown("#### Expense")
                        e_date = gr.Textbox(label="Date (YYYY-MM-DD)", value="2024-05-25")
                        e_desc = gr.Textbox(label="Description", placeholder="Adobe subscription…")
                        e_amt = gr.Number(label="Amount (before tax)", value=1200)
                        with gr.Row():
                            e_classify = gr.Button("Classify (MX)", size="sm")
                            e_book = gr.Button("Book expense", variant="primary", size="sm")
                    with gr.Column():
                        gr.Markdown("#### Income")
                        i_date = gr.Textbox(label="Date (YYYY-MM-DD)", value="2024-05-28")
                        i_desc = gr.Textbox(label="Description", placeholder="Logo design…")
                        i_amt = gr.Number(label="Amount (before tax)", value=10000)
                        i_iva = gr.Dropdown([("16%", "0.16"), ("0% (export)", "0.00")],
                                            value="0.16", label="IVA rate (Mexico only)")
                        i_book = gr.Button("Book income", variant="primary", size="sm")
                capture_out = gr.HTML()

            # ---- Ledger ----
            with gr.Tab("📒 Ledger"):
                with gr.Row():
                    l_year = gr.Dropdown(YEARS, value=2024, label="Year")
                    l_month = gr.Dropdown(MONTHS, value=5, label="Month")
                    l_refresh = gr.Button("Refresh", variant="primary", size="sm")
                ledger_cards = gr.HTML()
                ledger_table = gr.Dataframe(
                    headers=["Date", "Description", "Type", "Deductible", "Amount"],
                    datatype=["str", "str", "str", "str", "str"],
                    interactive=False, wrap=True)

            # ---- Taxes ----
            with gr.Tab("💸 Taxes"):
                with gr.Row():
                    t_year = gr.Dropdown(YEARS, value=2024, label="Year")
                    t_month = gr.Dropdown(MONTHS, value=5, label="Month")
                    t_go = gr.Button("Compute", variant="primary", size="sm")
                tax_cards = gr.HTML()
                tax_detail = gr.HTML()

            # ---- Statements ----
            with gr.Tab("📊 Statements"):
                with gr.Row():
                    s_year = gr.Dropdown(YEARS, value=2024, label="Year")
                    s_month = gr.Dropdown(MONTHS, value=5, label="Month")
                    s_go = gr.Button("Generate", variant="primary", size="sm")
                statements_out = gr.HTML()

            # ---- Ask ----
            with gr.Tab("💬 Ask your accountant"):
                gr.Markdown("Ask about your taxes — **pick your country above** (🇲🇽 or 🇺🇸). "
                            "Rules are **cited** from the regulation; the numbers come from the "
                            "deterministic engine.")
                with gr.Row():
                    a_year = gr.Dropdown(YEARS, value=2024, label="Year", scale=1)
                    a_month = gr.Dropdown(MONTHS, value=5, label="Month", scale=1)
                    a_lang = gr.Dropdown([("English", "en"), ("Español", "es")], value="en",
                                         label="Answer language", filterable=False, scale=1)
                a_q = gr.Textbox(label="Your question",
                                 placeholder="Which regime suits me? Can I deduct my laptop?")
                a_go = gr.Button("Ask", variant="primary")
                a_answer = gr.Markdown()
                a_tools = gr.Markdown()
                gr.Examples(
                    ["Which regime suits me this month?", "How much VAT do I owe?",
                     "Can I deduct my laptop?", "How much tax do I owe?",
                     "What is the QBI deduction?", "When do I file my taxes?"],
                    inputs=a_q)

        gr.HTML("<div class='pa-foot'>Educational assistant · not a substitute for a "
                "licensed accountant / CPA · 2024 tax tables (verify against DOF/IRS)</div>")

        # --- wiring ---
        e_classify.click(classify_only, [e_desc, country], [capture_out])
        e_book.click(book_expense_auto, [session, e_date, e_desc, e_amt, country], [capture_out])
        i_book.click(book_income, [session, i_date, i_desc, i_amt, i_iva, country], [capture_out])
        l_refresh.click(refresh_ledger, [session, l_year, l_month, country],
                        [ledger_cards, ledger_table])
        t_go.click(compute_taxes, [session, t_year, t_month, country], [tax_cards, tax_detail])
        s_go.click(show_statements, [session, s_year, s_month], [statements_out])
        a_go.click(ask, [a_q, session, a_year, a_month, a_lang, country], [a_answer, a_tools])

        # initial render
        demo.load(refresh_ledger, [session, l_year, l_month, country], [ledger_cards, ledger_table])
    return demo


demo = build_demo()

if __name__ == "__main__":
    demo.launch()