#!/usr/bin/env python3
"""Add price-escalation formulas to Finanzplan Excel files.

Replaces hard-coded Preis/Monat values in 'Umsatzerlöse' rows 4, 7, 10 with:
- Starting price (Aug 2026) read from a Treiber driver
- Annual price increase applied in January of each subsequent year
- Increase percentage configurable per year via Treiber driver

Treiber layout (appended after existing rows):
    blank
    'Preise'
    'Startpreis/Monat Starter'         → 300
    'Startpreis/Monat Professional'    → 2083
    'Startpreis/Monat Enterprise'      → 4167
    blank
    'Preiserhöhung pro Jahr (%)'
    'Preiserhöhung 2027'               → 0.03
    'Preiserhöhung 2028'               → 0.03
    'Preiserhöhung 2029'               → 0.03
    'Preiserhöhung 2030'               → 0.03

Usage:
    python3 pitch-deck/scripts/add-price-formulas.py --dry-run
    python3 pitch-deck/scripts/add-price-formulas.py
    python3 pitch-deck/scripts/add-price-formulas.py --only Finanzplan-Wandeldarlehen-400k.xlsx
"""

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"

# Segment → row in Umsatzerlöse where Preis/Monat lives
PRICE_ROWS = {
    "starter": 4,
    "professional": 7,
    "enterprise": 10,
}

# Default starting prices (Aug 2026) — used only if existing value isn't a number
DEFAULT_START_PRICES = {
    "starter": 300,
    "professional": 2083,
    "enterprise": 4167,
}

DEFAULT_PRICE_INCREASE = 0.03  # 3% per year by default

_SEG_LABEL = {"starter": "Starter", "professional": "Professional", "enterprise": "Enterprise"}


def year_columns(ws_kunden) -> dict[int, list[int]]:
    out: dict[int, list[int]] = {}
    for c in range(2, ws_kunden.max_column + 1):
        v = ws_kunden.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 existing_start_prices(ws_umsatz) -> dict[str, float]:
    """Read current Aug 2026 (column B) prices to use as Startpreis defaults."""
    res: dict[str, float] = {}
    for suffix, row in PRICE_ROWS.items():
        v = ws_umsatz.cell(row=row, column=2).value
        try:
            res[suffix] = float(v) if v is not None else float(DEFAULT_START_PRICES[suffix])
        except (TypeError, ValueError):
            # Already a formula? Use default.
            res[suffix] = float(DEFAULT_START_PRICES[suffix])
    return res


def already_processed(ws_treiber) -> tuple[bool, int]:
    """Return (already_processed, start_row). start_row is where to place the Preise block."""
    # Scan for an existing 'Preise' header
    for r in range(1, ws_treiber.max_row + 1):
        if ws_treiber.cell(row=r, column=1).value == "Preise":
            return True, r
    # First time: place after max_row with a blank separator
    return False, ws_treiber.max_row + 2


def read_existing_price_drivers(ws_treiber, header_row: int, years_after_start: list[int]) -> tuple[dict, dict, dict]:
    """Read existing price drivers we wrote previously. Returns (start_prices, increases, refs)."""
    start_prices: dict[str, float] = {}
    increases: dict[int, float] = {}
    refs: dict[str, int | dict[int, int]] = {}

    r = header_row + 1
    start_refs: dict[str, int] = {}
    for suffix in ("starter", "professional", "enterprise"):
        v = ws_treiber.cell(row=r, column=2).value
        try:
            start_prices[suffix] = float(v) if v is not None else float(DEFAULT_START_PRICES[suffix])
        except (TypeError, ValueError):
            start_prices[suffix] = float(DEFAULT_START_PRICES[suffix])
        start_refs[suffix] = r
        r += 1
    refs["start_refs"] = start_refs

    r += 1  # blank
    r += 1  # 'Preiserhöhung' header
    increase_refs: dict[int, int] = {}
    for yr in years_after_start:
        v = ws_treiber.cell(row=r, column=2).value
        try:
            increases[yr] = float(v) if v is not None else DEFAULT_PRICE_INCREASE
        except (TypeError, ValueError):
            increases[yr] = DEFAULT_PRICE_INCREASE
        increase_refs[yr] = r
        r += 1
    refs["increase_refs"] = increase_refs

    return start_prices, increases, refs


def write_treiber_price_drivers(ws_treiber, start_row: int, years_after_start: list[int],
                                 start_prices: dict[str, float], increases: dict[int, float]) -> dict:
    r = start_row
    ws_treiber.cell(row=r, column=1).value = "Preise"
    r += 1
    start_refs: dict[str, int] = {}
    for suffix in ("starter", "professional", "enterprise"):
        ws_treiber.cell(row=r, column=1).value = f"Startpreis/Monat {_SEG_LABEL[suffix]}"
        ws_treiber.cell(row=r, column=2).value = start_prices[suffix]
        start_refs[suffix] = r
        r += 1
    r += 1  # blank
    ws_treiber.cell(row=r, column=1).value = "Preiserhöhung pro Jahr (%)"
    r += 1
    increase_refs: dict[int, int] = {}
    for yr in years_after_start:
        ws_treiber.cell(row=r, column=1).value = f"Preiserhöhung {yr}"
        ws_treiber.cell(row=r, column=2).value = increases.get(yr, DEFAULT_PRICE_INCREASE)
        increase_refs[yr] = r
        r += 1
    return {"start_refs": start_refs, "increase_refs": increase_refs}


def write_umsatz_price_formulas(ws_umsatz, refs: dict, base_year: int) -> int:
    """Write price formulas to Umsatzerlöse rows 4/7/10 across all month columns.

    Formula structure:
      col B (Aug 2026, first month):  =Treiber!$B$startpreis
      col C+ same year:               =prev_col
      col where month=1 and year>base_year:
                                     =prev_col*(1+INDEX(Treiber!$B$inc_first:$B$inc_last,year-base_year))
    """
    start_refs = refs["start_refs"]
    increase_refs = refs["increase_refs"]
    inc_first = min(increase_refs.values())
    inc_last = max(increase_refs.values())

    cells_written = 0
    for suffix, row in PRICE_ROWS.items():
        # Clean label
        ws_umsatz.cell(row=row, column=1).value = f"Preis/Monat ({_SEG_LABEL[suffix]})"
        for c in range(2, ws_umsatz.max_column + 1):
            col_letter = get_column_letter(c)
            if c == 2:
                # Aug 2026 starting price
                ws_umsatz.cell(row=row, column=c).value = (
                    f"=Treiber!$B${start_refs[suffix]}"
                )
            else:
                prev_col = get_column_letter(c - 1)
                ws_umsatz.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},"
                    f"{col_letter}$1-{base_year})),"
                    f"{prev_col}{row})"
                )
            cells_written += 1
    return cells_written


def process_file(path: Path, dry_run: bool) -> dict | None:
    wb = load_workbook(path)
    if "Umsatzerlöse" not in wb.sheetnames or "Treiber" not in wb.sheetnames or "Kunden" not in wb.sheetnames:
        return None

    ws_u = wb["Umsatzerlöse"]
    ws_t = wb["Treiber"]
    ws_k = wb["Kunden"]

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

    processed, start_row = already_processed(ws_t)
    if processed:
        # Preserve existing driver values
        start_prices, increases, _ = read_existing_price_drivers(ws_t, start_row, years_after_start)
        source = "treiber (preserved)"
    else:
        # First time: read current Aug 2026 prices, default increases
        start_prices = existing_start_prices(ws_u)
        increases = {y: DEFAULT_PRICE_INCREASE for y in years_after_start}
        source = "umsatz column B"

    refs = write_treiber_price_drivers(ws_t, start_row, years_after_start, start_prices, increases)
    cells = write_umsatz_price_formulas(ws_u, refs, base_year)

    if not dry_run:
        wb.save(path)

    return {
        "years": years,
        "start_prices": start_prices,
        "increases": increases,
        "cells": cells,
        "source": source,
        "refs": refs,
    }


def backup(path: Path) -> Path:
    ts = datetime.now().strftime("%Y%m%d-%H%M%S")
    bk = path.with_name(f"{path.stem}.BACKUP-pre-price-formulas-{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 not all(s in wb_peek.sheetnames for s in ("Umsatzerlöse", "Treiber", "Kunden")):
            print(f"\n  ⨯ skip {path.name}: missing required 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"    Source: {info['source']}")
        print(f"    Startpreise (Aug {min(info['years'])}):")
        for seg, val in info["start_prices"].items():
            print(f"      {seg:13s}: {val:>8.2f} €/Monat")
        print("    Preiserhöhung pro Jahr (Jan jeweils):")
        for yr, inc in info["increases"].items():
            print(f"      {yr}: {inc*100:.1f}%")
        print(f"    Cells written: {info['cells']}")
    return 0


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