#!/usr/bin/env python3
"""Add inflation formulas to 'Betriebliche Aufwendungen' rows + update Büromiete.

What this does:
  1. Update row 27 (Büromiete): 0 in Aug 2026, then 1000 €/Monat from Sep 2026 to
     Dec 2026 (hardcoded), then inflation-adjusted formulas from Jan 2027 onwards.
  2. Add Treiber driver block 'Inflation Betriebskosten' (one rate per year).
  3. For each eligible row, keep 2026 values hardcoded (the 'Startwert') and
     replace cells from Jan 2027 onwards with a formula that:
        - carries the previous month's value forward in non-January months
        - multiplies by (1 + inflation_for_this_year) every January

Eligibility (which rows get inflation formulas):
  - Row contains numeric values (not formulas) — formula rows like KFZ that
    scale via Personalkosten are left alone.
  - Row has at least 2 non-zero months in 2026 (constant or 'later activation'
    patterns: D&O insurance, Recruiting, Buchführung, etc.).
  - Yearly-fee rows (exactly 1 non-zero month in 2026 — e.g., IHK in Sep)
    get a year-over-year carry: same month next year = last year × (1+inflation),
    other months stay 0.
  - Rows where all 2026 values are 0 are skipped.
  - Subtotal/summe rows are formulas already → automatically skipped.

Usage:
    python3 pitch-deck/scripts/add-inflation-formulas.py --dry-run
    python3 pitch-deck/scripts/add-inflation-formulas.py
"""

from __future__ import annotations

import argparse
import shutil
import sys
from datetime import datetime
from pathlib import Path

from openpyxl import load_workbook
from openpyxl.utils import get_column_letter

EXPORTS = Path(__file__).resolve().parent.parent / "exports"

BUEROMIETE_ROW = 27
BUEROMIETE_START_MONTH = 9  # September
BUEROMIETE_VALUE = 1000

DEFAULT_INFLATION = 0.03

_BA_SHEET = "Betriebliche Aufwendungen"


def year_columns(ws) -> dict[int, list[int]]:
    out: dict[int, list[int]] = {}
    for c in range(2, ws.max_column + 1):
        v = ws.cell(row=1, column=c).value
        if v is None:
            continue
        try:
            yr = int(v)
        except (TypeError, ValueError):
            continue
        out.setdefault(yr, []).append(c)
    return out


def cols_by_year_month(ws) -> dict[tuple[int, int], int]:
    """Return {(year, month): column_index} based on rows 1 and 2."""
    out: dict[tuple[int, int], int] = {}
    for c in range(2, ws.max_column + 1):
        y = ws.cell(row=1, column=c).value
        m = ws.cell(row=2, column=c).value
        try:
            out[(int(y), int(m))] = c
        except (TypeError, ValueError):
            continue
    return out


# ---------------------------------------------------------------------------
# Treiber: add or read inflation driver block
# ---------------------------------------------------------------------------

INFLATION_HEADER = "Inflation Betriebskosten"


def find_or_create_inflation_block(ws_t, years_after_start: list[int]) -> dict[int, int]:
    """Find existing inflation block by header, or append a new one. Return {year: row}."""
    # Look for existing header
    header_row = None
    for r in range(1, ws_t.max_row + 1):
        if ws_t.cell(row=r, column=1).value == INFLATION_HEADER:
            header_row = r
            break

    if header_row is None:
        # Append new block
        header_row = ws_t.max_row + 2  # blank then header
        ws_t.cell(row=header_row, column=1).value = INFLATION_HEADER

    refs: dict[int, int] = {}
    r = header_row + 1
    for yr in years_after_start:
        label_cell = ws_t.cell(row=r, column=1)
        if label_cell.value is None or not str(label_cell.value).startswith("Inflation"):
            # Fresh entry
            ws_t.cell(row=r, column=1).value = f"Inflation {yr}"
            ws_t.cell(row=r, column=2).value = DEFAULT_INFLATION
        # If already there, preserve existing value
        refs[yr] = r
        r += 1
    return refs


# ---------------------------------------------------------------------------
# Büromiete special handling (row 27)
# ---------------------------------------------------------------------------


def update_bueromiete(ws_ba, col_index: dict, base_year: int, inflation_refs: dict[int, int]) -> int:
    """Set row 27: 0 in Aug 2026, 1000 from Sep 2026 to Dec 2026 (hardcoded),
    then inflation formula from Jan 2027 onwards.
    """
    inc_first = min(inflation_refs.values())
    inc_last = max(inflation_refs.values())

    # Keep label free of the word 'Inflation' — number-formatting heuristics
    # treat that as a percent indicator and would mis-format Euro cells.
    ws_ba.cell(row=BUEROMIETE_ROW, column=1).value = "Büromiete (ab Sep 2026)"

    written = 0
    for c in range(2, ws_ba.max_column + 1):
        y = ws_ba.cell(row=1, column=c).value
        m = ws_ba.cell(row=2, column=c).value
        if y == base_year:
            # Hardcoded 2026 values
            val = BUEROMIETE_VALUE if m >= BUEROMIETE_START_MONTH else 0
            ws_ba.cell(row=BUEROMIETE_ROW, column=c).value = val
            written += 1
        else:
            # Inflation formula from Jan 2027 onwards
            col_letter = get_column_letter(c)
            prev_col = get_column_letter(c - 1)
            ws_ba.cell(row=BUEROMIETE_ROW, column=c).value = (
                f"=IF(AND({col_letter}$2=1,{col_letter}$1>{base_year}),"
                f"{prev_col}{BUEROMIETE_ROW}*(1+INDEX(Treiber!$B${inc_first}:$B${inc_last},{col_letter}$1-{base_year})),"
                f"{prev_col}{BUEROMIETE_ROW})"
            )
            written += 1
    return written


# ---------------------------------------------------------------------------
# General inflation application to qualifying rows
# ---------------------------------------------------------------------------


def is_formula(value) -> bool:
    return isinstance(value, str) and value.startswith("=")


def collect_eligible_rows(ws_ba, base_year_cols: list[int], skip_rows: set[int]) -> dict[int, str]:
    """For each row, decide if it's eligible and classify it.

    Returns {row: kind} where kind is one of:
        'monthly'   - ≥2 non-zero months in 2026, apply Jan-inflation carry-forward
        'yearly'    - exactly 1 non-zero month in 2026, apply same-month-next-year
        (rows not eligible are absent from the dict)
    """
    out: dict[int, str] = {}
    for r in range(4, ws_ba.max_row + 1):
        if r in skip_rows:
            continue
        label = ws_ba.cell(row=r, column=1).value
        if not label:
            continue
        if str(label).startswith(("summe —", "Summe ")) or "SUMME" in str(label):
            continue

        values_2026 = [ws_ba.cell(row=r, column=c).value for c in base_year_cols]
        if any(is_formula(v) for v in values_2026):
            continue
        nonzero_count = sum(1 for v in values_2026 if v not in (None, 0, ""))
        if nonzero_count == 0:
            continue
        if nonzero_count == 1:
            out[r] = "yearly"
        else:
            out[r] = "monthly"
    return out


def write_monthly_inflation(ws_ba, row: int, first_formula_col: int, base_year: int,
                            inflation_refs: dict[int, int]) -> int:
    """From first_formula_col onwards, carry previous + Jan inflation."""
    inc_first = min(inflation_refs.values())
    inc_last = max(inflation_refs.values())
    written = 0
    for c in range(first_formula_col, ws_ba.max_column + 1):
        col_letter = get_column_letter(c)
        prev_col = get_column_letter(c - 1)
        ws_ba.cell(row=row, column=c).value = (
            f"=IF(AND({col_letter}$2=1,{col_letter}$1>{base_year}),"
            f"{prev_col}{row}*(1+INDEX(Treiber!$B${inc_first}:$B${inc_last},{col_letter}$1-{base_year})),"
            f"{prev_col}{row})"
        )
        written += 1
    return written


def write_yearly_inflation(ws_ba, row: int, first_formula_col: int, base_year: int,
                           inflation_refs: dict[int, int], col_index: dict) -> int:
    """For yearly-fee rows: same month next year × (1 + inflation), else 0."""
    inc_first = min(inflation_refs.values())
    inc_last = max(inflation_refs.values())
    written = 0
    for c in range(first_formula_col, ws_ba.max_column + 1):
        y = ws_ba.cell(row=1, column=c).value
        m = ws_ba.cell(row=2, column=c).value
        if y is None or m is None:
            continue
        same_month_prev_year = col_index.get((int(y) - 1, int(m)))
        if same_month_prev_year is None:
            continue
        prev_letter = get_column_letter(same_month_prev_year)
        col_letter = get_column_letter(c)
        ws_ba.cell(row=row, column=c).value = (
            f"={prev_letter}{row}*(1+INDEX(Treiber!$B${inc_first}:$B${inc_last},{col_letter}$1-{base_year}))"
        )
        written += 1
    return written


# ---------------------------------------------------------------------------
# Main per-file processing
# ---------------------------------------------------------------------------


def process_file(path: Path, dry_run: bool) -> dict | None:
    wb = load_workbook(path)
    if _BA_SHEET not in wb.sheetnames or "Treiber" not in wb.sheetnames:
        return None
    ws_ba = wb[_BA_SHEET]
    ws_t = wb["Treiber"]

    yc = year_columns(ws_ba)
    years = sorted(yc.keys())
    base_year = min(years)
    years_after_start = [y for y in years if y > base_year]

    inflation_refs = find_or_create_inflation_block(ws_t, years_after_start)

    col_index = cols_by_year_month(ws_ba)
    base_year_cols = yc[base_year]
    first_year_after = min(years_after_start)
    first_formula_col = col_index.get((first_year_after, 1))
    if first_formula_col is None:
        return None

    # Büromiete (always)
    bueromiete_cells = update_bueromiete(ws_ba, col_index, base_year, inflation_refs)

    # Determine eligible rows (skip Büromiete since we just handled it)
    skip = {BUEROMIETE_ROW}
    eligible = collect_eligible_rows(ws_ba, base_year_cols, skip)

    cells_monthly = 0
    cells_yearly = 0
    n_monthly = n_yearly = 0
    for row, kind in eligible.items():
        if kind == "monthly":
            cells_monthly += write_monthly_inflation(ws_ba, row, first_formula_col, base_year, inflation_refs)
            n_monthly += 1
        else:
            cells_yearly += write_yearly_inflation(ws_ba, row, first_formula_col, base_year, inflation_refs, col_index)
            n_yearly += 1

    if not dry_run:
        wb.save(path)

    return {
        "years": years,
        "monthly_rows": n_monthly,
        "yearly_rows": n_yearly,
        "cells_bueromiete": bueromiete_cells,
        "cells_monthly": cells_monthly,
        "cells_yearly": cells_yearly,
        "eligible_rows": {kind: [r for r, k in eligible.items() if k == kind] for kind in ("monthly", "yearly")},
    }


def backup(path: Path) -> Path:
    ts = datetime.now().strftime("%Y%m%d-%H%M%S")
    bk = path.with_name(f"{path.stem}.BACKUP-pre-inflation-{ts}{path.suffix}")
    shutil.copy2(path, bk)
    return bk


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--dry-run", action="store_true")
    ap.add_argument("--only", help="Process only this filename")
    ap.add_argument("--no-backup", action="store_true")
    args = ap.parse_args()

    files = sorted(EXPORTS.glob("Finanzplan-*.xlsx"))
    files = [f for f in files if "BACKUP" not in f.name]
    if args.only:
        files = [f for f in files if f.name == args.only]

    for path in files:
        wb_peek = load_workbook(path, read_only=True)
        if _BA_SHEET not in wb_peek.sheetnames or "Treiber" not in wb_peek.sheetnames:
            print(f"\n  ⨯ skip {path.name}: missing sheets")
            continue
        if not args.dry_run and not args.no_backup:
            bk = backup(path)
            print(f"  ✓ backup: {bk.name}")
        info = process_file(path, dry_run=args.dry_run)
        if info is None:
            continue
        print(f"\n  === {path.name} ===")
        print(f"    Büromiete (row {BUEROMIETE_ROW}): {info['cells_bueromiete']} cells (Sep 2026 start)")
        print(f"    Monthly-inflation rows: {info['monthly_rows']} ({info['cells_monthly']} cells)")
        print(f"    Yearly-fee rows: {info['yearly_rows']} ({info['cells_yearly']} cells)")
        print(f"    Eligible monthly rows: {info['eligible_rows']['monthly']}")
        print(f"    Eligible yearly rows: {info['eligible_rows']['yearly']}")
    return 0


if __name__ == "__main__":
    sys.exit(main())