#!/usr/bin/env python3
"""
Update HFP Excel files to add HSA contributions, Debts, and Fixed Assets
for broader feature coverage in TOML regression tests.

Run from project root: python scripts/update_hfp_coverage.py
"""
import os
import pandas as pd
from pathlib import Path

EXAMPLES = Path("examples")
# All dollar amounts in HFP_CONFIG are in $k for readability. Convert to $ when writing.
HFP_CONFIG = {
    "HFP_jack+jill.xlsx": {
        "sheets": ["Jack", "Jill"],
        "hsa": {"Jack": (2026, 2027, 4.0), "Jill": (2026, 2030, 4.0)},  # year range, amt in $k
        "debt": {
            "active": True, "name": "Primary mortgage", "type": "mortgage",
            "year": 2020, "term": 20, "amount": 180, "rate": 4.5,
        },  # amount in $k
        "fixed_asset": {
            "active": True, "name": "Vacation condo", "type": "real estate",
            "year": 2026, "basis": 120, "value": 180, "rate": 3, "yod": 2035,
            "commission": 6,
        },  # basis, value in $k
    },
    "HFP_joe.xlsx": {
        "sheets": ["Joe"],
        "hsa": {"Joe": (2026, 2027, 4.5)},
        "fixed_asset": {
            "active": True, "name": "Rental property", "type": "real estate",
            "year": 2026, "basis": 200, "value": 280, "rate": 3, "yod": 2032,
            "commission": 5,
        },
    },
    "HFP_john+sally.xlsx": {
        "sheets": ["John", "Sally"],
        "hsa": {"John": (2026, 2027, 4.0), "Sally": (2026, 2027, 4.0)},
    },
    "HFP_kim+sam.xlsx": {
        "sheets": ["Kim", "Sam"],
        "hsa": {"Kim": (2026, 2028, 4.0), "Sam": (2026, 2029, 4.0)},
        "fixed_asset": {
            "active": True, "name": "Downsize home sale", "type": "residence",
            "year": 2026, "basis": 350, "value": 520, "rate": 3, "yod": 2038,
            "commission": 5,
        },
    },
    "HFP_alex+jamie.xlsx": {
        "sheets": ["Alex", "Jamie"],
        "hsa": {"Alex": (2026, 2027, 4.0), "Jamie": (2026, 2027, 4.0)},
    },
    "HFP_jon+jane.xlsx": {
        "sheets": ["Jon", "Jane"],
        "hsa": {"Jon": (2026, 2028, 4.0), "Jane": (2026, 2029, 4.0)},
    },
}

TIME_HORIZON_ITEMS = [
    "year", "anticipated wages", "other inc", "net inv",
    "taxable ctrb", "401k ctrb", "Roth 401k ctrb", "IRA ctrb", "Roth IRA ctrb",
    "HSA ctrb", "Roth conv", "big-ticket items",
]
DEBT_ITEMS = ["active", "name", "type", "year", "term", "amount", "rate"]
FIXED_ASSET_ITEMS = ["active", "name", "type", "year", "basis", "value", "rate", "yod", "commission"]


def add_hsa_column(df, sheet_name, hsa_spec):
    if sheet_name not in hsa_spec:
        return df
    start_yr, end_yr, amt = hsa_spec[sheet_name]
    if "HSA ctrb" not in df.columns:
        df["HSA ctrb"] = 0.0
    # Reset all to 0, then set in-range years (idempotent)
    df["HSA ctrb"] = 0.0
    for idx, row in df.iterrows():
        y = int(row["year"])
        if start_yr <= y <= end_yr:
            df.at[idx, "HSA ctrb"] = amt * 1000  # $k to $
    return df


def main():
    os.chdir(Path(__file__).resolve().parent.parent)
    for fname, cfg in HFP_CONFIG.items():
        path = EXAMPLES / fname
        if not path.exists():
            print(f"Skip {fname} (not found)")
            continue
        df_dict = pd.read_excel(path, sheet_name=None)
        for sh in cfg["sheets"]:
            if sh in df_dict:
                df_dict[sh] = add_hsa_column(df_dict[sh].copy(), sh, cfg.get("hsa", {}))
        if "debt" in cfg and "Debts" in df_dict:
            d = cfg["debt"].copy()
            if "amount" in d and d["amount"]:
                d["amount"] = float(d["amount"]) * 1000  # $k -> $
            df_dict["Debts"] = pd.DataFrame([{k: d.get(k) for k in DEBT_ITEMS}])
        if "fixed_asset" in cfg and "Fixed Assets" in df_dict:
            fa = cfg["fixed_asset"].copy()
            for col in ("basis", "value"):
                if col in fa and fa[col] is not None:
                    fa[col] = float(fa[col]) * 1000  # $k -> $
            df_dict["Fixed Assets"] = pd.DataFrame([{k: fa.get(k) for k in FIXED_ASSET_ITEMS}])
        with pd.ExcelWriter(path, engine="openpyxl") as w:
            for sh, df in df_dict.items():
                df.to_excel(w, sheet_name=sh, index=False)
        print(f"Updated {fname}")


if __name__ == "__main__":
    main()
