import sys
import time
_T0 = time.perf_counter()  # absolute zero — for measuring full window-open path
def _ms(t):
    return (time.perf_counter() - t) * 1000
import os
import copy
import sqlite3
import numpy as np
import pandas as pd
from datetime import datetime
print(f"[timing] imports (stdlib + pandas + numpy):       {_ms(_T0):7.1f} ms")
_T_AFTER_DATA_IMPORTS = time.perf_counter()
from PySide6.QtWidgets import (
    QApplication, QMainWindow, QTreeView, QVBoxLayout,
    QHBoxLayout, QPushButton, QWidget, QStackedWidget, QComboBox, QLabel, QCheckBox, QMessageBox,
    QDialog, QListWidget, QDialogButtonBox, QMenu, QAbstractItemView,
    QRadioButton, QButtonGroup, QTabWidget, QTabBar
)
from PySide6.QtGui import QStandardItemModel, QStandardItem, QColor, QBrush, QFont
from PySide6.QtCore import Qt, QTimer, QAbstractItemModel, QModelIndex, QEvent, Signal
print(f"[timing] PySide6 imports:                          {_ms(_T_AFTER_DATA_IMPORTS):7.1f} ms")

# ============================================================
# MRR MOVEMENT CATEGORY CODES
# ============================================================
#
#  2   Add On, New SOW / Amendment
#  21  Add On, New SOW / Amendment (incr. quantity)
#  22  Add On, New SOW / Amendment (incr. pricing)
#  23  Add On, BCO (standard)
#  24  Add On, BCO (incr. quantity)
#  25  Upgrade at Renewal
#  101 Overages / Underages
#  100 Accounting Driven
#  1   New Customer
#  11  New Customer (BCO)
#  12  New Customer (Amendment)           [manual]
#  13  New Customer on Ramp               [manual]
#  14  New Customer, Returning            [manual]
# -2   Downgrade, New SOW / Amendment
# -21  Downgrade, New SOW / Amendment (decr. quantity)
# -22  Downgrade, New SOW / Amendment (decr. pricing)
# -23  Downgrade, BCO (standard)
# -24  Downgrade, BCO (decr. quantity)
# -25  Downgrade at Renewal
# -1   Lost customer, voluntary
# -3   Lost customer, involuntary         [manual]

# Custom Qt item data roles for category cells
ROLE_CUSTOMER_ID     = Qt.UserRole + 1
ROLE_CONTRACT_ID     = Qt.UserRole + 2
ROLE_MONTH           = Qt.UserRole + 3
ROLE_IS_CATEGORY_CELL  = Qt.UserRole + 4
ROLE_CATEGORY_CODE     = Qt.UserRole + 5
ROLE_REF_NO              = Qt.UserRole + 6
ROLE_IS_UNASSIGNED_REF   = Qt.UserRole + 7
ROLE_ASSIGNED_CONTRACT_ID = Qt.UserRole + 8   # current contract_id for pre-selection in dialog
ROLE_CAN_CLEAR            = Qt.UserRole + 9   # True only when a manual DB assignment exists and can be cleared

# Sentinel values produced by heal_and_fill when source data lacks a real id.
# Used to decide whether a row is interactive (i.e. has a real customer/contract).
MISSING_SENTINELS = {"Missing ID", "Missing Info", "Uncategorized"}

# Ordered list of all valid category codes with descriptions
ALL_CATEGORIES = [
    (1,    "New Customer"),
    (11,   "New Customer (BCO)"),
    (12,   "New Customer (Amendment)"),
    (13,   "New Customer on Ramp"),
    (14,   "New Customer, Returning"),
    (2,    "Add On, New SOW / Amendment"),
    (21,   "Add On (increase quantity)"),
    (22,   "Add On (increase pricing)"),
    (23,   "Add On, BCO (standard)"),
    (24,   "Add On, BCO (increase quantity)"),
    (25,   "Upgrade at Renewal"),
    (100,  "Accounting Driven"),
    (101,  "Overages / Underages"),
    (-1,   "Lost Customer, Voluntary"),
    (-2,   "Downgrade, New SOW / Amendment"),
    (-21,  "Downgrade (decrease quantity)"),
    (-22,  "Downgrade (decrease pricing)"),
    (-23,  "Downgrade, BCO (standard)"),
    (-24,  "Downgrade, BCO (decrease quantity)"),
    (-25,  "Downgrade at Renewal"),
    (-3,   "Lost Customer, Involuntary"),
]

# ============================================================
# DATABASE & UTILS
# ============================================================

def init_database():
    db_path = "data_source/mrr_data.db"
    os.makedirs("data_source", exist_ok=True)
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='contracts'")
    exists = cursor.fetchone()

    if not exists:
        print("--- [CHECK] Database not found. Creating new DB from CSV sources... ---")
        try:
            df_l = pd.read_csv("data_source/ledger.csv", low_memory=False)
            df_c = pd.read_csv("data_source/contracts.csv", low_memory=False)
            df_l = clean_headers(df_l)
            df_c = clean_headers(df_c)
            df_l.to_sql("ledger", conn, index=False)
            df_c.to_sql("contracts", conn, index=False)
            conn.execute("CREATE TABLE manual_categorization (customer_name TEXT, contract_name TEXT, line_number TEXT, category TEXT, updated_at TIMESTAMP)")
            conn.execute("CREATE TABLE write_locks (lock_key TEXT PRIMARY KEY, user_name TEXT, timestamp DATETIME)")
            conn.commit()
            print("--- [CHECK] Database initialized successfully. ---")
        except Exception as e:
            print(f"--- [ERROR] Initial Database Creation Failed: {e} ---")
    else:
        print("--- [CHECK] Existing database found. Skipping CSV ingestion. ---")

    # Non-destructive top-up: load contractsbilling.csv into its own table if missing.
    # Lets older databases pick up the new billing source without a rebuild.
    cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='contracts_billing'")
    if not cursor.fetchone():
        billing_csv = "data_source/contractsbilling.csv"
        if os.path.exists(billing_csv):
            print("--- [CHECK] Loading contractsbilling.csv into contracts_billing table... ---")
            try:
                df_b = pd.read_csv(billing_csv, low_memory=False)
                df_b = clean_headers(df_b)
                df_b.to_sql("contracts_billing", conn, index=False)
                conn.commit()
                print(f"--- [CHECK] Loaded {len(df_b)} billing rows. ---")
            except Exception as e:
                print(f"--- [ERROR] Billing data load failed: {e} ---")
        else:
            print(f"--- [WARN] {billing_csv} not found; backlog view will be empty. ---")

    return conn

def clean_headers(df):
    df.columns = df.columns.str.strip().str.lower().str.replace(" ", "_")
    return df

def ensure_category_overrides_table():
    conn = sqlite3.connect("data_source/mrr_data.db")
    conn.execute("""
        CREATE TABLE IF NOT EXISTS category_overrides (
            customer_id TEXT,
            contract_id TEXT,
            month       TEXT,
            category    INTEGER,
            updated_at  TIMESTAMP,
            PRIMARY KEY (customer_id, contract_id, month)
        )
    """)
    conn.commit()
    conn.close()

def load_category_overrides() -> dict:
    """Return {(customer_id, contract_id): {month: category_code}} for all saved overrides."""
    conn = sqlite3.connect("data_source/mrr_data.db")
    try:
        df = pd.read_sql("SELECT customer_id, contract_id, month, category FROM category_overrides", conn)
    except Exception:
        return {}
    finally:
        conn.close()
    result = {}
    for _, row in df.iterrows():
        key = (str(row["customer_id"]), str(row["contract_id"]))
        result.setdefault(key, {})[row["month"]] = int(row["category"])
    return result

def save_category_override(customer_id: str, contract_id: str, month: str, category: int):
    conn = sqlite3.connect("data_source/mrr_data.db")
    conn.execute(
        "INSERT OR REPLACE INTO category_overrides (customer_id, contract_id, month, category, updated_at) VALUES (?, ?, ?, ?, ?)",
        (customer_id, contract_id, month, category, datetime.now())
    )
    conn.commit()
    conn.close()

def delete_category_override(customer_id: str, contract_id: str, month: str):
    conn = sqlite3.connect("data_source/mrr_data.db")
    conn.execute(
        "DELETE FROM category_overrides WHERE customer_id=? AND contract_id=? AND month=?",
        (customer_id, contract_id, month)
    )
    conn.commit()
    conn.close()

# ── Ledger → Contract manual assignments ─────────────────────────────────────

def ensure_ledger_assignments_table():
    conn = sqlite3.connect("data_source/mrr_data.db")
    conn.execute("""
        CREATE TABLE IF NOT EXISTS ledger_assignments (
            reference_no  TEXT,
            customer_id   TEXT,
            contract_id   TEXT,
            contract_name TEXT,
            updated_at    TIMESTAMP,
            PRIMARY KEY (reference_no, customer_id)
        )
    """)
    conn.commit()
    conn.close()

def load_ledger_assignments() -> dict:
    """Return {(reference_no, customer_id): (contract_id, contract_name)}."""
    conn = sqlite3.connect("data_source/mrr_data.db")
    try:
        df = pd.read_sql(
            "SELECT reference_no, customer_id, contract_id, contract_name FROM ledger_assignments",
            conn
        )
    except Exception:
        return {}
    finally:
        conn.close()
    return {
        (str(r["reference_no"]), str(r["customer_id"])): (str(r["contract_id"]), str(r["contract_name"]))
        for _, r in df.iterrows()
    }

def delete_ledger_assignment(reference_no: str, customer_id: str):
    conn = sqlite3.connect("data_source/mrr_data.db")
    conn.execute(
        "DELETE FROM ledger_assignments WHERE reference_no=? AND customer_id=?",
        (reference_no, customer_id)
    )
    conn.commit()
    conn.close()

def save_ledger_assignment(reference_no: str, customer_id: str, contract_id: str, contract_name: str):
    conn = sqlite3.connect("data_source/mrr_data.db")
    conn.execute(
        "INSERT OR REPLACE INTO ledger_assignments "
        "(reference_no, customer_id, contract_id, contract_name, updated_at) VALUES (?, ?, ?, ?, ?)",
        (reference_no, customer_id, contract_id, contract_name, datetime.now())
    )
    conn.commit()
    conn.close()

def apply_ledger_assignments(df_l: pd.DataFrame, assignments: dict) -> pd.DataFrame:
    """
    Patch contract_id / contract_name on raw ledger rows that have a saved manual assignment.

    Must run BEFORE build_ledger_pivot so the pivot groups the row under its new contract.

    The reference number and customer_id keys stored in the assignments dict come from the
    pivot (post heal_and_fill), where nulls have been collapsed to "Missing Info" /
    "Missing ID". We normalise the raw DataFrame the same way before matching, otherwise
    NaN-bearing rows silently fail to match (pandas.Series.astype(str) preserves NaN).
    """
    if not assignments:
        return df_l
    ref_col = "reference_no." if "reference_no." in df_l.columns else "reference_no"
    if ref_col not in df_l.columns or "customer_id" not in df_l.columns:
        return df_l

    def _norm(series, null_token):
        s = series.fillna(null_token).astype(str).str.strip()
        return s.replace(["", "nan", "None", "NaN", "<NA>"], null_token)

    df_l["_ref_norm"]  = _norm(df_l[ref_col],       "Missing Info")
    df_l["_cust_norm"] = _norm(df_l["customer_id"], "Missing ID")

    # Build a small DataFrame of assignments and merge it onto df_l in one pass,
    # instead of running a boolean mask scan per assignment.
    assign_df = pd.DataFrame(
        [(str(r), str(c), cid, cname)
         for (r, c), (cid, cname) in assignments.items()],
        columns=["_ref_norm", "_cust_norm", "_new_contract_id", "_new_contract_name"],
    )
    df_l = df_l.merge(assign_df, on=["_ref_norm", "_cust_norm"], how="left")

    matched = df_l["_new_contract_id"].notna()
    df_l.loc[matched, "contract_id"]   = df_l.loc[matched, "_new_contract_id"]
    df_l.loc[matched, "contract_name"] = df_l.loc[matched, "_new_contract_name"]

    return df_l.drop(columns=["_ref_norm", "_cust_norm",
                              "_new_contract_id", "_new_contract_name"])

# ============================================================
# DATA PROCESSING & MOM CALCULATION
# ============================================================

def heal_and_fill(df, index_cols):
    for id_col in ["customer_id", "contract_id"]:
        if id_col in df.columns:
            df[id_col] = df[id_col].astype(str).replace(["", "nan", "None"], "Missing ID")

    if "customer_id" in df.columns and "customer_name" in df.columns:
        valid_cust_names = df[df["customer_name"].notna() & (df["customer_name"] != "")].groupby("customer_id")["customer_name"].first()
        df["customer_name"] = df["customer_name"].fillna(df["customer_id"].map(valid_cust_names))

    if "contract_id" in df.columns and "contract_name" in df.columns:
        valid_cont_names = df[df["contract_name"].notna() & (df["contract_name"] != "")].groupby("contract_id")["contract_name"].first()
        df["contract_name"] = df["contract_name"].fillna(df["contract_id"].map(valid_cont_names))

    for col in index_cols:
        if col not in df.columns:
            df[col] = "Uncategorized" if "contract" in col else "Missing Info"
        elif "contract" in col:
            df[col] = df[col].fillna("Uncategorized").replace(["", "nan", "None"], "Uncategorized")
        else:
            df[col] = df[col].fillna("Missing Info").replace(["", "nan", "None"], "Missing Info")
    return df

# Items excluded from contracts processing (one-off / non-MRR SKUs).
# Used by the legacy/full contracts pipeline that feeds discrepancy highlighting
# and the categorisation engine. The per-tab views below use their own filters.
_EXCLUDED_CONTRACT_ITEMS = (
    'I590', 'I900', 'I901', 'I902', 'I904', 'I905',
    'L214', 'L217', 'L231', 'L310', 'L311', 'L198',
)
_CONTRACT_INDEX_COLS = ["customer_id", "customer_name", "contract_id", "contract_name", "line_number"]

# Contracts View sub-tabs. Each tab shows a pivot filtered to the listed items
# (or, for MS, the inverse of the listed exclusion set).
_CONTRACT_TAB_PS_ITEMS           = ('I900', 'I901', 'I902', 'I903', 'I904', 'I905', 'L231', 'L198')
_CONTRACT_TAB_MS_EXCLUDED_ITEMS  = (
    'I590', 'I900', 'I901', 'I902', 'I903', 'I904', 'I905',
    'L214', 'L217', 'L231', 'L310', 'L311', 'L198', 'L401',
)
_CONTRACT_TAB_CLOUD_ITEMS        = ('L203', 'L204', 'L401')
_CONTRACT_TAB_OTHER_RESELL_ITEMS = ('L214', 'L217', 'L310', 'L311')

# Professional Services items — source rows for the Backlog (project backlog) report.
_PS_ITEMS_GENERAL   = ('I900', 'I901', 'I902', 'I903', 'I905', 'L231')
_PS_ITEMS_RECURRING = ('I904',)


def build_backlog_data(df: pd.DataFrame, items=_PS_ITEMS_GENERAL):
    """
    Source: contracts_billing table (billing schedule, not revenue).
    Filter to the given PS item codes, drop status='Terminated' rows, parse month/amount,
    and aggregate to one row per (customer_id, customer_name, contract_id,
    contract_name, month).

    Also returns a set of (customer_id, contract_id) keys for contracts that
    still have unbilled PS rows (billed=False) — i.e. contracts with
    remaining scheduled billings.
    """
    if df is None or df.empty:
        return pd.DataFrame(), set()
    if 'item' in df.columns:
        df = df[df['item'].astype(str).isin(items)]
    if 'status' in df.columns:
        df = df[df['status'].astype(str).str.strip().str.lower() != 'terminated']
    if df.empty:
        return df.iloc[0:0], set()

    df = df.copy()
    parsed_dates = pd.to_datetime(df["scheduled_billing_date"], errors="coerce")
    df["month"]  = parsed_dates.dt.to_period("M").astype(str)
    df["amount"] = pd.to_numeric(df["amount"], errors="coerce").fillna(0)
    df = df.dropna(subset=["month"])
    df = heal_and_fill(df, ["customer_id", "customer_name", "contract_id", "contract_name"])

    # `billed` round-trips through SQLite as 0/1 INTEGER, but be defensive in
    # case the column ever lands as 'True'/'False' strings.
    billed_str = df["billed"].astype(str).str.strip().str.lower()
    is_billed = billed_str.isin(['true', '1'])
    df["_unbilled"] = ~is_billed
    df["_billed_amt"] = df["amount"].where(is_billed, 0.0)

    has_remaining = (
        df.groupby(["customer_id", "contract_id"])["_unbilled"].any()
    )
    has_remaining_keys = {
        (str(c), str(ct)) for (c, ct), v in has_remaining.items() if v
    }

    monthly = (
        df.groupby(
            ["customer_id", "customer_name", "contract_id", "contract_name", "month"],
            sort=False,
        ).agg(amount=("amount", "sum"), billed=("_billed_amt", "sum")).reset_index()
    )
    return monthly, has_remaining_keys

def _preprocess_contracts_df(df, item_filter=None):
    """
    Filter out excluded/terminated rows, parse month/amount, heal nullable id columns,
    and ensure mrr_category and item_description exist as strings.

    *item_filter* is an optional callable taking the item-column Series and
    returning a boolean mask of rows to keep. If None, the default exclusion
    list (_EXCLUDED_CONTRACT_ITEMS) is applied — preserving legacy behaviour.
    """
    if 'item' in df.columns:
        items = df['item'].astype(str)
        if item_filter is None:
            df = df[~items.isin(_EXCLUDED_CONTRACT_ITEMS)]
        else:
            df = df[item_filter(items)]
    if 'status' in df.columns:
        df = df[df['status'].astype(str).str.strip().str.lower() != 'terminated']

    df["month"]  = pd.to_datetime(df["scheduled_posting_date"], errors="coerce").dt.to_period("M").astype(str)
    df["amount"] = pd.to_numeric(df["amount"], errors="coerce").fillna(0)
    df = df.dropna(subset=["month"])

    df = heal_and_fill(df, _CONTRACT_INDEX_COLS)

    for col in ["mrr_category", "item_description"]:
        if col not in df.columns:
            df[col] = ""
        else:
            df[col] = df[col].fillna("").astype(str)

    return df

# Order is significant — drives the tab order in the Contracts View.
CONTRACT_TAB_KEYS = ('PS', 'MS', 'Cloud', 'Other Resell')

def build_contract_outputs(df):
    """
    Preprocess the contracts frame once and return both the pivot (used by the
    Contracts View) and the detail DataFrame (used by the categorisation engine).
    Avoids re-running _preprocess_contracts_df twice on the same input.
    """
    df = _preprocess_contracts_df(df)
    pivot = pd.pivot_table(df, values="amount", index=_CONTRACT_INDEX_COLS,
                           columns="month", aggfunc="sum", fill_value=0).reset_index()
    detail = (
        df.groupby(_CONTRACT_INDEX_COLS + ["mrr_category", "item_description", "month"])["amount"]
          .sum().reset_index()
    )
    return pivot, detail

# Ledger View sub-tabs. MS keeps the historical allowed-account list; PS / Cloud /
# Other Resell pull from different revenue accounts. Tab order drives the QTabWidget.
_LEDGER_TAB_MS_ACCOUNTS = (
    '40101', '40202', '40204', '40205', '40206', '40207',
    '40209', '40211', '40304', '42006', '43004', '44403',
)
_LEDGER_TAB_PS_ACCOUNTS = (
    '42001', '42002', '42003', '42004', '42005',
    '43005', '43006', '43007', '43008',
)
_LEDGER_TAB_CLOUD_ACCOUNTS = ('43001', '43002', '43003', '49100')
_LEDGER_TAB_OTHER_RESELL_ACCOUNTS = (
    '44101', '44102', '44103', '44201', '44202', '44203',
    '44301', '44302', '44401', '44402', '44404',
)
LEDGER_TAB_KEYS = ('PS', 'MS', 'Cloud', 'Other Resell')

_LEDGER_TAB_ACCOUNTS = {
    'PS':           _LEDGER_TAB_PS_ACCOUNTS,
    'MS':           _LEDGER_TAB_MS_ACCOUNTS,
    'Cloud':        _LEDGER_TAB_CLOUD_ACCOUNTS,
    'Other Resell': _LEDGER_TAB_OTHER_RESELL_ACCOUNTS,
}


def build_ledger_pivot(df, allowed_accounts=_LEDGER_TAB_MS_ACCOUNTS):
    df = df.copy()
    df['base_account'] = df['account'].astype(str).str.split('-').str[0]
    df = df[df['base_account'].isin(allowed_accounts)].copy()

    ref_col = "reference_no." if "reference_no." in df.columns else "reference_no"
    index_cols = ["customer_id", "customer_name", "contract_id", "contract_name", ref_col]
    df = heal_and_fill(df, index_cols)
    df = df[df["customer_id"] != "Missing Info"]   # drop rows with no identifiable customer

    df["month"] = pd.to_datetime(df["posting_date"], errors="coerce").dt.to_period("M").astype(str)
    df["amount"] = pd.to_numeric(df["amount"], errors="coerce").fillna(0) * (-1)

    pivot = pd.pivot_table(df, values="amount", index=index_cols, columns="month", aggfunc="sum", fill_value=0)
    return pivot.reset_index()

def calculate_mom_diff(df):
    meta_cols = [c for c in df.columns if "-" not in str(c)]
    month_cols = [c for c in df.columns if "-" in str(c)]
    diff_df = df[meta_cols].copy()
    for i in range(len(month_cols)):
        curr_m = month_cols[i]
        if i == 0:
            diff_df[curr_m] = df[curr_m]
        else:
            prev_m = month_cols[i-1]
            diff_df[curr_m] = df[curr_m] - df[prev_m]
    return diff_df


# ============================================================
# MRR CATEGORIZATION ENGINE
# ============================================================

def build_contract_category_map(detail_df: pd.DataFrame, all_months: list) -> dict:
    """
    Returns a nested dict:
        category_map[(customer_id, contract_id)][month] = dominant_category_code (int or 0)

    Logic per (customer_id, contract_id, month):
      1. Compute MoM diff per line (comparing to same line's amount in previous month).
      2. Classify each line's diff into a bucket via the np.select cascade below.
      3. The dominant bucket is whichever has the largest abs() contribution.
      4. If no movement → 0.

    Vectorised end-to-end (~12s → <1s on the production dataset).
    """
    if detail_df is None or detail_df.empty:
        return {}

    id_cols = ["customer_id", "contract_id", "line_number", "mrr_category", "item_description"]
    line_pivot = detail_df.pivot_table(
        values="amount", index=id_cols, columns="month",
        aggfunc="sum", fill_value=0,
    )

    # Initialise category_map with 0 for every (customer, contract) × every month.
    sorted_months = sorted(all_months)
    cust_contract_pairs = (
        line_pivot.index.to_frame(index=False)[["customer_id", "contract_id"]]
        .drop_duplicates()
    )
    category_map = {
        (str(c), str(ct)): {m: 0 for m in sorted_months}
        for c, ct in cust_contract_pairs.itertuples(index=False, name=None)
    }

    # Time axis: every month in `all_months`. Missing columns become zeros, so
    # diff against them naturally yields zero movement.
    amounts = line_pivot.reindex(columns=sorted_months, fill_value=0)
    prev = amounts.shift(1, axis=1).fillna(0)
    diffs = amounts - prev

    # "New customer" flag: fewer than 12 months of non-zero MRR anywhere in the pivot.
    lp_months = [c for c in line_pivot.columns if "-" in str(c)]
    if lp_months:
        cust_nonzero_months = (
            (line_pivot[lp_months].abs() > 0.01)
            .groupby(level="customer_id").any().sum(axis=1)
        )
        is_new_cust_series = cust_nonzero_months < 12
    else:
        is_new_cust_series = pd.Series(dtype=bool)

    # "Drops to zero" flag per (customer, contract, month).
    ct_totals = amounts.groupby(level=["customer_id", "contract_id"]).sum()
    ct_prev = ct_totals.shift(1, axis=1).fillna(0)
    drops_df = (ct_totals.abs() < 0.01) & (ct_prev.abs() > 0.01)

    # Melt to long form and keep only meaningful movements.
    diffs_long = (
        diffs.reset_index()
        .melt(id_vars=id_cols, value_vars=sorted_months,
              var_name="month", value_name="diff")
    )
    diffs_long = diffs_long[diffs_long["diff"].abs() >= 0.01]
    if diffs_long.empty:
        return category_map

    # Join the per-row flags.
    diffs_long["is_new_cust"] = (
        diffs_long["customer_id"].map(is_new_cust_series).fillna(True).astype(bool)
    )
    drops_long = drops_df.stack().rename("drops_to_zero")
    diffs_long = diffs_long.merge(
        drops_long,
        left_on=["customer_id", "contract_id", "month"],
        right_index=True, how="left",
    )
    diffs_long["drops_to_zero"] = diffs_long["drops_to_zero"].fillna(False).astype(bool)

    # Classification cascade. np.select picks the first matching condition,
    # so order encodes priority (new-customer ⟶ lost ⟶ renewal ⟶ BCO ⟶ standard).
    cat = diffs_long["mrr_category"].fillna("").astype(str).str.strip()
    desc = diffs_long["item_description"].fillna("").astype(str).str.strip().str.upper()
    is_bco_v = desc.str.startswith("RITM")
    increasing = diffs_long["diff"] > 0
    is_new = diffs_long["is_new_cust"]
    drops = diffs_long["drops_to_zero"]
    cat_qty   = cat == "Quantity"
    cat_price = cat == "Pricing"
    cat_up    = cat == "Upgrade at Renewal"
    cat_down  = cat == "Downgrade at Renewal"

    conditions = [
        is_new & is_bco_v,                       # 11  — new customer (BCO)
        is_new,                                  #  1  — new customer
        drops & ~increasing,                     # -1  — lost customer (voluntary)
        cat_up,                                  # 25  — upgrade at renewal
        cat_down,                                # -25 — downgrade at renewal
        is_bco_v & increasing & cat_qty,         # 24  — BCO add-on, quantity
        is_bco_v & increasing,                   # 23  — BCO add-on, standard
        is_bco_v & ~increasing & cat_qty,        # -24 — BCO downgrade, quantity
        is_bco_v & ~increasing,                  # -23 — BCO downgrade, standard
        increasing & cat_qty,                    # 21  — add-on, quantity
        increasing & cat_price,                  # 22  — add-on, pricing
        increasing,                              #  2  — add-on, amendment / SOW
        ~increasing & cat_qty,                   # -21 — downgrade, quantity
        ~increasing & cat_price,                 # -22 — downgrade, pricing
        # default (decreasing, neither qty nor pricing)  → -2
    ]
    codes = [11, 1, -1, 25, -25, 24, 23, -24, -23, 21, 22, 2, -21, -22]
    diffs_long["code"] = np.select(conditions, codes, default=-2)
    diffs_long["abs_diff"] = diffs_long["diff"].abs()

    # Dominant bucket per (customer, contract, month) = code with the largest sum |diff|.
    # sort=False preserves the row order from melt, so ties resolve to the code that
    # first appears in the pivot — matching the legacy's dict-insertion tiebreaker.
    # Round the bucket totals to cents before comparing to neutralise float-precision
    # noise between pandas' .sum() and the legacy's manual += accumulator.
    bucket_sums = (
        diffs_long.groupby(["customer_id", "contract_id", "month", "code"],
                           sort=False)["abs_diff"]
        .sum().reset_index()
    )
    bucket_sums["abs_diff"] = bucket_sums["abs_diff"].round(2)
    dominant_idx = (
        bucket_sums.groupby(["customer_id", "contract_id", "month"],
                            sort=False)["abs_diff"].idxmax()
    )
    for row in bucket_sums.loc[dominant_idx].itertuples(index=False):
        category_map[(str(row.customer_id), str(row.contract_id))][row.month] = int(row.code)

    return category_map


def augment_category_map_from_ledger(category_map: dict, ledger_diff_df: pd.DataFrame, all_months: list) -> dict:
    """
    For (customer_id, contract_id) pairs where the ledger shows a MoM change
    but the contracts-based category is 0, assign 2 (increase) or -2 (decrease).
    """
    if ledger_diff_df is None or ledger_diff_df.empty:
        return category_map

    all_months_set = set(all_months)
    relevant_months = [c for c in ledger_diff_df.columns
                       if "-" in str(c) and c in all_months_set]
    if not relevant_months:
        return category_map

    # One groupby+sum replaces the per-(cust,contract)×per-month .sum() loop.
    ct_sums = ledger_diff_df.groupby(["customer_id", "contract_id"], sort=False)[relevant_months].sum()

    for (cid, ctid), row in ct_sums.iterrows():
        key = (str(cid), str(ctid))
        if key not in category_map:
            category_map[key] = {}
        for month in relevant_months:
            if category_map[key].get(month, 0) != 0:
                continue
            ledger_sum = row[month]
            if abs(ledger_sum) < 0.01:
                continue
            category_map[key][month] = 2 if ledger_sum > 0 else -2

    return category_map

# ============================================================
# MODEL BUILDER
# ============================================================

CATEGORY_ROW_BG    = QColor(40, 50, 65)      # Dark blue-grey for category rows
CATEGORY_ROW_FG    = QColor(180, 210, 255)   # Light blue text for category codes
CATEGORY_ZERO_FG   = QColor(90, 100, 110)    # Muted grey for zero / no movement
CATEGORY_MANUAL_FG = QColor(255, 200, 100)   # Amber for manually overridden categories

SUMMARY_POS_FG   = QColor(100, 220, 140)   # Green for positive movements
SUMMARY_NEG_FG   = QColor(220, 100, 100)   # Red for negative movements
SUMMARY_TOTAL_BG = QColor(50, 58, 70)      # Slightly lifted bg for the total row
ASSIGNED_REF_FG  = QColor(100, 200, 190)   # Teal for manually-assigned (reassignable) refs


def _bold(item: QStandardItem) -> QStandardItem:
    """Set an item's font to bold. Returns the item for chaining."""
    f = item.font()
    f.setBold(True)
    item.setFont(f)
    return item

def _no_edit(item: QStandardItem) -> QStandardItem:
    """Make an item non-editable. Returns the item for chaining."""
    item.setFlags(item.flags() & ~Qt.ItemIsEditable)
    return item


def _amount_item(val: float, *, fmt: str = ",.0f", threshold: float = 0.01,
                 bold: bool = False, total_bg: bool = False) -> QStandardItem:
    """
    Build a money cell for QStandardItemModel: right-aligned, sign-coloured
    (green / red), blank for near-zero values, optionally bold and / or with
    the total-row background. Used across the drill-down dialogs.
    """
    text = format(val, fmt) if abs(val) >= threshold else ""
    item = _no_edit(QStandardItem(text))
    item.setTextAlignment(Qt.AlignRight | Qt.AlignVCenter)
    if bold:
        _bold(item)
    if total_bg:
        item.setBackground(QBrush(SUMMARY_TOTAL_BG))
    if val > threshold:
        item.setForeground(QBrush(SUMMARY_POS_FG))
    elif val < -threshold:
        item.setForeground(QBrush(SUMMARY_NEG_FG))
    return item


def _aggregate_summary_by_category(category_map: dict, contract_totals,
                                   month_cols: list) -> dict:
    """
    Sum per-contract MoM amounts into {category_code: {month: total}}.

    *contract_totals* is a DataFrame (or Series with multi-index) indexed by
    (customer_id, contract_id) with month columns. Used by both the in-app
    MRR Summary model and the Excel export.
    """
    cat_amounts: dict[int, dict[str, float]] = {}
    for (cid, ctid), row in contract_totals.iterrows():
        key = (str(cid), str(ctid))
        cat_months = category_map.get(key, {})
        for month in month_cols:
            code = cat_months.get(month, 0)
            if code == 0:
                continue
            val = float(row[month])
            if abs(val) < 0.01:
                continue
            cat_amounts.setdefault(code, {})
            cat_amounts[code][month] = cat_amounts[code].get(month, 0.0) + val
    return cat_amounts


# Backlog tree column → info-dict key. Column 0 holds the label.
_BACKLOG_COL_KEYS = {1: "tcv", 2: "billed", 3: "remaining"}

# ============================================================
# CUSTOM TREE MODEL (replaces QStandardItemModel for the large views)
# ============================================================
#
# Avoids allocating a QStandardItem per cell — the ledger view alone has
# ~25k rows × ~13 visible columns, which previously cost ~1.6s in Qt object
# allocation. The custom model holds a lightweight Python Node tree and
# computes per-cell role data on demand from cached per-row info.

# Pre-built QBrush instances. QBrush is cheap to construct but the views ask
# for the brush thousands of times — caching avoids per-call allocations.
_BR_CATEGORY_ROW_BG    = QBrush(CATEGORY_ROW_BG)
_BR_CATEGORY_ROW_FG    = QBrush(CATEGORY_ROW_FG)
_BR_CATEGORY_ZERO_FG   = QBrush(CATEGORY_ZERO_FG)
_BR_CATEGORY_MANUAL_FG = QBrush(CATEGORY_MANUAL_FG)
_BR_SUMMARY_POS_FG     = QBrush(SUMMARY_POS_FG)
_BR_SUMMARY_NEG_FG     = QBrush(SUMMARY_NEG_FG)
_BR_SUMMARY_TOTAL_BG   = QBrush(SUMMARY_TOTAL_BG)
_BR_ASSIGNED_REF_FG    = QBrush(ASSIGNED_REF_FG)
_BR_UNASSIGNED_REF_FG  = QBrush(QColor(220, 180, 80))

_FONT_CACHE = {}
_FONT_SIZE_OFFSET = 2  # Increase default font size by this many points

def _qfont(bold=False, italic=False):
    """Return a cached QFont with the given attributes."""
    key = (bold, italic)
    f = _FONT_CACHE.get(key)
    if f is None:
        f = QFont()
        if bold:   f.setBold(True)
        if italic: f.setItalic(True)
        f.setPointSize(f.pointSize() + _FONT_SIZE_OFFSET)
        _FONT_CACHE[key] = f
    return f


class Node:
    """Tree-model node. Holds a kind tag + free-form info dict for role lookups."""
    __slots__ = ("parent", "children", "kind", "info", "row_in_parent")

    def __init__(self, parent, kind, info):
        self.parent = parent
        self.kind = kind
        self.info = info
        self.children = []
        self.row_in_parent = 0
        if parent is not None:
            self.row_in_parent = len(parent.children)
            parent.children.append(self)


class _MRRModel(QAbstractItemModel):
    """
    Tree model backed by a Node tree. data() dispatches by node.kind.
    Used by all three views — both the hierarchical contracts/ledger views
    and the flat MRR summary (whose rows are direct children of root).
    """
    def __init__(self, root, headers, visible_months, manual_keys=None, parent=None):
        super().__init__(parent)
        self._root = root
        self._headers = headers
        self._visible_months = visible_months
        self._manual_keys = manual_keys or set()

    # ── Tree navigation ────────────────────────────────────────────

    def rowCount(self, parent=QModelIndex()):
        if not parent.isValid():
            return len(self._root.children)
        return len(parent.internalPointer().children)

    def columnCount(self, parent=QModelIndex()):
        return len(self._headers)

    def headerData(self, section, orientation, role=Qt.DisplayRole):
        if role == Qt.DisplayRole and orientation == Qt.Horizontal:
            return self._headers[section]
        return None

    def index(self, row, column, parent=QModelIndex()):
        if not self.hasIndex(row, column, parent):
            return QModelIndex()
        parent_node = self._root if not parent.isValid() else parent.internalPointer()
        return self.createIndex(row, column, parent_node.children[row])

    def parent(self, index):
        if not index.isValid():
            return QModelIndex()
        parent_node = index.internalPointer().parent
        if parent_node is None or parent_node is self._root:
            return QModelIndex()
        return self.createIndex(parent_node.row_in_parent, 0, parent_node)

    def flags(self, index):
        if not index.isValid():
            return Qt.NoItemFlags
        return Qt.ItemIsSelectable | Qt.ItemIsEnabled  # never editable

    # ── data() dispatcher ──────────────────────────────────────────

    def data(self, index, role=Qt.DisplayRole):
        if not index.isValid():
            return None
        node = index.internalPointer()
        col = index.column()
        kind = node.kind
        if kind == "customer":      return self._customer_data(node.info, col, role)
        if kind == "contract":      return self._contract_data(node.info, col, role)
        if kind == "category":      return self._category_data(node.info, col, role)
        if kind == "line":          return self._line_data(node.info, col, role)
        if kind == "total":         return self._total_data(node.info, col, role)
        if kind == "summary_row":   return self._summary_row_data(node.info, col, role)
        if kind == "summary_total": return self._summary_total_data(node.info, col, role)
        if kind == "backlog_customer": return self._backlog_customer_data(node.info, col, role)
        if kind == "backlog_contract": return self._backlog_contract_data(node.info, col, role)
        if kind == "backlog_month":    return self._backlog_month_data(node.info, col, role)
        if kind == "backlog_total":    return self._backlog_total_data(node.info, col, role)
        return None

    # ── Per-kind data lookups ──────────────────────────────────────

    def _customer_data(self, d, col, role):
        if role == Qt.BackgroundRole:
            return d.get("brush")
        if col == 0:
            if role == Qt.DisplayRole:
                return f"{d['cid']} - {d['cname']}"
            return None
        if role == Qt.DisplayRole:
            v = d["month_totals"].get(self._visible_months[col - 1], 0)
            return f"{v:,.2f}" if v != 0 else ""
        return None

    def _contract_data(self, d, col, role):
        if role == Qt.BackgroundRole:
            return d.get("brush")
        if col == 0:
            if role == Qt.DisplayRole:
                return f"{d['ctid']} - {d['ctname']}"
            return None
        if role == Qt.DisplayRole:
            v = d["month_totals"].get(self._visible_months[col - 1], 0)
            return f"{v:,.2f}" if v != 0 else ""
        return None

    def _category_data(self, d, col, role):
        if col == 0:
            if role == Qt.DisplayRole:    return "Category"
            if role == Qt.BackgroundRole: return _BR_CATEGORY_ROW_BG
            if role == Qt.ForegroundRole: return _BR_CATEGORY_ROW_FG
            if role == Qt.FontRole:       return _qfont(bold=True, italic=True)
            return None
        month = self._visible_months[col - 1]
        code = d["codes"].get(month, 0)
        if role == Qt.DisplayRole:
            return str(code) if code != 0 else "–"
        if role == Qt.BackgroundRole:
            return _BR_CATEGORY_ROW_BG
        if role == Qt.ForegroundRole:
            if code == 0:
                return _BR_CATEGORY_ZERO_FG
            is_manual = (d["cid"], d["ctid"], month) in self._manual_keys
            return _BR_CATEGORY_MANUAL_FG if is_manual else _BR_CATEGORY_ROW_FG
        if role == Qt.FontRole and code != 0:
            return _qfont(bold=True)
        if role == Qt.TextAlignmentRole:
            return Qt.AlignCenter
        # Custom roles consumed by on_category_clicked:
        if role == ROLE_IS_CATEGORY_CELL: return True
        if role == ROLE_CUSTOMER_ID:      return d["cid"]
        if role == ROLE_CONTRACT_ID:      return d["ctid"]
        if role == ROLE_MONTH:            return month
        return None

    def _line_data(self, d, col, role):
        if col == 0:
            if role == Qt.DisplayRole: return d["label"]
            if role == Qt.ToolTipRole: return d.get("tooltip")
            if role == Qt.ForegroundRole:
                return d.get("label_fg")
            if role == Qt.FontRole and d.get("italic"):
                return _qfont(italic=True)
        else:
            if role == Qt.DisplayRole:
                v = d["month_values"].get(self._visible_months[col - 1], 0)
                return f"{v:,.2f}" if v != 0 else ""
        # Stamp the interactive-row roles on EVERY column so any clicked cell works.
        if d.get("is_interactive"):
            if role == ROLE_IS_UNASSIGNED_REF:    return True
            if role == ROLE_CUSTOMER_ID:          return d["cid"]
            if role == ROLE_REF_NO:               return d["ref_text"]
            if role == ROLE_ASSIGNED_CONTRACT_ID: return d.get("assigned_ctid")
            if role == ROLE_CAN_CLEAR:            return d.get("can_clear", False)
        return None

    def _total_data(self, d, col, role):
        if role == Qt.BackgroundRole: return _BR_SUMMARY_TOTAL_BG
        if role == Qt.FontRole:       return _qfont(bold=True)
        if col == 0:
            if role == Qt.DisplayRole: return "Total"
            return None
        if role == Qt.DisplayRole:
            v = d["month_sums"].get(self._visible_months[col - 1], 0)
            return f"{v:,.2f}" if v != 0 else ""
        return None

    def _summary_row_data(self, d, col, role):
        if col == 0:
            if role == Qt.DisplayRole: return d["label"]
            return None
        month = self._visible_months[col - 1]
        val = d["amounts"].get(month, 0.0)
        if role == Qt.DisplayRole:
            return f"{val:,.0f}" if abs(val) >= 0.01 else ""
        if role == Qt.TextAlignmentRole:
            return Qt.AlignRight | Qt.AlignVCenter
        if role == Qt.ForegroundRole:
            if val > 0.01:  return _BR_SUMMARY_POS_FG
            if val < -0.01: return _BR_SUMMARY_NEG_FG
        if role == ROLE_CATEGORY_CODE: return d["code"]
        if role == ROLE_MONTH:         return month
        return None

    def _summary_total_data(self, d, col, role):
        if role == Qt.BackgroundRole: return _BR_SUMMARY_TOTAL_BG
        if role == Qt.FontRole:       return _qfont(bold=True)
        if col == 0:
            if role == Qt.DisplayRole: return "Net MRR Change"
            return None
        month = self._visible_months[col - 1]
        val = d["totals_by_month"].get(month, 0.0)
        if role == Qt.DisplayRole:
            return f"{val:,.0f}" if abs(val) >= 0.01 else ""
        if role == Qt.TextAlignmentRole:
            return Qt.AlignRight | Qt.AlignVCenter
        if role == Qt.ForegroundRole:
            if val > 0.01:  return _BR_SUMMARY_POS_FG
            if val < -0.01: return _BR_SUMMARY_NEG_FG
        return None

    # ── Backlog view (Customer → Contract → Month) ─────────────────
    # Columns: 0=Name, 1=TCV, 2=Invoice Amount, 3=Remaining Balance

    def _backlog_customer_data(self, d, col, role):
        if col == 0:
            if role == Qt.DisplayRole:
                return f"{d['cid']} - {d['cname']}"
            return None
        key = _BACKLOG_COL_KEYS.get(col)
        if key is None:
            return None
        if role == Qt.DisplayRole:
            v = d.get(key, 0.0)
            return f"{v:,.2f}" if abs(v) >= 0.01 else ""
        if role == Qt.TextAlignmentRole:
            return Qt.AlignLeft | Qt.AlignVCenter
        return None

    def _backlog_contract_data(self, d, col, role):
        if col == 0:
            if role == Qt.DisplayRole:
                return f"{d['ctid']} - {d['ctname']}"
            return None
        key = _BACKLOG_COL_KEYS.get(col)
        if key is None:
            return None
        if role == Qt.DisplayRole:
            v = d.get(key, 0.0)
            return f"{v:,.2f}" if abs(v) >= 0.01 else ""
        if role == Qt.TextAlignmentRole:
            return Qt.AlignLeft | Qt.AlignVCenter
        return None

    def _backlog_total_data(self, d, col, role):
        if role == Qt.BackgroundRole: return _BR_SUMMARY_TOTAL_BG
        if role == Qt.FontRole:       return _qfont(bold=True)
        if col == 0:
            if role == Qt.DisplayRole: return "Total"
            if role == Qt.TextAlignmentRole:
                return Qt.AlignLeft | Qt.AlignVCenter
            return None
        key = _BACKLOG_COL_KEYS.get(col)
        if key is None:
            return None
        if role == Qt.DisplayRole:
            v = d.get(key, 0.0)
            return f"{v:,.2f}" if abs(v) >= 0.01 else ""
        if role == Qt.TextAlignmentRole:
            return Qt.AlignLeft | Qt.AlignVCenter
        return None

    def _backlog_month_data(self, d, col, role):
        if col == 0:
            if role == Qt.DisplayRole:
                return d.get("month", "")
            if role == Qt.TextAlignmentRole:
                return Qt.AlignLeft | Qt.AlignVCenter
            return None
        if col == 2:
            if role == Qt.DisplayRole:
                v = d.get("amount", 0.0)
                return f"{v:,.2f}" if abs(v) >= 0.01 else ""
            if role == Qt.TextAlignmentRole:
                return Qt.AlignLeft | Qt.AlignVCenter
        if col == 3:
            if role == Qt.DisplayRole:
                v = d.get("remaining", 0.0)
                return f"{v:,.2f}"
            if role == Qt.TextAlignmentRole:
                return Qt.AlignLeft | Qt.AlignVCenter
        return None


def build_qt_model(pivot_df, selected_highlight_month, visible_months,
                   is_ledger=False, is_diff_view=False,
                   reference_totals=None, category_map=None, manual_keys=None,
                   assigned_refs=None):
    root = Node(None, "root", {})
    headers = ["Name"] + list(visible_months)
    if pivot_df.empty:
        return _MRRModel(root, headers, visible_months, manual_keys)

    cid_col, cname_col   = pivot_df.columns[0], pivot_df.columns[1]
    ctid_col, ctname_col = pivot_df.columns[2], pivot_df.columns[3]
    ref_col              = pivot_df.columns[4]

    pivot_df = pivot_df.sort_values([cname_col, ctname_col], kind="stable")

    for (c_id, c_name), c_group in pivot_df.groupby([cid_col, cname_col], sort=False):
        if is_diff_view:
            hv = c_group[selected_highlight_month].sum() if selected_highlight_month in c_group.columns else 0
            if abs(hv) < 0.01:
                continue

        color = None
        if is_ledger and reference_totals:
            color = reference_totals.get(c_id) or reference_totals.get(c_name)
        brush = QBrush(color) if color else None

        c_totals = {m: float(c_group[m].sum()) if m in c_group.columns else 0
                    for m in visible_months}
        customer_node = Node(root, "customer", {
            "cid": c_id, "cname": c_name,
            "month_totals": c_totals, "brush": brush,
        })

        for (ct_id, ct_name), sub_group in c_group.groupby([ctid_col, ctname_col], sort=False):
            if is_diff_view:
                sh = sub_group[selected_highlight_month].sum() if selected_highlight_month in sub_group.columns else 0
                if abs(sh) < 0.01:
                    continue

            # Category row (ledger view only) is a sibling of the contract row.
            if is_ledger and category_map is not None:
                key = (str(c_id), str(ct_id))
                Node(customer_node, "category", {
                    "cid": str(c_id), "ctid": str(ct_id),
                    "codes": category_map.get(key, {}),
                })

            ct_totals = {m: float(sub_group[m].sum()) if m in sub_group.columns else 0
                         for m in visible_months}
            contract_node = Node(customer_node, "contract", {
                "ctid": ct_id, "ctname": ct_name,
                "month_totals": ct_totals, "brush": brush,
            })

            is_interactive = is_ledger and str(c_id) not in MISSING_SENTINELS
            is_unassigned_ref_row = is_interactive and str(ct_id) == "Uncategorized"

            for _, row in sub_group.iterrows():
                if is_diff_view:
                    lh = row[selected_highlight_month] if selected_highlight_month in row.index else 0
                    if abs(lh) < 0.01:
                        continue
                label    = f"Ref {row[ref_col]}" if is_ledger else f"Line {row[ref_col]}"
                ref_text = str(row[ref_col])
                has_assignment = (
                    is_interactive
                    and assigned_refs is not None
                    and (ref_text, str(c_id)) in assigned_refs
                )

                if is_unassigned_ref_row:
                    label_fg, italic = _BR_UNASSIGNED_REF_FG, True
                    tooltip = "No contract assigned. Right-click or double-click to assign."
                elif has_assignment:
                    label_fg, italic = _BR_ASSIGNED_REF_FG, True
                    tooltip = "Manually assigned. Right-click or double-click to change or clear."
                elif is_interactive:
                    label_fg, italic = None, False
                    tooltip = "Right-click or double-click to move to a different contract."
                else:
                    label_fg, italic, tooltip = None, False, None

                line_values = {m: float(row[m]) if m in row.index else 0
                               for m in visible_months}
                Node(contract_node, "line", {
                    "label": label,
                    "ref_text": ref_text,
                    "month_values": line_values,
                    "cid": str(c_id),
                    "ctid": str(ct_id),
                    "is_interactive": is_interactive,
                    "assigned_ctid": (str(ct_id) if (is_interactive and not is_unassigned_ref_row) else None),
                    "can_clear": has_assignment,
                    "label_fg": label_fg,
                    "italic": italic,
                    "tooltip": tooltip,
                })

    total_sums = {m: float(pivot_df[m].sum()) if m in pivot_df.columns else 0
                  for m in visible_months}
    Node(root, "total", {"month_sums": total_sums})

    return _MRRModel(root, headers, visible_months, manual_keys)


def build_category_summary_model(category_map, p_l_diff, visible_months):
    root = Node(None, "root", {})
    headers = ["Category"] + list(visible_months)

    if p_l_diff is None or p_l_diff.empty or not visible_months:
        return _MRRModel(root, headers, visible_months)

    cid_col  = p_l_diff.columns[0]
    ctid_col = p_l_diff.columns[2]
    month_cols = [m for m in visible_months if m in p_l_diff.columns]

    contract_totals = p_l_diff.groupby([cid_col, ctid_col])[month_cols].sum()
    cat_amounts = _aggregate_summary_by_category(category_map, contract_totals, month_cols)

    if not cat_amounts:
        return _MRRModel(root, headers, visible_months)

    totals_by_month: dict[str, float] = {}
    for code, label in ALL_CATEGORIES:
        if code not in cat_amounts:
            continue
        amounts = cat_amounts[code]
        Node(root, "summary_row", {
            "code": code,
            "label": f"{code}  —  {label}",
            "amounts": amounts,
        })
        for m in month_cols:
            totals_by_month[m] = totals_by_month.get(m, 0.0) + amounts.get(m, 0.0)

    Node(root, "summary_total", {"totals_by_month": totals_by_month})
    return _MRRModel(root, headers, visible_months)


def build_backlog_model(backlog_df, has_remaining_keys=None, outstanding_only=False):
    """
    Build the Backlog tree model: Customer → Contract → Month.

    For each contract, TCV is the sum of all PS-line amounts on it. Months
    under the contract are listed chronologically; the Remaining Balance
    column shows TCV decreasing by each month's invoice amount until 0.

    When *outstanding_only* is True, hide any contract whose key is not in
    *has_remaining_keys* — i.e. contracts with no scheduled_posting_date
    in the future. Customers left with no visible contracts are dropped.
    """
    headers = ["Name", "TCV", "Billed to Date", "Remaining Balance"]
    root = Node(None, "root", {})

    if backlog_df is None or backlog_df.empty:
        return _MRRModel(root, headers, [])

    df = backlog_df.sort_values(
        ["customer_name", "contract_name", "month"], kind="stable"
    )

    keys = has_remaining_keys or set()

    total_tcv = 0.0
    total_billed = 0.0
    total_remaining = 0.0

    for (c_id, c_name), c_group in df.groupby(
        ["customer_id", "customer_name"], sort=False
    ):
        # Build contracts first so we can skip empty customers when filtering.
        contracts = []
        for (ct_id, ct_name), ct_group in c_group.groupby(
            ["contract_id", "contract_name"], sort=False
        ):
            if outstanding_only and (str(c_id), str(ct_id)) not in keys:
                continue
            tcv = float(ct_group["amount"].sum())
            billed_to_date = float(ct_group["billed"].sum())
            remaining = tcv - billed_to_date
            month_rows = ct_group.groupby("month", sort=True).agg(
                amount=("amount", "sum"), billed=("billed", "sum")
            )
            contracts.append((ct_id, ct_name, tcv, billed_to_date, remaining, month_rows))

        if not contracts:
            continue

        cust_tcv       = sum(t for _, _, t, _, _, _ in contracts)
        cust_billed    = sum(b for _, _, _, b, _, _ in contracts)
        cust_remaining = sum(r for _, _, _, _, r, _ in contracts)
        total_tcv       += cust_tcv
        total_billed    += cust_billed
        total_remaining += cust_remaining
        cust_node = Node(root, "backlog_customer", {
            "cid": c_id, "cname": c_name,
            "tcv": cust_tcv,
            "billed": cust_billed,
            "remaining": cust_remaining,
        })

        for ct_id, ct_name, tcv, billed_to_date, remaining, month_rows in contracts:
            contract_node = Node(cust_node, "backlog_contract", {
                "ctid": ct_id, "ctname": ct_name,
                "tcv": tcv,
                "billed": billed_to_date,
                "remaining": remaining,
            })
            running = tcv
            for month, row in month_rows.iterrows():
                billed_m = float(row["billed"])
                running -= billed_m
                Node(contract_node, "backlog_month", {
                    "month": str(month),
                    "amount": billed_m,
                    "remaining": running,
                })

    if root.children:
        Node(root, "backlog_total", {
            "tcv": total_tcv,
            "billed": total_billed,
            "remaining": total_remaining,
        })

    return _MRRModel(root, headers, [])


# ============================================================
# DRILL-DOWN DIALOG  (double-click a summary cell)
# ============================================================

class DrillDownDialog(QDialog):
    """
    First-level drill-down from an MRR Summary cell: shows the list of
    customers / contracts contributing to that category × month, with monthly
    diff values across the visible range. Double-clicking a number cell opens
    a second-level dialog (TransactionBreakdownDialog) with the underlying
    ledger postings for that specific (customer, contract, month).
    """
    def __init__(self, category_code: int, clicked_month: str,
                 category_map: dict, p_l_diff, visible_months: list,
                 df_ledger, parent=None):
        super().__init__(parent)
        label = next((lbl for c, lbl in ALL_CATEGORIES if c == category_code), str(category_code))
        self.setWindowTitle(f"{category_code} — {label}  ·  {clicked_month}")
        self.resize(1000, 420)

        self._df_ledger = df_ledger

        layout = QVBoxLayout(self)
        layout.addWidget(QLabel(
            f"Customers in <b>{category_code} — {label}</b> for <b>{clicked_month}</b>"
            f", amounts shown for all selected months. "
            f"Double-click a number to see the underlying ledger transactions."
        ))

        cid_col   = p_l_diff.columns[0]
        cname_col = p_l_diff.columns[1]
        ctid_col  = p_l_diff.columns[2]
        ctname_col = p_l_diff.columns[3]
        month_cols = [m for m in visible_months if m in p_l_diff.columns]

        contract_totals = p_l_diff.groupby([cid_col, ctid_col])[month_cols].sum()
        name_df = p_l_diff.groupby([cid_col, ctid_col])[[cname_col, ctname_col]].first()

        model = QStandardItemModel()
        model.setHorizontalHeaderLabels(["Customer / Contract"] + visible_months)

        totals: dict[str, float] = {}
        any_rows = False

        # Collect matching rows, resolve names, then sort by customer name
        matching = []
        for (cid, ctid), row in contract_totals.iterrows():
            key = (str(cid), str(ctid))
            if category_map.get(key, {}).get(clicked_month) != category_code:
                continue
            val_in_month = float(row[clicked_month]) if clicked_month in row.index else 0.0
            if abs(val_in_month) < 0.01:
                continue
            try:
                c_name  = name_df.loc[(cid, ctid), cname_col]
                ct_name = name_df.loc[(cid, ctid), ctname_col]
            except KeyError:
                c_name, ct_name = str(cid), str(ctid)
            matching.append((str(c_name), str(ct_name), row, str(cid), str(ctid)))

        matching.sort(key=lambda x: (x[0].lower(), x[1].lower()))

        for c_name, ct_name, row, cid, ctid in matching:
            any_rows = True
            tooltip = f"{c_name}  /  {ct_name}"
            row_items = [_no_edit(QStandardItem(tooltip))]

            for m in visible_months:
                val = float(row[m]) if m in row.index else 0.0
                cell = _amount_item(val)
                # Tag the cell so double-click can drill into it.
                cell.setData(cid, ROLE_CUSTOMER_ID)
                cell.setData(ctid, ROLE_CONTRACT_ID)
                cell.setData(m,    ROLE_MONTH)
                cell.setData(tooltip, Qt.ToolTipRole)
                row_items.append(cell)
                if m in month_cols:
                    totals[m] = totals.get(m, 0.0) + val

            model.appendRow(row_items)

        if not any_rows:
            model.appendRow([QStandardItem("No contributing customers found.")])

        # ── Total row ─────────────────────────────────────────────────────
        total_label = _no_edit(_bold(QStandardItem("Total")))
        total_label.setBackground(QBrush(SUMMARY_TOTAL_BG))
        total_row = [total_label] + [
            _amount_item(totals.get(m, 0.0), bold=True, total_bg=True)
            for m in visible_months
        ]
        model.appendRow(total_row)

        tree = QTreeView()
        tree.setRootIsDecorated(False)
        tree.setModel(model)
        tree.resizeColumnToContents(0)
        tree.doubleClicked.connect(self._on_cell_double_clicked)
        layout.addWidget(tree)

        btn = QPushButton("Close")
        btn.clicked.connect(self.accept)
        layout.addWidget(btn)

    def _on_cell_double_clicked(self, index):
        if not index.isValid() or index.column() == 0:
            return
        cid   = index.data(ROLE_CUSTOMER_ID)
        ctid  = index.data(ROLE_CONTRACT_ID)
        month = index.data(ROLE_MONTH)
        if cid is None or ctid is None or month is None:
            return
        TransactionBreakdownDialog(cid, ctid, month, self._df_ledger,
                                   parent=self).exec()


class TransactionBreakdownDialog(QDialog):
    """
    Second-level drill-down: shows the raw ledger postings for a single
    (customer, contract, month) cell from DrillDownDialog. Includes the prior
    month so the MoM diff is visible (clicked - prior = the cell value).
    """
    def __init__(self, customer_id: str, contract_id: str, clicked_month: str,
                 df_ledger, parent=None):
        super().__init__(parent)
        try:
            prior_month = str(pd.Period(clicked_month, freq='M') - 1)
        except Exception:
            prior_month = ""

        ref_col   = 'reference_no.' if 'reference_no.' in df_ledger.columns else 'reference_no'
        title_col = 'account_title' if 'account_title' in df_ledger.columns else None

        # Filter to the specific (customer, contract) for the two months.
        df = df_ledger.copy()
        df['_base_account'] = df['account'].astype(str).str.split('-').str[0]
        df = df[df['_base_account'].isin(_LEDGER_TAB_MS_ACCOUNTS)]
        df = df[df['customer_id'].astype(str) == str(customer_id)]
        df = df[df['contract_id'].astype(str) == str(contract_id)]
        df['_month']  = pd.to_datetime(df['posting_date'], errors='coerce').dt.to_period('M').astype(str)
        df['_amount'] = pd.to_numeric(df['amount'], errors='coerce').fillna(0) * (-1)
        df = df[df['_month'].isin([clicked_month, prior_month])]

        # Resolve display names from the filtered slice (or fall back to ids).
        if not df.empty:
            customer_name = str(df['customer_name'].iloc[0])
            contract_name = str(df['contract_name'].iloc[0])
        else:
            customer_name, contract_name = str(customer_id), str(contract_id)

        self.setWindowTitle(f"{customer_name} / {contract_name}  ·  {clicked_month}")
        self.resize(1000, 460)

        clicked_total = df.loc[df['_month'] == clicked_month, '_amount'].sum()
        prior_total   = df.loc[df['_month'] == prior_month,   '_amount'].sum()
        diff = clicked_total - prior_total

        layout = QVBoxLayout(self)
        layout.addWidget(QLabel(
            f"<b>{customer_name}</b> / <b>{contract_name}</b><br>"
            f"{prior_month}: {prior_total:,.0f}  →  "
            f"{clicked_month}: {clicked_total:,.0f}  "
            f"(MoM diff: <b>{diff:,.0f}</b>)"
        ))

        headers = ["Posting Date", "Month", "Reference", "Account"]
        if title_col is not None:
            headers.append("Description")
        headers.append("Amount")

        model = QStandardItemModel()
        model.setHorizontalHeaderLabels(headers)

        if df.empty:
            model.appendRow([QStandardItem("No ledger postings found for this cell.")])
        else:
            df = df.sort_values(['posting_date'])
            for _, tx in df.iterrows():
                row = [
                    _no_edit(QStandardItem(str(tx['posting_date']))),
                    _no_edit(QStandardItem(str(tx['_month']))),
                    _no_edit(QStandardItem(str(tx[ref_col]))),
                    _no_edit(QStandardItem(str(tx['account']))),
                ]
                if title_col is not None:
                    row.append(_no_edit(QStandardItem(str(tx[title_col]))))
                row.append(_amount_item(float(tx['_amount']),
                                        fmt=",.2f", threshold=0.005))
                model.appendRow(row)

            # Per-month total rows + diff total row.
            def _total_row(label_text, val):
                lbl = _no_edit(_bold(QStandardItem(label_text)))
                lbl.setBackground(QBrush(SUMMARY_TOTAL_BG))
                blanks = []
                for _ in range(len(headers) - 2):
                    b = _no_edit(QStandardItem(""))
                    b.setBackground(QBrush(SUMMARY_TOTAL_BG))
                    blanks.append(b)
                amt = _amount_item(val, fmt=",.2f", threshold=0.005,
                                   bold=True, total_bg=True)
                return [lbl, *blanks, amt]

            model.appendRow(_total_row(f"Total {prior_month}",   prior_total))
            model.appendRow(_total_row(f"Total {clicked_month}", clicked_total))
            model.appendRow(_total_row("MoM diff",               diff))

        tree = QTreeView()
        tree.setRootIsDecorated(False)
        tree.setModel(model)
        for col in range(model.columnCount()):
            tree.resizeColumnToContents(col)
        layout.addWidget(tree)

        btn = QPushButton("Close")
        btn.clicked.connect(self.accept)
        layout.addWidget(btn)


# ============================================================
# ASSIGN UNASSIGNED LEDGER TRANSACTION TO A CONTRACT
# ============================================================

class AssignContractDialog(QDialog):
    def __init__(self, customer_id: str, reference_no: str,
                 current_contract_id: str = None, can_clear: bool = False, parent=None):
        super().__init__(parent)
        self.selected_contract  = None   # (contract_id, contract_name) on accept
        self.clear_requested    = False
        self._current_contract_id = current_contract_id
        self._customer_contracts  = self._load_contracts(customer_id)
        self._all_contracts       = self._load_all_contracts()

        is_reassign = current_contract_id is not None
        self.setWindowTitle("Change Contract Assignment" if is_reassign else "Assign Transaction to Contract")
        self.setMinimumWidth(560)
        self.setMinimumHeight(400)

        layout = QVBoxLayout(self)
        action_word = "Reassigning" if is_reassign else "Assigning"
        layout.addWidget(QLabel(
            f"{action_word} reference <b>{reference_no}</b> "
            f"(customer <b>{customer_id}</b>) to a contract:"
        ))

        self.list_widget = QListWidget()
        layout.addWidget(self.list_widget)

        self.chk_all = QCheckBox("Show contracts from all customers")
        self.chk_all.toggled.connect(self._populate_list)
        layout.addWidget(self.chk_all)

        btn_box    = QDialogButtonBox()
        self._btn_ok = btn_box.addButton("Assign", QDialogButtonBox.AcceptRole)
        btn_cancel = btn_box.addButton("Cancel",  QDialogButtonBox.RejectRole)
        self._btn_ok.clicked.connect(self._on_ok)
        btn_cancel.clicked.connect(self.reject)
        if can_clear:
            btn_clear = btn_box.addButton("Clear Assignment", QDialogButtonBox.ResetRole)
            btn_clear.clicked.connect(self._on_clear)
        self.list_widget.itemDoubleClicked.connect(lambda _: self._on_ok())
        layout.addWidget(btn_box)

        # Populate with same-customer contracts initially
        self._populate_list(show_all=False)

    def _populate_list(self, show_all: bool = False):
        self.list_widget.clear()
        if show_all:
            # Each entry: (customer_name, contract_id, contract_name)
            for cust_name, cid, cname in self._all_contracts:
                self.list_widget.addItem(f"{cust_name}   /   {cid}   —   {cname}")
            data = [(cid, cname) for _, cid, cname in self._all_contracts]
        else:
            for cid, cname in self._customer_contracts:
                self.list_widget.addItem(f"{cid}   —   {cname}")
            data = self._customer_contracts

        self._active_data = data
        # Pre-select current contract if visible
        preselect = 0
        for i, (cid, _) in enumerate(data):
            if cid == self._current_contract_id:
                preselect = i
                break
        if data:
            self.list_widget.setCurrentRow(preselect)
        self._btn_ok.setEnabled(bool(data))

    def _load_contracts(self, customer_id: str):
        try:
            conn = sqlite3.connect("data_source/mrr_data.db")
            df = pd.read_sql(
                "SELECT DISTINCT contract_id, contract_name FROM contracts "
                "WHERE customer_id = ? ORDER BY contract_name",
                conn, params=(customer_id,)
            )
            conn.close()
        except Exception:
            return []
        result = []
        for _, r in df.iterrows():
            cid   = "" if pd.isna(r["contract_id"])   else str(r["contract_id"]).strip()
            cname = "" if pd.isna(r["contract_name"]) else str(r["contract_name"]).strip()
            if cid and cid.lower() not in ("nan", "none"):
                result.append((cid, cname or "(unnamed)"))
        return result

    def _load_all_contracts(self):
        try:
            conn = sqlite3.connect("data_source/mrr_data.db")
            df = pd.read_sql(
                "SELECT DISTINCT customer_name, contract_id, contract_name FROM contracts "
                "ORDER BY customer_name, contract_name",
                conn
            )
            conn.close()
        except Exception:
            return []
        result = []
        for _, r in df.iterrows():
            cid   = "" if pd.isna(r["contract_id"])   else str(r["contract_id"]).strip()
            cname = "" if pd.isna(r["contract_name"]) else str(r["contract_name"]).strip()
            cust  = "" if pd.isna(r["customer_name"]) else str(r["customer_name"]).strip()
            if cid and cid.lower() not in ("nan", "none"):
                result.append((cust or "(unknown)", cid, cname or "(unnamed)"))
        return result

    def _on_ok(self):
        row = self.list_widget.currentRow()
        if 0 <= row < len(self._active_data):
            self.selected_contract = self._active_data[row]
            self.accept()

    def _on_clear(self):
        self.clear_requested = True
        self.accept()


# ============================================================
# CATEGORY OVERRIDE DIALOG
# ============================================================

class CategoryDialog(QDialog):
    def __init__(self, current_code: int, is_manual: bool = False, parent=None):
        super().__init__(parent)
        self.selected_code = None
        self.clear_requested = False
        self.setWindowTitle("Set Category Override")
        self.setMinimumWidth(380)

        layout = QVBoxLayout(self)

        status = "manual override" if is_manual else "auto"
        layout.addWidget(QLabel(f"Current category: {current_code}  ({status})"))

        self.list_widget = QListWidget()
        for code, label in ALL_CATEGORIES:
            self.list_widget.addItem(f"{code:>4}  —  {label}")
            if code == current_code:
                self.list_widget.setCurrentRow(self.list_widget.count() - 1)
        layout.addWidget(self.list_widget)

        btn_box = QDialogButtonBox()
        btn_ok = btn_box.addButton("OK", QDialogButtonBox.AcceptRole)
        btn_cancel = btn_box.addButton("Cancel", QDialogButtonBox.RejectRole)
        if is_manual:
            btn_clear = btn_box.addButton("Clear Override", QDialogButtonBox.ResetRole)
            btn_clear.clicked.connect(self._on_clear)
        btn_ok.clicked.connect(self._on_ok)
        btn_cancel.clicked.connect(self.reject)
        layout.addWidget(btn_box)

    def _on_ok(self):
        row = self.list_widget.currentRow()
        if row >= 0:
            self.selected_code = ALL_CATEGORIES[row][0]
            self.accept()

    def _on_clear(self):
        self.clear_requested = True
        self.accept()

# ============================================================
# CHECKABLE MULTI-SELECT COMBO BOX (used by the per-tab item/account filters)
# ============================================================

class CheckableComboBox(QComboBox):
    """
    QComboBox with multi-select checkable items.

    A pinned top row labelled "All" toggles every other row at once. The line
    edit shows a compact summary of the current selection. The popup stays open
    while the user toggles individual items; clicking outside (or pressing Esc)
    closes it.
    """

    selectionChanged = Signal()

    def __init__(self, parent=None):
        super().__init__(parent)
        self.setEditable(True)
        self.lineEdit().setReadOnly(True)
        self.lineEdit().installEventFilter(self)
        self._model = QStandardItemModel(self)
        self.setModel(self._model)
        self.view().viewport().installEventFilter(self)
        self._all_items: list[str] = []
        # Suppress hidePopup() once after an item click so toggling stays open.
        self._block_close = False

    # ---- public API -------------------------------------------------------

    def set_items(self, items, default_selected=None, all_label='All'):
        """
        Populate the dropdown with *items*.

        *default_selected* is an optional iterable of item values that should
        start checked. If None, every item starts checked. Items not present
        in *default_selected* (when provided) start unchecked.
        """
        self._all_items = [str(i) for i in items]
        default_set = None if default_selected is None else {str(x) for x in default_selected}
        self._model.clear()
        all_item = QStandardItem(all_label)
        all_item.setFlags(Qt.ItemIsEnabled)
        self._model.appendRow(all_item)
        for it in self._all_items:
            row = QStandardItem(it)
            row.setFlags(Qt.ItemIsEnabled)
            checked = (default_set is None) or (it in default_set)
            row.setData(Qt.Checked if checked else Qt.Unchecked, Qt.CheckStateRole)
            self._model.appendRow(row)
        # Sync the "All" row to reflect the aggregate of the per-item states.
        states = [self._model.item(r).checkState()
                  for r in range(1, self._model.rowCount())]
        if not states or all(s == Qt.Checked for s in states):
            all_item.setCheckState(Qt.Checked)
        elif all(s == Qt.Unchecked for s in states):
            all_item.setCheckState(Qt.Unchecked)
        else:
            all_item.setCheckState(Qt.PartiallyChecked)
        self._refresh_text()

    def selected_items(self):
        m = self._model
        return [m.item(r).text() for r in range(1, m.rowCount())
                if m.item(r).checkState() == Qt.Checked]

    def all_items(self):
        return list(self._all_items)

    # ---- internals --------------------------------------------------------

    def eventFilter(self, obj, event):
        if obj is self.view().viewport():
            if event.type() == QEvent.MouseButtonRelease:
                idx = self.view().indexAt(event.pos())
                if idx.isValid():
                    self._toggle_row(idx.row())
                    self._block_close = True
                    return True
        elif obj is self.lineEdit():
            if event.type() == QEvent.MouseButtonPress:
                self.showPopup()
                return True
        return super().eventFilter(obj, event)

    def hidePopup(self):
        if self._block_close:
            self._block_close = False
            return
        super().hidePopup()

    def _toggle_row(self, row):
        m = self._model
        item = m.item(row)
        if item is None:
            return
        new_state = Qt.Unchecked if item.checkState() == Qt.Checked else Qt.Checked
        if row == 0:  # "All" toggled — apply to every row
            for r in range(m.rowCount()):
                m.item(r).setCheckState(new_state)
        else:
            item.setCheckState(new_state)
            # Sync the "All" row to reflect the aggregate state.
            states = [m.item(r).checkState() for r in range(1, m.rowCount())]
            if all(s == Qt.Checked for s in states):
                m.item(0).setCheckState(Qt.Checked)
            elif all(s == Qt.Unchecked for s in states):
                m.item(0).setCheckState(Qt.Unchecked)
            else:
                m.item(0).setCheckState(Qt.PartiallyChecked)
        self._refresh_text()
        self.selectionChanged.emit()

    def _refresh_text(self):
        sel = self.selected_items()
        n_sel, n_all = len(sel), len(self._all_items)
        if n_all == 0:
            txt = '(none available)'
        elif n_sel == n_all:
            txt = f'All ({n_all})'
        elif n_sel == 0:
            txt = 'None selected'
        elif n_sel <= 3:
            txt = ', '.join(sel)
        else:
            txt = f'{n_sel} of {n_all} selected'
        self.lineEdit().setText(txt)


# ============================================================
# MAIN WINDOW
# ============================================================

class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("MRR Reconciliation Tool")
        self.resize(1500, 800)

        self.p_c = None
        self.p_l = None
        self.p_l_diff = None
        self.all_months = []
        self.category_map = {}                    # (customer_id, contract_id) → {month → code}
        self._contracts_only_category_map = {}    # contracts-only base map (no ledger augmentation)
        self._base_category_map = {}              # programmatic map (pre user overrides)
        self.manual_override_keys = set()         # {(customer_id, contract_id, month)}
        self._raw_df_l = None                     # raw ledger DataFrame (before assignments / pivoting)
        self.ledger_assignments = {}              # {(reference_no, customer_id): (contract_id, contract_name)}

        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        layout = QVBoxLayout(central_widget)

        # ── View navigation (real tabs along the top) ────────────────────
        self.main_tabs = QTabBar()
        self.main_tabs.setDrawBase(False)
        self.main_tabs.setExpanding(False)
        for name in ("Contracts", "Ledger", "MRR Summary", "Backlog"):
            self.main_tabs.addTab(name)

        # ── Filter / action toolbar (below tabs, above content) ──────────
        ctrl_layout = QHBoxLayout()
        self.period_selector = QComboBox()
        self.from_selector = QComboBox()
        self.to_selector = QComboBox()
        self.lbl_status = QLabel("")
        self.lbl_status.setFixedWidth(25)
        self.lbl_status.setStyleSheet("color: red; font-weight: bold; font-size: 16px;")

        self.check_diff = QCheckBox("Show Ledger MoM Differences")
        self.check_diff.toggled.connect(self.update_views)

        self.check_highlight = QCheckBox("Highlight Discrepancies")
        self.check_highlight.setChecked(True)
        self.check_highlight.toggled.connect(self.update_views)

        self.btn_export = QPushButton("Export to Excel")
        self.btn_export.setStyleSheet("background-color: #2c7d32; color: white; font-weight: bold;")
        self.btn_export.clicked.connect(self.export_to_excel)

        ctrl_layout.addWidget(QLabel("Highlight Month:"))
        ctrl_layout.addWidget(self.period_selector)
        ctrl_layout.addSpacing(20)
        ctrl_layout.addWidget(QLabel("Range:"))
        ctrl_layout.addWidget(self.from_selector)
        ctrl_layout.addWidget(QLabel("-"))
        ctrl_layout.addWidget(self.to_selector)
        ctrl_layout.addSpacing(30)
        ctrl_layout.addWidget(self.check_diff)
        ctrl_layout.addWidget(self.check_highlight)
        ctrl_layout.addWidget(self.lbl_status)
        ctrl_layout.addStretch()
        ctrl_layout.addWidget(self.btn_export)
        ctrl_layout.addSpacing(10)

        layout.addWidget(self.main_tabs)
        layout.addLayout(ctrl_layout)

        for cb in [self.period_selector, self.from_selector, self.to_selector]:
            cb.currentTextChanged.connect(self.update_views)

        self.stack = QStackedWidget()
        # Contracts View: one QTreeView per item category, each with its own
        # item-filter dropdown above the tree.
        self.contracts_tabs = QTabWidget()
        self.contract_trees       = {key: QTreeView()          for key in CONTRACT_TAB_KEYS}
        self.contract_item_combos = {key: CheckableComboBox()  for key in CONTRACT_TAB_KEYS}
        for key in CONTRACT_TAB_KEYS:
            self.contracts_tabs.addTab(
                self._build_filtered_tab_page(self.contract_trees[key],
                                              self.contract_item_combos[key],
                                              "Items:"),
                key,
            )
        # Default to MS — it carries the bulk of recurring revenue and matches
        # the prior single-pane contracts view most closely.
        self.contracts_tabs.setCurrentIndex(CONTRACT_TAB_KEYS.index('MS'))
        # Ledger View mirrors the contracts pattern, with per-tab GL-account filters.
        self.ledger_tabs = QTabWidget()
        self.ledger_trees          = {key: QTreeView()         for key in LEDGER_TAB_KEYS}
        self.ledger_account_combos = {key: CheckableComboBox() for key in LEDGER_TAB_KEYS}
        for key in LEDGER_TAB_KEYS:
            tree = self.ledger_trees[key]
            # Belt-and-braces: prevent Qt's default inline editor from opening on
            # double-click. doubleClicked must fire cleanly for ref reassignment.
            tree.setEditTriggers(QAbstractItemView.NoEditTriggers)
            # Right-click anywhere in the ledger view to assign an unassigned ref.
            tree.setContextMenuPolicy(Qt.CustomContextMenu)
            self.ledger_tabs.addTab(
                self._build_filtered_tab_page(tree, self.ledger_account_combos[key],
                                              "Accounts:"),
                key,
            )
        self.ledger_tabs.setCurrentIndex(LEDGER_TAB_KEYS.index('MS'))
        self.tree_summary = QTreeView()
        self.tree_backlog = QTreeView()
        self.tree_summary.setRootIsDecorated(False)

        # Backlog page wraps the tree with its own toolbar.
        self.rb_backlog_general   = QRadioButton("General PS")
        self.rb_backlog_recurring = QRadioButton("Recurring PS (I904)")
        self.rb_backlog_general.setChecked(True)
        _backlog_ps_group = QButtonGroup(self)
        _backlog_ps_group.addButton(self.rb_backlog_general)
        _backlog_ps_group.addButton(self.rb_backlog_recurring)
        self.rb_backlog_general.toggled.connect(self._rebuild_backlog_tree)

        self.chk_backlog_outstanding = QCheckBox("Show only contracts with outstanding amounts")
        self.chk_backlog_outstanding.toggled.connect(self._rebuild_backlog_tree)

        backlog_toolbar = QHBoxLayout()
        backlog_toolbar.addWidget(self.rb_backlog_general)
        backlog_toolbar.addWidget(self.rb_backlog_recurring)
        backlog_toolbar.addSpacing(20)
        backlog_toolbar.addWidget(self.chk_backlog_outstanding)
        backlog_toolbar.addStretch()

        backlog_page = QWidget()
        backlog_layout = QVBoxLayout(backlog_page)
        backlog_layout.setContentsMargins(0, 0, 0, 0)
        backlog_layout.addLayout(backlog_toolbar)
        backlog_layout.addWidget(self.tree_backlog)

        contracts_page = QWidget()
        contracts_layout = QVBoxLayout(contracts_page)
        contracts_layout.setContentsMargins(0, 0, 0, 0)
        contracts_layout.addWidget(self.contracts_tabs)

        ledger_page = QWidget()
        ledger_layout = QVBoxLayout(ledger_page)
        ledger_layout.setContentsMargins(0, 0, 0, 0)
        ledger_layout.addWidget(self.ledger_tabs)

        self.stack.addWidget(contracts_page)
        self.stack.addWidget(ledger_page)
        self.stack.addWidget(self.tree_summary)
        self.stack.addWidget(backlog_page)
        layout.addWidget(self.stack)

        self.main_tabs.currentChanged.connect(self.stack.setCurrentIndex)
        # Start on Contracts to match prior default
        self.main_tabs.setCurrentIndex(0)
        self.stack.setCurrentIndex(0)

        for tree in self.ledger_trees.values():
            tree.clicked.connect(self.on_category_clicked)
            tree.doubleClicked.connect(self.on_ledger_ref_double_clicked)
            tree.customContextMenuRequested.connect(self.on_ledger_context_menu)
        self.tree_summary.doubleClicked.connect(self.on_summary_double_clicked)

        self.initial_load()

    # ── Per-tab item / account filter dropdowns ────────────────────────

    def _populate_filter_combos(self):
        """
        Populate every dropdown with the full universe of codes present in
        the data, with each tab's default selection pre-checked. The user can
        then check items from any other category to include them in the tab.
        """
        # Contracts: every distinct item in the data; tab default = its category list.
        if 'item' in self._raw_df_c.columns:
            items_all = sorted(self._raw_df_c['item'].astype(str).dropna().unique())
        else:
            items_all = []
        items_set = set(items_all)
        contract_defaults = {
            'PS':           items_set & set(_CONTRACT_TAB_PS_ITEMS),
            'MS':           items_set - set(_CONTRACT_TAB_MS_EXCLUDED_ITEMS),
            'Cloud':        items_set & set(_CONTRACT_TAB_CLOUD_ITEMS),
            'Other Resell': items_set & set(_CONTRACT_TAB_OTHER_RESELL_ITEMS),
        }
        for key in CONTRACT_TAB_KEYS:
            self.contract_item_combos[key].set_items(
                items_all, default_selected=contract_defaults[key]
            )

        # Ledger: every distinct base account in the data; tab default = its account list.
        if 'account' in self._df_l_for_pivot.columns:
            bases = self._df_l_for_pivot['account'].astype(str).str.split('-').str[0]
            accounts_all = sorted(bases.dropna().unique())
        else:
            accounts_all = []
        accounts_set = set(accounts_all)
        for key in LEDGER_TAB_KEYS:
            defaults = accounts_set & set(_LEDGER_TAB_ACCOUNTS[key])
            self.ledger_account_combos[key].set_items(
                accounts_all, default_selected=defaults
            )

    def _connect_filter_combos(self):
        for key in CONTRACT_TAB_KEYS:
            self.contract_item_combos[key].selectionChanged.connect(
                lambda k=key: self._on_contract_filter_changed(k)
            )
        for key in LEDGER_TAB_KEYS:
            self.ledger_account_combos[key].selectionChanged.connect(
                lambda k=key: self._on_ledger_filter_changed(k)
            )

    def _build_contract_tab_pivot(self, key):
        """Pivot the raw contracts data filtered to the items the tab's dropdown allows."""
        selected = set(self.contract_item_combos[key].selected_items())
        if not selected:
            return pd.DataFrame(columns=_CONTRACT_INDEX_COLS)
        sub = _preprocess_contracts_df(
            self._raw_df_c,
            item_filter=lambda items: items.isin(selected),
        )
        if sub.empty:
            return pd.DataFrame(columns=_CONTRACT_INDEX_COLS)
        return pd.pivot_table(
            sub, values="amount", index=_CONTRACT_INDEX_COLS,
            columns="month", aggfunc="sum", fill_value=0,
        ).reset_index()

    def _build_ledger_tab_pivot(self, key):
        """Build the ledger pivot for *key* using the tab's selected accounts."""
        selected = list(self.ledger_account_combos[key].selected_items())
        if not selected:
            return pd.DataFrame()
        return build_ledger_pivot(self._df_l_for_pivot, selected)

    def _rebuild_all_contract_tab_pivots(self):
        for key in CONTRACT_TAB_KEYS:
            self._contract_tab_pivots[key] = self._build_contract_tab_pivot(key)

    def _rebuild_all_ledger_tab_pivots(self):
        for key in LEDGER_TAB_KEYS:
            pivot = self._build_ledger_tab_pivot(key)
            self._ledger_tab_pivots[key] = pivot
            self._ledger_tab_diffs[key]  = calculate_mom_diff(pivot)

    def _on_contract_filter_changed(self, key):
        self._contract_tab_pivots[key] = self._build_contract_tab_pivot(key)
        self._refresh_contract_tab(key)

    def _on_ledger_filter_changed(self, key):
        pivot = self._build_ledger_tab_pivot(key)
        self._ledger_tab_pivots[key] = pivot
        self._ledger_tab_diffs[key]  = calculate_mom_diff(pivot)
        self._refresh_ledger_tab(key)

    def _refresh_contract_tab(self, key):
        tree = self.contract_trees[key]
        visible_months = self.get_visible_range()
        highlight_m    = self.period_selector.currentText()
        pivot = self._contract_tab_pivots.get(key, pd.DataFrame())
        filtered = self.filter_active_rows(pivot, visible_months)
        exp = self._save_expanded(tree)
        tree.setModel(
            build_qt_model(filtered, highlight_m, visible_months,
                           is_ledger=False, is_diff_view=False)
        )
        tree.resizeColumnToContents(0)
        self._restore_expanded(tree, exp)

    def _refresh_ledger_tab(self, key):
        tree = self.ledger_trees[key]
        visible_months = self.get_visible_range()
        highlight_m    = self.period_selector.currentText()
        use_diff       = self.check_diff.isChecked()
        source = (self._ledger_tab_diffs if use_diff
                  else self._ledger_tab_pivots).get(key, pd.DataFrame())
        filtered = self.filter_active_rows(source, visible_months, is_diff=use_diff)
        # MS is the only tab participating in discrepancy + categorisation.
        is_ms = (key == 'MS')
        discrepancy_colors = (self._compute_discrepancy_colors(highlight_m)
                              if is_ms and self.check_highlight.isChecked() else None)
        exp = self._save_expanded(tree)
        tree.setModel(
            build_qt_model(filtered, highlight_m, visible_months,
                           is_ledger=True, is_diff_view=use_diff,
                           reference_totals=discrepancy_colors,
                           category_map=self.category_map if is_ms else None,
                           manual_keys=self.manual_override_keys if is_ms else None,
                           assigned_refs=set(self.ledger_assignments.keys()))
        )
        tree.resizeColumnToContents(0)
        self._restore_expanded(tree, exp)

    def _compute_discrepancy_colors(self, highlight_m):
        """Per-customer contracts-vs-ledger discrepancy palette for the MS tab."""
        cid_col  = self.p_c.columns[0]
        lcid_col = self.p_l.columns[0]
        c_abs_data = (self.p_c.groupby(cid_col)[highlight_m].sum().to_dict()
                      if highlight_m in self.p_c.columns else {})
        l_abs_data = (self.p_l.groupby(lcid_col)[highlight_m].sum().to_dict()
                      if highlight_m in self.p_l.columns else {})
        colors = {}
        for c_id, c_val in c_abs_data.items():
            l_val = l_abs_data.get(c_id, 0)
            if abs(c_val) > 0.01:
                diff_pct = abs(l_val - c_val) / abs(c_val)
                if diff_pct > 0.10:
                    colors[c_id] = QColor(150, 100, 100)
                elif diff_pct > 0.05:
                    colors[c_id] = QColor(155, 120, 50)
            elif abs(l_val) > 1.00:
                colors[c_id] = QColor(150, 100, 100)
        return colors

    # ── Tab-page construction helper ───────────────────────────────────

    def _build_filtered_tab_page(self, tree, combo, label_text):
        """Wrap *tree* in a QWidget that puts *combo* (with *label_text*) above it."""
        page = QWidget()
        layout = QVBoxLayout(page)
        layout.setContentsMargins(0, 0, 0, 0)
        toolbar = QHBoxLayout()
        toolbar.addWidget(QLabel(label_text))
        combo.setMinimumWidth(220)
        toolbar.addWidget(combo)
        toolbar.addStretch()
        layout.addLayout(toolbar)
        layout.addWidget(tree)
        return page

    # Only the columns the data pipeline actually touches. Narrowing the SELECT
    # dramatically reduces the SQLite → pandas read cost on a ~100MB DB.
    _LEDGER_SQL = (
        'SELECT account, amount, customer_id, customer_name, '
        'contract_id, contract_name, posting_date, "reference_no." '
        'FROM ledger'
    )
    _CONTRACTS_SQL = (
        'SELECT customer_id, customer_name, contract_id, contract_name, '
        'line_number, item, item_description, scheduled_posting_date, '
        'amount, status, mrr_category '
        'FROM contracts'
    )
    _BILLING_SQL = (
        'SELECT customer_id, customer_name, contract_id, contract_name, '
        'item, scheduled_billing_date, amount, billed, status '
        'FROM contracts_billing'
    )

    def initial_load(self):
        _t_il = time.perf_counter()
        try:
            conn = init_database()
            df_l = pd.read_sql(self._LEDGER_SQL, conn)
            df_c = pd.read_sql(self._CONTRACTS_SQL, conn)
            try:
                df_b = pd.read_sql(self._BILLING_SQL, conn)
            except Exception as e:
                print(f"--- [WARN] Could not read contracts_billing table: {e} ---")
                df_b = pd.DataFrame()
            conn.close()

            # Keep raw ledger so _rebuild_ledger can re-apply assignments from a clean base.
            # Keep raw contracts so the per-tab item-filter dropdowns can re-pivot subsets.
            self._raw_df_l = df_l.copy()
            self._raw_df_c = df_c.copy()

            # Load any previously saved manual ledger→contract assignments and patch the ledger
            # BEFORE pivoting, so the row groups under its assigned contract from the start.
            ensure_ledger_assignments_table()
            self.ledger_assignments = load_ledger_assignments()
            self._df_l_for_pivot = apply_ledger_assignments(df_l.copy(), self.ledger_assignments)
            print(f"--- [CHECK] Applied {len(self.ledger_assignments)} ledger assignment(s). ---")

            # Canonical MS pivots (built from the full MS account list, ignoring the
            # MS-tab dropdown filter) — these drive category map, discrepancy, MRR
            # summary, and Excel export. The per-tab pivots used for display are
            # built separately below and follow the dropdown selections.
            self.p_l = build_ledger_pivot(self._df_l_for_pivot, _LEDGER_TAB_MS_ACCOUNTS)
            self.p_l_diff = calculate_mom_diff(self.p_l)
            self.p_c, detail_df = build_contract_outputs(df_c)

            # Populate the per-tab item/account filter dropdowns with whatever
            # codes actually appear in the data for each tab, then build the
            # initial filtered pivots (everything checked → matches full tab pivot).
            self._populate_filter_combos()
            self._connect_filter_combos()
            self._contract_tab_pivots = {}
            self._ledger_tab_pivots, self._ledger_tab_diffs = {}, {}
            self._rebuild_all_contract_tab_pivots()
            self._rebuild_all_ledger_tab_pivots()
            self._backlog_df_general, self._backlog_has_remaining_general = build_backlog_data(df_b, _PS_ITEMS_GENERAL)
            self._backlog_df_recurring, self._backlog_has_remaining_recurring = build_backlog_data(df_b, _PS_ITEMS_RECURRING)
            self.tree_backlog.setModel(build_backlog_model(
                self._backlog_df_general, self._backlog_has_remaining_general
            ))
            for col in range(4):
                self.tree_backlog.resizeColumnToContents(col)

            l_months = [c for c in self.p_l.columns if "-" in str(c)]
            c_months = [c for c in self.p_c.columns if "-" in str(c)]
            self.all_months = sorted(list(set(l_months) | set(c_months)))

            # Contracts-only category map (cached so _rebuild_ledger can re-augment from a clean base)
            self._contracts_only_category_map = build_contract_category_map(detail_df, self.all_months)
            self.category_map = augment_category_map_from_ledger(
                copy.deepcopy(self._contracts_only_category_map), self.p_l_diff, self.all_months
            )
            self._base_category_map = copy.deepcopy(self.category_map)
            print(f"--- [CHECK] Category map built for {len(self.category_map)} (customer, contract) pairs. ---")

            # Apply saved user overrides on top of the programmatic map
            ensure_category_overrides_table()
            self.manual_override_keys = set()
            for key, month_codes in load_category_overrides().items():
                self.category_map.setdefault(key, {})
                for month, code in month_codes.items():
                    self.category_map[key][month] = code
                    self.manual_override_keys.add((key[0], key[1], month))
            print(f"--- [CHECK] Loaded {len(self.manual_override_keys)} manual category override(s). ---")

            for cb in [self.period_selector, self.from_selector, self.to_selector]:
                cb.blockSignals(True)
                cb.addItems(self.all_months)

            now = datetime.now()
            current_m  = now.strftime("%Y-%m")
            year_start = now.strftime("%Y-01")
            year_end   = now.strftime("%Y-12")

            if current_m in self.all_months:
                self.period_selector.setCurrentText(current_m)
            elif self.all_months:
                self.period_selector.setCurrentIndex(len(self.all_months) - 1)

            if year_start in self.all_months:
                self.from_selector.setCurrentText(year_start)
            else:
                self.from_selector.setCurrentIndex(0)

            if year_end in self.all_months:
                self.to_selector.setCurrentText(year_end)
            elif self.all_months:
                self.to_selector.setCurrentIndex(len(self.all_months) - 1)

            for cb in [self.period_selector, self.from_selector, self.to_selector]:
                cb.blockSignals(False)

            _t_data = time.perf_counter()
            print(f"[timing] initial_load (data pipeline only):      {_ms(_t_il):7.1f} ms")
            self.update_views()
            print(f"[timing] update_views (Qt model building):       {_ms(_t_data):7.1f} ms")
            print(f"[timing] initial_load total:                     {_ms(_t_il):7.1f} ms")

        except Exception as e:
            import traceback
            traceback.print_exc()
            print(f"Data Load Error: {e}")

    def get_visible_range(self):
        start_m = self.from_selector.currentText()
        end_m   = self.to_selector.currentText()
        try:
            idx_s = self.all_months.index(start_m)
            idx_e = self.all_months.index(end_m)
            if idx_s > idx_e:
                idx_s, idx_e = idx_e, idx_s
            return self.all_months[idx_s : idx_e + 1]
        except ValueError:
            return self.all_months

    def filter_active_rows(self, df, requested_months, is_diff=False):
        if df is None or df.empty:
            return df
        existing_months = [m for m in requested_months if m in df.columns]
        if not existing_months:
            return df.iloc[0:0]
        meta_cols = [c for c in df.columns if "-" not in str(c)]
        if is_diff:
            # In diff view, keep rows that have at least one month with meaningful movement.
            # Using sum() would incorrectly drop contracts where e.g. +500 in Jan and -500 in Feb
            # cancel out — those rows have real changes and must be shown.
            mask = df[existing_months].abs().max(axis=1) > 0.01
        else:
            mask = df[existing_months].abs().sum(axis=1) > 0.01
        return df[mask][meta_cols + existing_months]

    # ── Expanded-state helpers ──────────────────────────────────────────

    def _save_expanded(self, tree) -> set:
        """Return a set of text-path tuples for every expanded node (depth ≤ 2)."""
        model = tree.model()
        if model is None:
            return set()
        expanded = set()
        for r0 in range(model.rowCount()):
            i0 = model.index(r0, 0)
            if not tree.isExpanded(i0):
                continue
            t0 = model.data(i0) or ""
            expanded.add((t0,))
            for r1 in range(model.rowCount(i0)):
                i1 = model.index(r1, 0, i0)
                if tree.isExpanded(i1):
                    expanded.add((t0, model.data(i1) or ""))
        return expanded

    def _restore_expanded(self, tree, expanded: set):
        """Re-expand nodes whose text-path appears in *expanded*."""
        model = tree.model()
        if model is None or not expanded:
            return
        for r0 in range(model.rowCount()):
            i0 = model.index(r0, 0)
            t0 = model.data(i0) or ""
            if (t0,) not in expanded:
                continue
            tree.expand(i0)
            for r1 in range(model.rowCount(i0)):
                i1 = model.index(r1, 0, i0)
                if (t0, model.data(i1) or "") in expanded:
                    tree.expand(i1)

    # ── View refresh ───────────────────────────────────────────────────

    def update_views(self):
        visible_months = self.get_visible_range()

        # Range reversal indicator
        idx_s = self.all_months.index(self.from_selector.currentText())
        idx_e = self.all_months.index(self.to_selector.currentText())
        self.lbl_status.setText("●" if idx_s > idx_e else "")

        for key in CONTRACT_TAB_KEYS:
            self._refresh_contract_tab(key)
        for key in LEDGER_TAB_KEYS:
            self._refresh_ledger_tab(key)

        self.tree_summary.setModel(
            build_category_summary_model(self.category_map, self.p_l_diff, visible_months)
        )
        self.tree_summary.resizeColumnToContents(0)

    def on_category_clicked(self, index):
        if not index.isValid() or not index.data(ROLE_IS_CATEGORY_CELL):
            return

        customer_id = index.data(ROLE_CUSTOMER_ID)
        contract_id = index.data(ROLE_CONTRACT_ID)
        month       = index.data(ROLE_MONTH)
        current_code = self.category_map.get((customer_id, contract_id), {}).get(month, 0)
        is_manual    = (customer_id, contract_id, month) in self.manual_override_keys

        dialog = CategoryDialog(current_code, is_manual=is_manual, parent=self)
        if dialog.exec() != QDialog.Accepted:
            return

        if dialog.clear_requested:
            delete_category_override(customer_id, contract_id, month)
            self.manual_override_keys.discard((customer_id, contract_id, month))
            base_code = self._base_category_map.get((customer_id, contract_id), {}).get(month, 0)
            self.category_map.setdefault((customer_id, contract_id), {})[month] = base_code
        elif dialog.selected_code is not None:
            save_category_override(customer_id, contract_id, month, dialog.selected_code)
            self.category_map.setdefault((customer_id, contract_id), {})[month] = dialog.selected_code
            self.manual_override_keys.add((customer_id, contract_id, month))

        self.update_views()

    # ── Manual ledger→contract assignment ──────────────────────────────

    def _get_assignable_index(self, index):
        """Return an index bearing assignment roles if the row is assignable/reassignable, else None."""
        if not index.isValid():
            return None

        def _valid(idx):
            if idx.isValid() and idx.data(ROLE_IS_UNASSIGNED_REF):
                cust = idx.data(ROLE_CUSTOMER_ID)
                if cust and cust not in MISSING_SENTINELS:
                    return idx
            return None

        valid = _valid(index)
        if valid is not None:
            return valid
        # Fallback: column 0 of the same row. With the custom model, the roles are
        # returned identically across columns of a line row, but this guards against
        # any future divergence.
        return _valid(index.sibling(index.row(), 0))

    def on_ledger_ref_double_clicked(self, index):
        idx = self._get_assignable_index(index)
        if idx is None:
            return
        self._open_assign_dialog(
            idx.data(ROLE_CUSTOMER_ID),
            idx.data(ROLE_REF_NO),
            idx.data(ROLE_ASSIGNED_CONTRACT_ID),
            bool(idx.data(ROLE_CAN_CLEAR)),
        )

    def on_ledger_context_menu(self, pos):
        # Signal fires on whichever ledger tab tree the user right-clicked.
        tree = self.sender()
        if not isinstance(tree, QTreeView):
            return
        index = tree.indexAt(pos)
        idx   = self._get_assignable_index(index)
        if idx is None:
            return
        customer_id         = idx.data(ROLE_CUSTOMER_ID)
        ref_no              = idx.data(ROLE_REF_NO)
        current_contract_id = idx.data(ROLE_ASSIGNED_CONTRACT_ID)
        can_clear           = bool(idx.data(ROLE_CAN_CLEAR))
        menu   = QMenu(self)
        label  = f"Change contract for \"{ref_no}\"…" if current_contract_id else f"Assign \"{ref_no}\" to a contract…"
        action = menu.addAction(label)
        if menu.exec(tree.viewport().mapToGlobal(pos)) == action:
            self._open_assign_dialog(customer_id, ref_no, current_contract_id, can_clear)

    def _open_assign_dialog(self, customer_id: str, ref_no: str,
                             current_contract_id: str = None, can_clear: bool = False):
        dialog = AssignContractDialog(customer_id, ref_no,
                                      current_contract_id=current_contract_id,
                                      can_clear=can_clear, parent=self)
        if dialog.exec() != QDialog.Accepted:
            return
        if dialog.clear_requested:
            delete_ledger_assignment(ref_no, customer_id)
            self.ledger_assignments.pop((str(ref_no), str(customer_id)), None)
            self._rebuild_ledger()
        elif dialog.selected_contract is not None:
            contract_id, contract_name = dialog.selected_contract
            save_ledger_assignment(ref_no, customer_id, contract_id, contract_name)
            self.ledger_assignments[(str(ref_no), str(customer_id))] = (contract_id, contract_name)
            self._rebuild_ledger()

    def _rebuild_ledger(self):
        """Re-pivot the ledger after a new assignment and cascade through the category map."""
        QApplication.setOverrideCursor(Qt.WaitCursor)
        try:
            self._rebuild_ledger_inner()
        finally:
            QApplication.restoreOverrideCursor()

    def _rebuild_ledger_inner(self):
        self._df_l_for_pivot = apply_ledger_assignments(self._raw_df_l.copy(), self.ledger_assignments)
        # Canonical MS pivot — full account list, independent of the dropdown filter.
        self.p_l = build_ledger_pivot(self._df_l_for_pivot, _LEDGER_TAB_MS_ACCOUNTS)
        self.p_l_diff = calculate_mom_diff(self.p_l)
        # Per-tab display pivots follow the current dropdown selections.
        self._rebuild_all_ledger_tab_pivots()

        # Re-augment category map from updated ledger; contracts-only base stays unchanged.
        self.category_map = augment_category_map_from_ledger(
            copy.deepcopy(self._contracts_only_category_map),
            self.p_l_diff,
            self.all_months
        )
        self._base_category_map = copy.deepcopy(self.category_map)

        # Re-apply user category overrides (these are independent of ledger assignments).
        for key, month_codes in load_category_overrides().items():
            self.category_map.setdefault(key, {})
            for month, code in month_codes.items():
                self.category_map[key][month] = code

        self.update_views()

    def _rebuild_backlog_tree(self, _=False):
        """Rebuild the backlog tree based on the current PS-type toggle and outstanding filter."""
        if not hasattr(self, "_backlog_df_general"):
            return
        if self.rb_backlog_recurring.isChecked():
            df, remaining = self._backlog_df_recurring, self._backlog_has_remaining_recurring
        else:
            df, remaining = self._backlog_df_general, self._backlog_has_remaining_general
        self.tree_backlog.setModel(build_backlog_model(
            df, remaining,
            outstanding_only=self.chk_backlog_outstanding.isChecked(),
        ))
        for col in range(4):
            self.tree_backlog.resizeColumnToContents(col)

    def on_summary_double_clicked(self, index):
        if not index.isValid():
            return
        code  = index.data(ROLE_CATEGORY_CODE)
        month = index.data(ROLE_MONTH)
        if code is None or month is None:
            return
        visible_months = self.get_visible_range()
        DrillDownDialog(code, month, self.category_map, self.p_l_diff,
                        visible_months, self._df_l_for_pivot,
                        parent=self).exec()

    def _write_grouped_sheet(self, writer, sheet_name, df, summary_fn=None):
        """
        Write *df* to *sheet_name* with three-level collapsible row groups.

        Input columns: customer_id, customer_name, contract_id, contract_name,
        <detail>, value_col1, value_col2, ...

        Output columns: customer_name | contract_name | detail | value_cols...
        ID columns are dropped. Each name only appears on its own summary row:
        customer_name on level-0 rows, contract_name on level-1 rows,
        detail on level-2 leaf rows. Contracts and detail rows are hidden by
        default so the sheet opens fully collapsed.
        """
        from openpyxl.worksheet.properties import Outline
        from openpyxl.styles import Font

        if df is None or df.empty:
            (df if df is not None else pd.DataFrame()).to_excel(
                writer, sheet_name=sheet_name, index=False
            )
            return

        cols = list(df.columns)
        cust_id_col, cust_name_col, ct_id_col, ct_name_col = cols[:4]
        detail_col = cols[4]
        value_cols = cols[5:]
        out_cols = [cust_name_col, ct_name_col, detail_col] + value_cols

        df_sorted = df.sort_values([cust_name_col, ct_name_col], kind="stable")

        def _default_summary(group, vcols):
            return {v: group[v].sum() for v in vcols}
        agg = summary_fn or _default_summary

        out_rows = []
        levels = []
        for (cid, cname), c_grp in df_sorted.groupby(
            [cust_id_col, cust_name_col], sort=False
        ):
            # Customer summary row — only customer_name populated
            row = {c: "" for c in out_cols}
            row[cust_name_col] = cname
            row.update(agg(c_grp, value_cols))
            out_rows.append(row)
            levels.append(0)

            for (ctid, ctname), ct_grp in c_grp.groupby(
                [ct_id_col, ct_name_col], sort=False
            ):
                # Contract summary row — only contract_name populated
                row = {c: "" for c in out_cols}
                row[ct_name_col] = ctname
                row.update(agg(ct_grp, value_cols))
                out_rows.append(row)
                levels.append(1)

                # Detail rows — only detail_col and value_cols populated
                for _, r in ct_grp.iterrows():
                    row = {c: "" for c in out_cols}
                    row[detail_col] = r[detail_col]
                    for v in value_cols:
                        row[v] = r[v]
                    out_rows.append(row)
                    levels.append(2)

        pd.DataFrame(out_rows, columns=out_cols).to_excel(
            writer, sheet_name=sheet_name, index=False
        )
        ws = writer.sheets[sheet_name]
        # summaryBelow=False puts the parent row above its detail rows, matching
        # the in-app tree view where a customer/contract sits above its children.
        ws.sheet_properties.outlinePr = Outline(summaryBelow=False)
        bold_font = Font(bold=True)
        for i, lvl in enumerate(levels):
            excel_row = i + 2
            if lvl > 0:
                rd = ws.row_dimensions[excel_row]
                rd.outline_level = lvl
                rd.hidden = True
            if lvl == 0:
                for cell in ws[excel_row]:
                    cell.font = bold_font

    def export_to_excel(self):
        visible_months = self.get_visible_range()
        filename = f"MRR_Export_{datetime.now().strftime('%Y%m%d_%H%M%S')}.xlsx"

        try:
            with pd.ExcelWriter(filename, engine='openpyxl') as writer:
                # 1. Contracts View
                df_exp_c = self.filter_active_rows(self.p_c, visible_months)
                self._write_grouped_sheet(writer, "Contracts", df_exp_c)

                # 2. Ledger View (Full)
                df_exp_l = self.filter_active_rows(self.p_l, visible_months)
                self._write_grouped_sheet(writer, "Ledger_Actuals", df_exp_l)

                # 3. Ledger MoM Differences
                df_exp_diff = self.filter_active_rows(self.p_l_diff, visible_months)
                self._write_grouped_sheet(writer, "Ledger_MoM_Changes", df_exp_diff)

                # 4. MRR Summary (category × month)
                cid_col  = self.p_l_diff.columns[0]
                ctid_col = self.p_l_diff.columns[2]
                month_cols = [m for m in visible_months if m in self.p_l_diff.columns]
                contract_totals = self.p_l_diff.groupby([cid_col, ctid_col])[month_cols].sum()
                cat_amounts = _aggregate_summary_by_category(
                    self.category_map, contract_totals, month_cols
                )
                summary_rows = []
                totals: dict = {}
                for code, label in ALL_CATEGORIES:
                    if code not in cat_amounts:
                        continue
                    r = {"Category": f"{code}  —  {label}"}
                    for m in visible_months:
                        r[m] = cat_amounts[code].get(m, 0.0)
                        totals[m] = totals.get(m, 0.0) + r[m]
                    summary_rows.append(r)
                if summary_rows:
                    total_row = {"Category": "Net MRR Change"}
                    total_row.update(totals)
                    summary_rows.append(total_row)
                    pd.DataFrame(summary_rows).to_excel(writer, sheet_name="MRR_Summary", index=False)

                # 5. Backlog reports — General PS and Recurring PS (I904).
                def _backlog_summary(group, _vcols):
                    sched  = float(group["scheduled_amount"].sum())
                    billed = float(group["billed_amount"].sum())
                    return {
                        "scheduled_amount":  sched,
                        "billed_amount":     billed,
                        "remaining_balance": sched - billed,
                    }
                for sheet_name, bl_df in (
                    ("Backlog_General_PS",   getattr(self, "_backlog_df_general", None)),
                    ("Backlog_Recurring_PS", getattr(self, "_backlog_df_recurring", None)),
                ):
                    if bl_df is None or bl_df.empty:
                        continue
                    bl_sorted = bl_df.sort_values(
                        ["customer_name", "contract_name", "month"], kind="stable"
                    ).copy()
                    bl_sorted["remaining_balance"] = (
                        bl_sorted.groupby(["customer_id", "contract_id"], sort=False)["amount"].transform("sum")
                        - bl_sorted.groupby(["customer_id", "contract_id"], sort=False)["billed"].cumsum()
                    )
                    bl_sorted = bl_sorted.rename(columns={
                        "amount": "scheduled_amount",
                        "billed": "billed_amount",
                    })[[
                        "customer_id", "customer_name", "contract_id", "contract_name",
                        "month", "scheduled_amount", "billed_amount", "remaining_balance",
                    ]]
                    self._write_grouped_sheet(writer, sheet_name, bl_sorted, summary_fn=_backlog_summary)

            QMessageBox.information(self, "Export Successful",
                                    f"Data exported successfully to:\n{os.path.abspath(filename)}")
        except Exception as e:
            QMessageBox.critical(self, "Export Failed",
                                 f"An error occurred while exporting:\n{str(e)}")


if __name__ == "__main__":
    _t = time.perf_counter()
    app = QApplication(sys.argv)
    print(f"[timing] QApplication(sys.argv):                 {_ms(_t):7.1f} ms")

    # Increase application-wide font size
    app_font = QApplication.font()
    app_font.setPointSize(app_font.pointSize() + _FONT_SIZE_OFFSET)
    QApplication.setFont(app_font)

    _t = time.perf_counter()
    window = MainWindow()
    print(f"[timing] MainWindow() (incl. initial_load):      {_ms(_t):7.1f} ms")

    _t = time.perf_counter()
    window.show()
    print(f"[timing] window.show():                          {_ms(_t):7.1f} ms")

    def _first_event_tick():
        print(f"[timing] first event-loop tick (≈ first paint):  {_ms(_T0):7.1f} ms TOTAL since process start")
    QTimer.singleShot(0, _first_event_tick)

    sys.exit(app.exec())
