#!/usr/bin/env python3
"""Statement loader: parses bank/card statements and adds new transactions to Bills_Tracker.xlsx.

Sources understood: Chase Visa PDF, Ally Bank PDF, CVFCU PDF, Capital One CSV download.
Everything runs locally. The workbook is backed up before every write.

CLI:  .venv/bin/python statement_loader.py preview     # what would be added (writes nothing)
      .venv/bin/python statement_loader.py import      # back up, then add new rows
"""
import re, csv, sys, shutil, zipfile, datetime as dt, calendar, collections
from dataclasses import dataclass, field
from pathlib import Path
import pdfplumber
from seed_rules import seed_match

import os
HERE = Path(__file__).parent                      # code lives here
DATA = Path(os.environ.get("DATA_DIR", HERE))     # statements, workbook and backups live here (a mounted volume in Docker)
WORKBOOK = DATA / "Bills_Tracker.xlsx"
BACKUPS = DATA / "backups"

# ---------------------------------------------------------------- accounts
# name, text on page 1 that identifies it, day-of-month the statement closes (None = month end)
ACCOUNTS = [
    ("Chase Visa",     ["chase.com", "Prime Visa"],                  15),
    ("Ally Checking",  ["COMBINED CUSTOMER STATEMENT", "ally.com"],  10),
    ("CVFCU Checking", ["cvfcu.com"],                                None),
]
CSV_ACCOUNTS = {}   # optional: last 4 of card no. in the CSV -> account name (set in loader_rules.json)
CSV_COLUMNS = ("Transaction Date", "Debit", "Credit", "Description")

# Personal matching rules live in DATA/loader_rules.json (git-ignored, never baked into the image), e.g.
#   {"own_names": ["JANE DOE"], "own_bank_hints": ["MY CREDIT UNION"], "csv_accounts": {"1234": "Capital One Card"}}
# own_names: your name as it appears in transfer descriptions, so transfers between your own accounts are not counted as spending.
import json
def _rules():
    try: return json.loads((DATA / "loader_rules.json").read_text())
    except (OSError, ValueError): return {}
def own(d):  return any(n.upper() in d for n in _rules().get("own_names", []))
def bank_hint(d): return any(n.upper() in d for n in _rules().get("own_bank_hints", []))

def money(s):
    s = s.replace(",", "").replace("$", "")
    return float(s)
def pdate(s):
    for f in ("%m/%d/%Y", "%m/%d/%y", "%Y-%m-%d"):
        try: return dt.datetime.strptime(s, f).date()
        except ValueError: pass

@dataclass
class Txn:
    account: str
    date: dt.date
    desc: str            # statement description as shown
    detail: str          # continuation lines (merchant detail on Ally etc.)
    amount: float        # positive = money out / charge, negative = money in / refund
    source: str = ""     # file name
    extra: str = ""      # e.g. mortgage principal/interest note
    category_hint: str = ""

@dataclass
class Statement:
    file: str
    account: str
    start: dt.date | None
    end: dt.date
    txns: list = field(default_factory=list)
    stated_out: float | None = None   # total "money out" the statement itself reports (for checking)
    kind: str = "pdf"
    snaps: list = field(default_factory=list)   # savings / loan balances (kept on the separate Accounts tab, never in spending)

# ---------------------------------------------------------------- PDF parsers
def _pages(path):
    with pdfplumber.open(path) as pdf:
        return [(pg.extract_text() or "") for pg in pdf.pages]

def identify_pdf(pages):
    first = pages[0].lower()
    for name, keys, _ in ACCOUNTS:
        if any(k.lower() in first for k in keys): return name
    return None

def parse_chase(path, pages):
    full = "\n".join(pages)
    m = re.search(r"Opening/Closing Date\s+(\S+) - (\S+)", full)
    start, end = pdate(m.group(1)), pdate(m.group(2))
    out = {}
    for k in ("Purchases", "Fees Charged", "Interest Charged"):
        mm = re.search(rf"^{k}\s+\+?\$([\d,]+\.\d\d)", full, re.M)
        out[k] = money(mm.group(1)) if mm else 0.0
    st = Statement(Path(path).name, "Chase Visa", start, end, stated_out=sum(out.values()))
    in_rewards = False
    for line in full.splitlines():
        if re.search(r"\$ Amount Rewards", line): in_rewards = True; continue
        if re.match(r"ACCOUNT ACTIVITY|PAYMENTS AND OTHER|PURCHASE$|FEES CHARGED|INTEREST CHARGED", line): in_rewards = False
        if in_rewards: continue
        mm = re.match(r"^(\d\d)/(\d\d) (.+?) (-?(?:[\d,]+)?\.\d\d)$", line)
        if not mm: continue
        mo, dy = int(mm.group(1)), int(mm.group(2))
        yr = end.year if mo <= end.month or (mo, dy) <= (end.month, end.day) else end.year - 1
        if dt.date(yr, mo, dy) > end + dt.timedelta(days=31): yr -= 1
        st.txns.append(Txn("Chase Visa", dt.date(yr, mo, dy), mm.group(3).strip(), "", money(mm.group(4)), st.file))
    return st

ALLY_TX = re.compile(r"^(\d\d/\d\d/\d{4}) (.+?) \$([\d,]+\.\d\d) -\$([\d,]+\.\d\d) \$([\d,]+\.\d\d)$")
def parse_ally(path, pages):
    full = "\n".join(pages)
    m = re.search(r"Statement Date.{0,120}?(\d\d/\d\d/\d{4})", pages[0], re.S)
    end = pdate(m.group(1))
    mb = re.search(r"Beginning Balance, as of (\d\d/\d\d/\d{4})", full)
    start = pdate(mb.group(1)) + dt.timedelta(days=1) if mb else None
    md = re.search(r"Withdrawals and Other Debits\s+-?\$([\d,]+\.\d\d)", full)
    st = Statement(Path(path).name, "Ally Checking", start, end, stated_out=money(md.group(1)) if md else None)
    acct, cur = "Ally Checking", None
    for line in full.splitlines():
        if re.match(r"^(Spending|Savings) Account\b", line) and "xxxx" not in line.lower() and "www" in line:
            acct = "Ally Checking" if line.startswith("Spending") else "Ally Savings"
        mm = ALLY_TX.match(line)
        if mm:
            if mm.group(2) in ("Beginning Balance", "Ending Balance"): cur = None; continue
            cr, db = money(mm.group(3)), money(mm.group(4))
            cur = Txn(acct, pdate(mm.group(1)), re.sub(r"^(\w) (\w)", r"\1\2", mm.group(2)).strip(), "", round(db - cr, 2), st.file)
            st.txns.append(cur); continue
        if re.match(r"^\d\d/\d\d/\d{4} (Beginning|Ending) Balance", line): cur = None; continue
        if cur and not re.match(r"^(\d{6}/|Ally Bank Member|Page \d|Statement Date|COMBINED|\d{6}-\d\d)", line):
            cur.detail = (cur.detail + " | " + line.strip()).strip(" |")
        if re.match(r"^\d{6}/\d+//", line): cur = None
    st.txns = [t for t in st.txns if t.account == "Ally Checking"]   # savings has no spending
    heads = list(re.finditer(r"Product: (.+?) Account Ownership", full))
    for i, h in enumerate(heads):
        prod = h.group(1).strip()
        if prod.lower().startswith("spending"): continue
        sec = full[h.end(): heads[i + 1].start() if i + 1 < len(heads) else h.end() + 3000]
        g = lambda rx: (lambda m: money(m.group(1)) if m else None)(re.search(rx, sec))
        bal = g(r"Ending Balance, as of \S+ \$([\d,]+\.\d\d)")
        if bal is None: continue
        apy = re.search(r"Annual Percentage Yield Earned ([\d.]+)%", sec)
        st.snaps.append(dict(account="Ally " + re.sub(r" Account$", "", prod), group="Savings", date=end, balance=bal,
                             interest=g(r"Interest Paid This Period \$([\d,]+\.\d\d)"), ytd=g(r"Interest Paid Year to Date \$([\d,]+\.\d\d)"),
                             deposits=g(r"Deposits and Other Credits \$([\d,]+\.\d\d)"), withdrawals=g(r"Withdrawals and Other Debits -?\$([\d,]+\.\d\d)"),
                             rate=float(apy.group(1)) if apy else None))
    return st

CV_TX = re.compile(r"^(\d\d/\d\d/\d{4}) (.+?) \$([\d,]+\.\d\d) \$([\d,]+\.\d\d)$")
def parse_cvfcu(path, pages):
    full = "\n".join(pages)
    end = pdate(re.search(r"Statement Ending (\d\d/\d\d/\d{4})", pages[0]).group(1))
    st = Statement(Path(path).name, "CVFCU Checking", dt.date(end.year, end.month, 1), end)
    block, prev, cur = None, None, None
    pr = fin = None
    for line in full.splitlines():
        h = re.match(r"^([A-Z0-9][A-Z0-9 &/\.\-]+?) - XXXXXX\d{4}$", line)
        if h:
            if h.group(1) != block: prev = None
            block, cur = h.group(1), None; continue
        if re.match(r"^30 YR MORTGAGE|^Statement Year to Date", line): block = None
        if block is None:
            mp = re.match(r"^Principal Amount -?\$([\d,.]+)", line); mf = re.match(r"^Finance Charge \$([\d,.]+)", line)
            if mp: pr = money(mp.group(1))
            if mf: fin = money(mf.group(1))
            continue
        if block != "PREFERRED CHECKING": continue
        md = re.match(r"^([\d,]+) Debit\(s\) This Period -\$([\d,]+\.\d\d)", line)
        if md: st.stated_out = money(md.group(2)); continue
        mb = re.match(r"^\d\d/\d\d/\d{4} Beginning Balance \$([\d,]+\.\d\d)$", line)
        if mb: prev = money(mb.group(1)); continue
        mm = CV_TX.match(line)
        if mm and prev is not None:
            amt, bal = money(mm.group(3)), money(mm.group(4))
            sign = 1 if abs(prev - amt - bal) < 0.005 else -1 if abs(prev + amt - bal) < 0.005 else 1
            cur = Txn("CVFCU Checking", pdate(mm.group(1)), mm.group(2).strip(), "", round(sign * amt, 2), st.file)
            st.txns.append(cur); prev = bal; continue
        if re.match(r"^\d\d/\d\d/\d{4} Ending Balance", line): cur = None
    for t in st.txns:
        if "SCHEDULED TRANSFER TO LOAN" in t.desc.upper() and pr and fin:
            t.extra = f"Principal ${pr:,.2f} / Interest ${fin:,.2f}"
    st.snaps = cvfcu_snapshots(full, end)
    return st

def cvfcu_snapshots(full, end):
    """Savings/share and loan accounts from a CVFCU statement (checking is handled as transactions)."""
    bal, kind, loan = {}, {}, False
    for line in full.splitlines():
        if line.startswith("Loan Account Account Number"): loan = True; continue
        if line.startswith("Account Type Account Number"): loan = False; continue
        m = re.match(r"^([A-Z0-9][A-Z0-9 &]+?) XXXXXX\d{4} \$(-?[\d,]+\.\d\d)$", line)
        if m and "CHECKING" not in m.group(1):
            bal[m.group(1)] = money(m.group(2)); kind[m.group(1)] = "Loan" if loan else "Savings"
    info, blk = {n: {} for n in bal}, None
    pats = dict(interest=r"Interest Paid This Period \$([\d,]+\.\d\d)", ytd=r"Interest Paid (?:Year-to-Date|YTD) \$([\d,]+\.\d\d)",
                rate=r"Interest Rate ([\d.]+)%", payment=r"Payment Amount \$([\d,]+\.\d\d)", due=r"Payment Due Date (\d\d/\d\d/\d{4})",
                principal=r"Principal Amount -?\$([\d,]+\.\d\d)")
    for line in full.splitlines():
        h = re.match(r"^([A-Z0-9][A-Z0-9 &/\.]+?) - XXXXXX\d{4}$", line)
        if h: blk = h.group(1); continue
        if blk in info:
            for k, rx in pats.items():
                m = re.search(rx, line)
                if m and k not in info[blk]: info[blk][k] = pdate(m.group(1)) if k == "due" else float(m.group(1).replace(",", ""))
    out = []
    for n, b in bal.items():
        i = info[n]; note = None
        if kind[n] == "Loan" and i.get("payment") and i.get("principal") is not None:
            note = f"Payment ${i['payment']:,.2f}: principal ${i['principal']:,.2f} + interest ${i.get('interest', 0):,.2f}"
        out.append(dict(account="CVFCU " + n.title(), group=kind[n], date=end, balance=b, interest=i.get("interest"), ytd=i.get("ytd"),
                        principal=i.get("principal"), rate=i.get("rate"), payment=i.get("payment"), due=i.get("due"), note=note))
    return out

# Online "Account Details" page saved as PDF (recent activity, no statement dates). Account is named by its product line.
ACTIVITY_ACCOUNTS = {"PREFERRED CHECKING": "CVFCU Checking"}
ACT_TX = re.compile(r"^([A-Z][a-z]{2}) (\d\d) (.+?) ([+-])\$([\d,]+\.\d\d) \$(-?[\d,]+\.\d\d)$")
MONTHS = {m: i for i, m in enumerate(["Jan", "Feb", "Mar", "Apr", "May", "Jun", "Jul", "Aug", "Sep", "Oct", "Nov", "Dec"], 1)}

def identify_activity(pages):
    first = pages[0]
    if "Account Details" in first and "Past Transactions" in first:
        for k, name in ACTIVITY_ACCOUNTS.items():
            if k in first.upper(): return name
    return None

def parse_activity(path, pages, account):
    """Recent-activity page: newest first, signed amounts and a running balance. Every row is checked against the balance chain."""
    full = "\n".join(pages)
    rng = re.search(r"(\d\d/\d\d/\d{4}) (\d\d/\d\d/\d{4})", full)
    lo, hi = (pdate(rng.group(1)), pdate(rng.group(2))) if rng else (None, None)
    cur = re.search(r"Current Balance \$(-?[\d,]+\.\d\d)", full)
    rows = []
    for line in full.split("Past Transactions", 1)[-1].splitlines():
        m = ACT_TX.match(line.strip())
        if not m: continue
        mon, day = MONTHS[m.group(1)], int(m.group(2))
        yr = (hi or dt.date.today()).year
        d = dt.date(yr, mon, day)
        if hi and d > hi: d = dt.date(yr - 1, mon, day)
        amt = money(m.group(5)) * (1 if m.group(4) == "+" else -1)     # signed: + money in
        rows.append((d, m.group(3).strip(), amt, money(m.group(6))))
    if not rows: raise ValueError("no transactions found")
    if cur and abs(rows[0][3] - money(cur.group(1))) > 0.005: raise ValueError("latest balance doesn't match Current Balance")
    for (d, _, a, bal), (_, _, _, nxt) in zip(rows, rows[1:]):
        if abs(nxt - (bal - a)) > 0.005: raise ValueError(f"balance doesn't add up near {d:%b %-d}: a row is missing or misread")
    st = Statement(Path(path).name, account, min(r[0] for r in rows), max(r[0] for r in rows), kind="activity")
    for d, desc, a, _ in rows:
        desc = re.sub(r"^ACH TRANS : ", "EFT ", desc)
        st.txns.append(Txn(account, d, desc, "", round(-a, 2), st.file))      # app convention: + = money out
    st.stated_out = round(sum(-a for _, _, a, _ in rows if a < 0), 2)
    return st

def parse_capone_csv(path):
    rows = list(csv.DictReader(open(path, newline="", encoding="utf-8-sig")))
    rows = [r for r in rows if r.get("Transaction Date") and pdate(r["Transaction Date"])]
    cards = {r.get("Card No.", "") for r in rows}
    mapped = {**CSV_ACCOUNTS, **_rules().get("csv_accounts", {})}
    unknown = sorted(c for c in cards if c not in mapped)
    acct = next((mapped[c] for c in cards if c in mapped), None) or ("Capital One Card" if len(cards) == 1 else f"Capital One Card ****{unknown[0]}")
    ds = [pdate(r["Transaction Date"]) for r in rows]
    st = Statement(Path(path).name, acct, min(ds), max(ds), kind="csv")
    st.stated_out = round(sum(money(r["Debit"]) for r in rows if r.get("Debit")), 2)
    for r in rows:
        amt = money(r["Debit"]) if r.get("Debit") else -money(r["Credit"])
        st.txns.append(Txn(acct, pdate(r["Transaction Date"]), r["Description"].strip(), "", round(amt, 2), st.file, category_hint=r.get("Category", "")))
    return st

_CACHE = {}
_DISK = DATA / ".statement_cache.pkl"      # parsed statements survive restarts (re-reading ~30 PDFs takes several seconds)
_CACHE_VERSION = 2
_loaded = False

def _load_cache():
    global _loaded
    if _loaded: return
    _loaded = True
    try:
        import pickle
        v, data = pickle.loads(_DISK.read_bytes())
        if v == _CACHE_VERSION: _CACHE.update(data)
    except Exception: pass

def _save_cache():
    try:
        import pickle
        tmp = _DISK.with_suffix(".tmp"); tmp.write_bytes(pickle.dumps((_CACHE_VERSION, _CACHE))); tmp.replace(_DISK)
    except Exception: pass
def load_statements(folder=DATA):
    """Parse every statement file in folder. Returns (statements, skipped, duplicates)."""
    sts, skipped, dups, seen = [], [], [], {}
    _load_cache(); parsed_new = False
    for p in sorted(Path(folder).glob("*")):
        if p.suffix.lower() not in (".pdf", ".csv"): continue
        ck = (p.name, p.stat().st_mtime_ns, p.stat().st_size)
        if ck in _CACHE:
            st, err = _CACHE[ck]
            if err: skipped.append((p.name, err))
            elif st:
                key = (st.account, st.end, st.kind)
                if key in seen and st.kind == "pdf": dups.append((st.file, seen[key]))
                else: seen[key] = st.file; sts.append(st)
            continue
        st, err = None, None
        try:
            if p.suffix.lower() == ".pdf":
                pages = _pages(p); name = identify_pdf(pages); act = None if name else identify_activity(pages)
                if act: st = parse_activity(p, pages, act)
                elif not name: err = "not recognized (new account? see ACCOUNTS in statement_loader.py)"
                else: st = {"Chase Visa": parse_chase, "Ally Checking": parse_ally, "CVFCU Checking": parse_cvfcu}[name](p, pages)
            elif p.suffix.lower() == ".csv":
                head = open(p, encoding="utf-8-sig").readline()
                if not all(c in head for c in CSV_COLUMNS): _CACHE[ck] = (None, None); continue
                st = parse_capone_csv(p)
        except Exception as e:
            err = f"could not read: {e}"
        _CACHE[ck] = (st, err); parsed_new = True
        if err: skipped.append((p.name, err)); continue
        if st is None: continue
        key = (st.account, st.end, st.kind)
        if key in seen and st.kind == "pdf": dups.append((st.file, seen[key])); continue
        seen[key] = st.file; sts.append(st)
    if parsed_new:
        live = {(p.name, p.stat().st_mtime_ns, p.stat().st_size) for p in Path(folder).glob("*")}
        for k in [k for k in _CACHE if k not in live]: del _CACHE[k]      # forget deleted / replaced files
        _save_cache()
    return sts, skipped, dups

# ---------------------------------------------------------------- classification
def norm(desc):
    toks = []
    for t in re.sub(r"[*,;#\-]", " ", desc.upper()).split():
        if re.search(r"\d|#", t): break
        toks.append(t)
        if len(toks) == 3: break
    return " ".join(toks)

def _is_chase_payment(t): return t.account == "Chase Visa" and re.match(r"(AUTOMATIC )?PAYMENT", t.desc.upper()) and t.amount < 0

def exclusion(t):
    """Return (reason, direction) if this is NOT spending (card payment, transfer, income, interest), else None."""
    d = (t.desc + " " + t.detail).upper(); out = "Money out" if t.amount > 0 else "Money in"
    CARD = "Card payment (spending counted when purchased)"
    if _is_chase_payment(t): return CARD, out
    if t.account == "Capital One Card" and t.category_hint.startswith("Payment") and t.amount < 0: return CARD, out
    if t.account == "Ally Checking":
        if t.desc.upper().startswith("INTEREST PAID"): return "Interest earned", out
        if "CASH APP*" in d and own(d): return "Cash App transfer to your own Cash App (confirm destination)", out
        if "CASH APP" in d and t.amount < 0: return "Deposit (Cash App / Amazon) – not spending", out
        if "AMAZON.COM SERVI" in d and t.amount < 0: return "Deposit (Cash App / Amazon) – not spending", out
        if "SAVINGS" in d and "TRANSFER" in d: return "Transfer to savings account", out
        if ("TRANSFER TO" in d and own(d)) or bank_hint(d): return "Transfer to CVFCU checking", out
        if ("TRANSFER FROM" in d and own(d)) or ("ALLY BANK TRANSFER" in d and t.amount < 0): return "Transfer from CVFCU checking", out
    if t.account == "CVFCU Checking":
        if "PAYROLL" in d: return "Income (payroll deposit)", out
        if "CHASE CREDIT CRD" in d or "DISCOVER" in d: return CARD, out
        if "ALLY BANK" in d: return "Transfer to Ally account", out
        if "P2P" in d and own(d): return "Transfer from Ally account", out
        if d.startswith("DIVIDEND EARNED"): return "Interest earned", out
    return None

SEED_RULES = [   # (text to find in desc+detail, payee, what for)  for sources whose log description is generic
    ("T-MOBILE", "T-Mobile", "Cell phone service"),
    ("SCHEDULED TRANSFER TO LOAN", "CVFCU Mortgage", "Monthly mortgage payment (principal + interest)"),
]

@dataclass
class Row:
    txn: Txn
    kind: str                # "log" or "excluded"
    payee: str = ""
    what: str = ""
    reason: str = ""
    direction: str = ""
    note: str = ""
    is_new_payee: bool = False
    needs_review: bool = False
    category: str = ""
    type: str = ""

def learn_rules(wb):
    """payee/what-for per normalized description, learned from rows already in the Bill Log."""
    c = collections.defaultdict(collections.Counter)
    for r in wb["Bill Log"].iter_rows(min_row=6, values_only=True):
        if r[0] and r[7] and r[1]:
            c[norm(r[7])][(r[1], r[4])] += 1
    return {k: v.most_common(1)[0][0] for k, v in c.items() if k}

def classify(t, rules, known_payees, cats=None):
    ex = exclusion(t)
    if ex: return Row(t, "excluded", reason=ex[0], direction=ex[1])
    d = (t.desc + " " + t.detail).upper()
    payee = what = None
    for needle, p, w in SEED_RULES:
        if needle in d: payee, what = p, w; break
    if not payee and t.account == "Ally Checking" and "ATM" in d:
        payee, what = "ATM Cash", "Cash withdrawal (net of ATM-fee reimbursement)"
    if not payee:
        k = norm(t.desc)
        hit = rules.get(k)
        if not hit and k:   # fall back to the first two words
            two = " ".join(k.split()[:2])
            cands = {v for kk, v in rules.items() if " ".join(kk.split()[:2]) == two}
            if len(cands) == 1: hit = next(iter(cands))
        if hit: payee, what = hit
    sm = seed_match(d)                       # built-in generic merchant rules
    row = Row(t, "log")
    if payee: row.payee, row.what = payee, what
    else:
        row.payee = (sm[0] if sm and sm[0] else " ".join(w.capitalize() for w in re.sub(r"^EFT ", "", norm(t.desc) or t.desc.upper()).split())[:40]) or "Unknown"
        row.what = sm[3] if sm else "NEW – please categorize"
        row.is_new_payee = True
        if not sm: row.needs_review = True; row.note = "NEW – review payee/category"
    if payee and payee not in known_payees: row.is_new_payee = True
    if row.is_new_payee and row.payee in known_payees: row.is_new_payee = False
    if row.is_new_payee:
        row.category, row.type = ("Other / Unsorted", "Spending")
        if sm and (cats is None or sm[1] in cats): row.category, row.type = sm[1], sm[2]
    if t.amount < 0: row.note = (row.note + "; " if row.note else "") + "Refund / credit"
    if t.extra: row.note = (row.note + "; " if row.note else "") + t.extra
    return row

# ---------------------------------------------------------------- workbook access
def _categories(wb):
    return [r[0] for r in wb["Lists"].iter_rows(min_row=4, max_row=40, min_col=5, max_col=5, values_only=True)
            if r[0] and not str(r[0]).startswith("Slots")]

def existing_keys(wb):
    """Multiset of (account, date, |amount|) already in Bill Log + Excluded Transfers."""
    keys = collections.Counter()
    for r in wb["Bill Log"].iter_rows(min_row=6, values_only=True):
        if r[0]: keys[(r[6], r[0].date(), round(abs(r[5]), 2))] += 1
    for r in wb["Excluded Transfers"].iter_rows(min_row=5, values_only=True):
        if r[0] and isinstance(r[0], dt.datetime): keys[(r[1], r[0].date(), round(abs(r[4]), 2))] += 1
    return keys

def plan(wb, statements):
    """Work out which transactions are new. ATM withdrawal + same-day fee line are merged like the existing log does."""
    rules = learn_rules(wb)
    lists = wb["Lists"]
    known = {r[0] for r in lists.iter_rows(min_row=4, max_row=153, max_col=1, values_only=True) if r[0]}
    cats = _categories(wb)
    have = existing_keys(wb)
    rows, taken = [], collections.Counter()
    for st in sorted(statements, key=lambda x: x.kind == "activity"):      # real statements win over recent-activity pages
        local = collections.Counter()
        for t in st.txns:
            k = (t.account, t.date, round(abs(t.amount), 2)); local[k] += 1
            if local[k] > taken[k]: rows.append(classify(t, rules, known, cats))      # same charge in two overlapping files counts once
        for k, n in local.items(): taken[k] = max(taken[k], n)
    new, already = [], 0
    for r in rows:
        t = r.txn; k = (t.account, t.date, round(abs(t.amount), 2))
        if have[k] > 0: have[k] -= 1; already += 1; continue
        new.append(r)
    return new, already

def check_statements(statements):
    """Compare what we parsed to the totals each statement prints. Returns [(file, parsed, stated, ok)]."""
    res = []
    for st in statements:
        if st.stated_out is None: continue
        parsed = round(sum(t.amount for t in st.txns if t.amount > 0), 2)
        res.append((st.file, parsed, st.stated_out, abs(parsed - st.stated_out) < 0.005))
    return res

# ---------------------------------------------------------------- writing
def _chart_sig(xml):
    refs = re.findall(rb"<(?:c:)?f>(.*?)</(?:c:)?f>", xml)
    kinds = tuple(re.findall(rb"<(?:c:)?(bar|line|pie|doughnut|area)Chart>", xml))
    return kinds + tuple(sorted(re.sub(rb"&apos;|'|\$| ", b"", r) for r in refs))

def _restore_charts(original, saved):
    """openpyxl rewrites charts and drops some styling; put the original chart XML back (matched by plotted ranges)."""
    with zipfile.ZipFile(original) as zo, zipfile.ZipFile(saved) as zs:
        orig = {_chart_sig(zo.read(n)): zo.read(n) for n in zo.namelist() if re.match(r"xl/charts/chart\d+\.xml$", n)}
        out = {}
        for n in zs.namelist():
            data = zs.read(n)
            if re.match(r"xl/charts/chart\d+\.xml$", n):
                data = orig.get(_chart_sig(data), data)
            out[n] = data
        infos = {i.filename: i for i in zs.infolist()}
    tmp = Path(str(saved) + ".tmp")
    with zipfile.ZipFile(tmp, "w", zipfile.ZIP_DEFLATED) as z:
        for n, data in out.items(): z.writestr(n, data)
    tmp.replace(saved)

def _copy_style(src, dst):
    from copy import copy
    dst.font, dst.fill, dst.border, dst.alignment, dst.number_format = copy(src.font), copy(src.fill), copy(src.border), copy(src.alignment), src.number_format

ACCOUNT_COLS = ["Statement date", "Account", "Type", "Balance", "Interest this period", "Interest year to date", "Deposits", "Withdrawals",
                "Principal this period", "Rate %", "Payment", "Next due", "Note"]
ACCOUNT_KEYS = ["date", "account", "group", "balance", "interest", "ytd", "deposits", "withdrawals", "principal", "rate", "payment", "due", "note"]

def account_rows(wb):
    """Existing snapshots on the Accounts tab as dicts (empty if the tab doesn't exist yet)."""
    if "Accounts" not in wb.sheetnames: return []
    out = []
    for r in wb["Accounts"].iter_rows(min_row=2, max_col=len(ACCOUNT_KEYS), values_only=True):
        if r[0] is None or r[1] is None: continue
        d = dict(zip(ACCOUNT_KEYS, r))
        for k in ("date", "due"):
            if isinstance(d[k], dt.datetime): d[k] = d[k].date()
        out.append(d)
    return out

def pending_snaps(wb, statements):
    have = {(r["account"], r["date"]) for r in account_rows(wb)}; out = []
    for st in statements:
        for sn in st.snaps:
            if (sn["account"], sn["date"]) not in have: have.add((sn["account"], sn["date"])); out.append(sn)
    return sorted(out, key=lambda x: (x["date"], x["account"]))

def _write_accounts(wb, snaps):
    from openpyxl.styles import Font, PatternFill
    if "Accounts" not in wb.sheetnames:
        ws = wb.create_sheet("Accounts")
        for c, h in enumerate(ACCOUNT_COLS, 1):
            cell = ws.cell(1, c, h); cell.font = Font(name="Arial", bold=True, color="FFFFFF"); cell.fill = PatternFill("solid", fgColor="1F3864")
        ws["O1"] = "Savings and loan balances from statements. Kept apart from the Bill Log, spending totals and budget; see the Accounts page in the app."
        ws["O1"].font = Font(name="Arial", italic=True)
        for c, w in zip("ABCDEFGHIJKLM", (14, 26, 10, 14, 12, 12, 12, 12, 12, 8, 12, 12, 44)): ws.column_dimensions[c].width = w
        ws.freeze_panes = "A2"
    ws = wb["Accounts"]
    rows = account_rows(wb) + list(snaps)
    for r in range(2, ws.max_row + 1):
        for c in range(1, len(ACCOUNT_KEYS) + 1): ws.cell(r, c).value = None
    for i, d in enumerate(sorted(rows, key=lambda x: (x["date"], x["account"])), 2):
        for c, k in enumerate(ACCOUNT_KEYS, 1):
            v = d.get(k)
            if isinstance(v, dt.date) and not isinstance(v, dt.datetime): v = dt.datetime.combine(v, dt.time())
            cell = ws.cell(i, c, v); cell.font = Font(name="Arial", size=10)
            if k in ("date", "due"): cell.number_format = "m/d/yyyy"
            elif k in ("balance", "interest", "ytd", "deposits", "withdrawals", "principal", "payment"): cell.number_format = "$#,##0.00;($#,##0.00);-"

def apply(new_rows, workbook=WORKBOOK, snaps=None):
    """Back up, then append new rows. Returns a summary dict."""
    from openpyxl import load_workbook
    BACKUPS.mkdir(exist_ok=True)
    backup = BACKUPS / f"Bills_Tracker_{dt.datetime.now():%Y%m%d_%H%M%S_%f}.xlsx"
    shutil.copy2(workbook, backup)
    wb = load_workbook(workbook)
    bl, ex, ls, rc, rv = wb["Bill Log"], wb["Excluded Transfers"], wb["Lists"], wb["Reconciliation"], wb["Review Items"]
    log_rows = sorted([r for r in new_rows if r.kind == "log"], key=lambda r: (r.txn.date, r.txn.account))
    ex_rows = sorted([r for r in new_rows if r.kind == "excluded"], key=lambda r: (r.txn.date, r.txn.account))
    # Bill Log: fill the next empty rows (formulas in C/D are already there)
    r = 6
    while bl.cell(r, 1).value is not None: r += 1
    if r + len(log_rows) - 1 > 2505: raise RuntimeError("Bill Log is full (2,500 rows). Archive a year first.")
    for row in log_rows:
        t = row.txn
        for c, v in ((1, dt.datetime.combine(t.date, dt.time())), (2, row.payee), (5, row.what), (6, t.amount), (7, t.account), (8, t.desc), (9, row.note or None)):
            bl.cell(r, c).value = v
        r += 1
    # Excluded Transfers
    er = 5
    while ex.cell(er, 1).value is not None: er += 1
    for row in ex_rows:
        t = row.txn
        for c, v in enumerate((dt.datetime.combine(t.date, dt.time()), t.account, t.desc, t.detail or None, abs(t.amount), row.direction, row.reason), 1):
            ex.cell(er, c).value = v
            _copy_style(ex.cell(5, c), ex.cell(er, c))
        er += 1
    # reconciliation formulas read 'Excluded Transfers'!E5:E111; widen them
    for rr in rc.iter_rows():
        for c in rr:
            if isinstance(c.value, str) and "$111" in c.value: c.value = c.value.replace("$111", "$2000")
    ex.auto_filter.ref = f"A4:G{max(er - 1, 111)}"
    # Lists: new payees
    lr = 4
    while ls.cell(lr, 1).value is not None and lr <= 153: lr += 1
    added_payees = []
    for row in log_rows:
        if row.is_new_payee and row.payee not in added_payees and row.payee not in {ls.cell(i, 1).value for i in range(4, lr)}:
            if lr > 153: break
            ls.cell(lr, 1).value, ls.cell(lr, 2).value, ls.cell(lr, 3).value = row.payee, row.category or "Other / Unsorted", row.type or "Spending"
            for c in (1, 2, 3): _copy_style(ls.cell(lr - 1, c), ls.cell(lr, c))
            added_payees.append(row.payee); lr += 1
    # Reconciliation: statement money-out added per account (statements are verified against their printed totals before import)
    by_acct = collections.Counter()
    for row in new_rows:
        if row.txn.amount > 0: by_acct[row.txn.account] += row.txn.amount
    for rr in range(5, 9):
        a = rc.cell(rr, 1).value
        if by_acct.get(a): rc.cell(rr, 2).value = round(rc.cell(rr, 2).value + by_acct[a], 2)
    # coverage lines
    lo, hi = {}, {}
    for sheet, a_col, d_col, first in ((bl, 7, 1, 6), (ex, 2, 1, 5)):
        for rr in sheet.iter_rows(min_row=first, values_only=True):
            d, a = rr[d_col - 1], rr[a_col - 1]
            if isinstance(d, dt.datetime) and a:
                lo[a] = min(lo.get(a, d), d); hi[a] = max(hi.get(a, d), d)
    for rr in range(1, rc.max_row + 1):
        a = rc.cell(rr, 1).value
        if a in lo and rc.cell(rr, 2).value and ("–" in str(rc.cell(rr, 2).value) or str(rc.cell(rr, 2).value) == "no data yet"):
            rc.cell(rr, 2).value = f"{lo[a]:%-m/%-d/%Y} – {hi[a]:%-m/%-d/%Y}"
    # Review Items
    if added_payees or any(r.needs_review for r in log_rows):
        n = max([x[0] for x in rv.iter_rows(min_row=5, max_col=1, values_only=True) if isinstance(x[0], int)] + [0]) + 1
        rr = rv.max_row + 1
        for c, v in enumerate((n, f"{len(added_payees)} new payee(s) from imported statements: " + ", ".join(added_payees[:12]) + ("…" if len(added_payees) > 12 else ""),
                               "Added to the Lists tab (category from the built-in merchant rules, or 'Other / Unsorted' when unknown) so they appear in every summary.",
                               "Set the right Category/Type on the Lists tab (and rename the payee if you like)."), 1):
            rv.cell(rr, c).value = v; _copy_style(rv.cell(rr - 1, c), rv.cell(rr, c))
    if snaps: _write_accounts(wb, snaps)
    wb.calculation.fullCalcOnLoad = True
    tmp = Path(str(workbook) + ".new")
    wb.save(tmp)
    _restore_charts(workbook, tmp)
    tmp.replace(workbook)
    return dict(backup=str(backup), log_rows=len(log_rows), excluded_rows=len(ex_rows), new_payees=added_payees, accounts=len(snaps or []))


# ---------------------------------------------------------------- categorizing existing payees
def _backup(workbook):
    BACKUPS.mkdir(exist_ok=True)
    b = BACKUPS / f"Bills_Tracker_{dt.datetime.now():%Y%m%d_%H%M%S_%f}.xlsx"
    shutil.copy2(workbook, b); return b

def _save_workbook(wb, workbook):
    wb.calculation.fullCalcOnLoad = True
    tmp = Path(str(workbook) + ".new"); wb.save(tmp); _restore_charts(workbook, tmp); tmp.replace(workbook)

def uncategorized(workbook=WORKBOOK):
    """Payees on the Lists tab still set to Other / Unsorted, with how much they account for."""
    from openpyxl import load_workbook
    wb = load_workbook(workbook, read_only=True)
    tot, cnt = collections.Counter(), collections.Counter()
    for r in wb["Bill Log"].iter_rows(min_row=6, max_col=6, values_only=True):
        if r[0] and r[1]: tot[r[1]] += r[5] or 0; cnt[r[1]] += 1
    cats = _categories(wb)
    out = [dict(payee=r[0], category=r[1] or "Other / Unsorted", type=r[2] or "Spending", n=cnt[r[0]], total=round(tot[r[0]], 2))
           for r in wb["Lists"].iter_rows(min_row=4, max_row=153, max_col=3, values_only=True)
           if r[0] and (not r[1] or r[1] == "Other / Unsorted")]
    wb.close()
    return dict(payees=sorted(out, key=lambda x: -x["total"]), categories=cats, types=["Bill", "Spending"])

def set_categories(changes, workbook=WORKBOOK):
    """changes: [{payee, category, type}] -> written to the Lists tab (backup first)."""
    from openpyxl import load_workbook
    b = _backup(workbook); wb = load_workbook(workbook); ls = wb["Lists"]
    cats = set(_categories(wb)); rows = {ls.cell(r, 1).value: r for r in range(4, 154) if ls.cell(r, 1).value}
    n = 0
    for ch in changes:
        r = rows.get(ch.get("payee"))
        if r and ch.get("category") in cats and ch.get("type") in ("Bill", "Spending"):
            ls.cell(r, 2).value, ls.cell(r, 3).value = ch["category"], ch["type"]; n += 1
    _save_workbook(wb, workbook)
    return dict(updated=n, backup=b.name)

def add_category(name, workbook=WORKBOOK):
    """Add a category (envelope) to the Lists tab: first free slot in E4:E23, else E25:E40 (the web pages show them all;
    the Excel summary tabs only have 20 slots)."""
    from openpyxl import load_workbook
    from copy import copy
    name = re.sub(r"\s+", " ", (name or "")).strip()
    if not name or len(name) > 40: raise ValueError("Give the category a name (40 characters or fewer).")
    b = _backup(workbook); wb = load_workbook(workbook); ls = wb["Lists"]
    if name.lower() in {str(c).lower() for c in _categories(wb)}: raise ValueError(f"'{name}' already exists.")
    row = next((r for r in list(range(4, 24)) + list(range(25, 41)) if ls.cell(r, 5).value is None), None)
    if row is None: raise ValueError("The Lists tab has no free category slots left.")
    cell = ls.cell(row, 5); cell.value = name
    src = ls.cell(4, 5); cell.font, cell.border, cell.alignment, cell.fill = copy(src.font), copy(src.border), copy(src.alignment), copy(src.fill)
    if row > 23:
        for dv in ls.data_validations.dataValidation:
            if dv.formula1 == "$E$4:$E$23": dv.formula1 = "$E$4:$E$40"
    _save_workbook(wb, workbook)
    return dict(name=name, backup=b.name, beyond_excel=row > 23)

def autocategorize(workbook=WORKBOOK):
    """Apply the built-in merchant rules to payees still set to Other / Unsorted."""
    from openpyxl import load_workbook
    b = _backup(workbook); wb = load_workbook(workbook); ls = wb["Lists"]; cats = set(_categories(wb))
    descs = collections.defaultdict(list)
    for r in wb["Bill Log"].iter_rows(min_row=6, max_col=8, values_only=True):
        if r[0] and r[1] and r[7] and len(descs[r[1]]) < 3: descs[r[1]].append(str(r[7]))
    n = left = 0
    for r in range(4, 154):
        p = ls.cell(r, 1).value
        if not p or (ls.cell(r, 2).value and ls.cell(r, 2).value != "Other / Unsorted"): continue
        sm = seed_match(" ".join([p] + descs[p]))
        if sm and sm[1] in cats: ls.cell(r, 2).value, ls.cell(r, 3).value = sm[1], sm[2]; n += 1
        else: left += 1
    _save_workbook(wb, workbook)
    return dict(updated=n, remaining=left, backup=b.name)


def set_budget(month, budgets, workbook=WORKBOOK):
    """Replace the envelope amounts for one month (YYYY-MM) on the Budget tab. budgets: {envelope: amount}."""
    from openpyxl import load_workbook
    from openpyxl.styles import Font, PatternFill, Alignment
    if not re.fullmatch(r"\d{4}-\d\d", month or ""): raise ValueError("bad month")
    b = _backup(workbook); wb = load_workbook(workbook)
    if "Budget" not in wb.sheetnames:
        ws = wb.create_sheet("Budget")
        ws["A1"], ws["B1"], ws["C1"] = "Month", "Envelope", "Budgeted"
        for c in ("A1", "B1", "C1"):
            ws[c].font = Font(name="Arial", bold=True, color="FFFFFF"); ws[c].fill = PatternFill("solid", fgColor="1F3864")
        ws["E1"] = "Envelope budget amounts, one row per month and envelope. Edit them on the app's Budget page; spending comes from the Bill Log."
        ws["E1"].font = Font(name="Arial", italic=True)
        ws.column_dimensions["A"].width = 10; ws.column_dimensions["B"].width = 28; ws.column_dimensions["C"].width = 12
        ws.freeze_panes = "A2"
    ws = wb["Budget"]
    keep = []
    for r in ws.iter_rows(min_row=2, max_col=3, values_only=True):
        if r[0] is None or r[1] is None: continue
        m = r[0] if isinstance(r[0], str) else f"{r[0]:%Y-%m}"
        if m != month: keep.append((m, r[1], r[2] or 0))
    for r in range(2, ws.max_row + 1):
        for c in (1, 2, 3): ws.cell(r, c).value = None
    new = [(month, k, round(float(v), 2)) for k, v in budgets.items() if v not in (None, "") and round(float(v), 2) != 0]
    for i, row in enumerate(sorted(keep + new), 2):
        for c, v in enumerate(row, 1):
            ws.cell(i, c).value = v
            ws.cell(i, c).font = Font(name="Arial", size=10)
        ws.cell(i, 3).number_format = '$#,##0.00;($#,##0.00);-'
    _save_workbook(wb, workbook)
    return dict(rows=len(new), backup=b.name)

def summarize(rows):
    return dict(log=sum(r.kind == "log" for r in rows), excluded=sum(r.kind == "excluded" for r in rows),
                new_payees=sorted({r.payee for r in rows if r.is_new_payee}))

# ---------------------------------------------------------------- filing statements away
ARCHIVE = DATA / "Statements"                 # Statements/YYYY/MM/<file>, by statement end date
INDEX = DATA / "statements_index.json"        # what each filed statement covered, so coverage needs no re-reading

def _index():
    try: return json.loads(INDEX.read_text())
    except Exception: return []

def archived_statements():
    """Light Statement objects (no transactions) for everything already filed away."""
    out = []
    for i in _index():
        try: out.append(Statement(i["file"], i["account"], dt.date.fromisoformat(i["start"]) if i.get("start") else None,
                                  dt.date.fromisoformat(i["end"]), kind=i["kind"]))
        except Exception: pass
    return out

def archive_statements():
    """Once everything read is in the workbook, move the statement files into Statements/YYYY/MM and remember what they covered."""
    res = run_preview(DATA)
    if res["new"] or res["snaps"]: raise ValueError("Add the new transactions to the workbook first.")
    idx, batch, moved = _index(), f"{dt.datetime.now():%Y%m%d%H%M%S}", 0
    for st in res["statements"]:
        if st.file in res["bad"]: continue
        src = DATA / st.file
        if not src.exists(): continue
        dest = ARCHIVE / f"{st.end.year}" / f"{st.end.month:02d}"; dest.mkdir(parents=True, exist_ok=True)
        target, n = dest / src.name, 1
        while target.exists(): target = dest / f"{src.stem}_{n}{src.suffix}"; n += 1
        shutil.move(str(src), str(target))
        if st.kind == "pdf": idx = [i for i in idx if (i["account"], i["end"], i["kind"]) != (st.account, st.end.isoformat(), st.kind)]
        idx.append(dict(file=str(target.relative_to(DATA)), account=st.account, start=st.start.isoformat() if st.start else None,
                        end=st.end.isoformat(), kind=st.kind, batch=batch)); moved += 1
    INDEX.write_text(json.dumps(idx, indent=1))
    if moved and res["dups"]: moved += archive_statements()["moved"]      # identical copies read as duplicates until the original is filed
    return dict(moved=moved)

def unarchive_last():
    """Put the most recently filed batch of statements back in the data folder (used by Undo)."""
    idx = _index()
    if not idx or not any(i.get("batch") for i in idx): return 0
    last = max(i.get("batch", "") for i in idx); n = 0
    for i in [i for i in idx if i.get("batch") == last]:
        src = DATA / i["file"]
        if src.exists():
            target = DATA / src.name
            if not target.exists(): shutil.move(str(src), str(target)); n += 1
    INDEX.write_text(json.dumps([i for i in idx if i.get("batch") != last], indent=1))
    return n

def run_preview(folder=DATA):
    from openpyxl import load_workbook
    sts, skipped, dups = load_statements(folder)
    checks = check_statements(sts)
    bad = {f for f, p, s, ok in checks if not ok}
    good = [s for s in sts if s.file not in bad]
    wb = load_workbook(WORKBOOK)
    new, already = plan(wb, good)
    return dict(statements=sts, skipped=skipped, dups=dups, checks=checks, bad=bad, new=new, already=already, snaps=pending_snaps(wb, good))

if __name__ == "__main__":
    cmd = sys.argv[1] if len(sys.argv) > 1 else "preview"
    res = run_preview()
    for f, p, st, ok in res["checks"]:
        if not ok: print(f"!! {f}: parsed {p:,.2f} but statement says {st:,.2f} - NOT imported")
    for f, why in res["skipped"]: print(f"skipped {f}: {why}")
    new = res["new"]; sm = summarize(new)
    print(f"{len(res['statements'])} statements read, {res['already']} transactions already in workbook, {len(new)} new "
          f"({sm['log']} to Bill Log, {sm['excluded']} to Excluded Transfers).")
    for r in new[:200]:
        print(f"  {r.txn.date} {r.txn.account:16} {r.txn.amount:>10,.2f}  {r.kind:8} {r.payee or r.reason}  <- {r.txn.desc[:40]}")
    if sm["new_payees"]: print("New payees to categorize:", ", ".join(sm["new_payees"]))
    if cmd == "import" and new:
        print(apply(new))
    elif cmd == "import": print("Nothing new to import.")