#!/usr/bin/env python3
"""Analyze ToE lambda output XLSX for calculation correctness.

Reads the Adjustments and Summary sheets, verifies each row against the
calculation rules (percent / flat_rate, negative flag, projected balance),
and appends an "Analysis" column to the Adjustments sheet.

Usage:
    python scripts/analyze_output.py [path_to_xlsx]

Defaults to ./ToE-lambda-output.xlsx if no path is given.
"""

from __future__ import annotations

import sys
from dataclasses import dataclass

import openpyxl

DEFAULT_OUTPUT = 'ToE-lambda-output.xlsx'


def _to_float(val: object) -> float | None:
    """Coerce a cell value to float; returns None for non-numeric/missing values (e.g. 'N/A')."""
    if isinstance(val, (int, float)):
        return float(val)
    return None


@dataclass
class SummaryRow:
    """A row from the Summary sheet of the ToE output workbook."""

    et_id: int
    transfer_type: str
    rate_type: str
    rate_amount: object
    input: str
    negative: bool
    from_account: str
    from_contract_id: int
    to_account: str
    to_contract_id: int
    currency: str
    closing_balance: float | None
    net_revenue: float | None
    gross_revenue: float | None
    to_closing_balance: float | None
    to_net_revenue: float | None
    to_gross_revenue: float | None
    selected_balance: float | None
    calculated_amount: float | None
    error: str | None

    def from_input_balance(self) -> float | None:
        """Return the from-contract balance for the configured input field."""
        if self.input == 'net_revenue':
            return self.net_revenue
        if self.input == 'gross_revenue':
            return self.gross_revenue
        return self.closing_balance

    def to_input_balance(self) -> float | None:
        """Return the to-contract balance for the configured input field."""
        if self.input == 'net_revenue':
            return self.to_net_revenue
        if self.input == 'gross_revenue':
            return self.to_gross_revenue
        return self.to_closing_balance


def load_summary(ws) -> dict[int, SummaryRow]:
    """Parse the Summary sheet into a dict keyed by earnings_transfer_id."""
    rows: dict[int, SummaryRow] = {}
    for row in ws.iter_rows(min_row=2, values_only=True):
        if row[0] is None:
            continue
        s = SummaryRow(
            et_id=row[0],
            transfer_type=row[1],
            rate_type=row[2],
            rate_amount=row[3],
            input=row[4],
            negative=row[5],
            from_account=row[6],
            from_contract_id=row[7],
            to_account=row[8],
            to_contract_id=row[9],
            currency=row[10],
            closing_balance=_to_float(row[11]),
            net_revenue=_to_float(row[12]),
            gross_revenue=_to_float(row[13]),
            to_closing_balance=_to_float(row[14]),
            to_net_revenue=_to_float(row[15]),
            to_gross_revenue=_to_float(row[16]),
            selected_balance=_to_float(row[17]),
            calculated_amount=_to_float(row[18]),
            error=row[21] if len(row) > 21 and row[21] else None,
        )
        rows[s.et_id] = s
    return rows


def _rate_as_float(rate_amount: object) -> float | None:
    """Coerce a rate_amount cell (may be '25%' or numeric) to a decimal."""
    if isinstance(rate_amount, str):
        try:
            return float(rate_amount.replace('%', '')) / 100
        except ValueError:
            return None
    return _to_float(rate_amount)


def expected_amount(s: SummaryRow) -> float | None:
    """Recompute the expected calculated amount from the summary data."""
    if s.rate_type == 'flat_rate':
        amt = _to_float(s.rate_amount)
        if amt is None:
            amt = 0.0
        if s.negative:
            return round(amt, 2)
        selected_balance = 0.0 if s.selected_balance is None else s.selected_balance
        if selected_balance < 0:
            return 0.0
        return round(amt, 2)
    # percent
    selected_balance = 0.0 if s.selected_balance is None else s.selected_balance
    if selected_balance <= 0:
        return 0.0
    pct = _rate_as_float(s.rate_amount)
    if pct is None:
        return None
    return round(selected_balance * pct, 2)


def _verify_percent_amount(s: SummaryRow) -> str | None:
    """Validate a percent row's calculated_amount against the displayed rate.

    Returns an issue string if the amount looks wrong, or None if it's consistent.

    When multiple transfers draw from the same from-contract, the engine
    cascades: each transfer operates on the balance remaining after prior
    transfers.  ``selected_balance`` is a pre-transfer snapshot, so
    ``calc / selected_balance`` will be less than the nominal rate for any
    transfer that isn't first in the chain.  Instead we derive the effective
    balance the engine must have used (``calc / rate``) and verify it is
    positive.  For ``closing_balance`` inputs the effective balance can
    legitimately exceed the snapshot when prior cascade transfers include
    net inflows, so we only check for negative values.  For revenue inputs
    (``net_revenue`` / ``gross_revenue``) there is no cascade, so we also
    verify the effective balance does not exceed the snapshot.
    """
    bal = s.selected_balance
    calc = s.calculated_amount
    if bal is None or calc is None:
        return None

    # The engine blocks percent transfers when closing_balance < 0 and
    # negative=False, regardless of which input source is selected.
    if not s.negative and s.closing_balance is not None and s.closing_balance < 0:
        if round(calc, 2) != 0:
            return (
                f'calculated_amount {calc} is non-zero '
                f'but closing_balance {s.closing_balance} is negative '
                f'and negative=False'
            )
        return None

    if bal <= 0:
        if round(calc, 2) != 0:
            return (
                f'calculated_amount {calc} is non-zero '
                f'but selected_balance {bal} is non-positive'
            )
        return None

    displayed_pct = _rate_as_float(s.rate_amount)
    if displayed_pct is None or displayed_pct == 0:
        return None
    if displayed_pct == 0:
        # The displayed rate rounds to 0% so we can't derive an effective
        # balance.  Sanity-check that calc isn't negative, and for revenue
        # inputs (no cascade) verify the upper bound using the maximum rate
        # that would still display as 0%.
        if round(calc, 2) < 0:
            return (
                f'calculated_amount {calc} is negative '
                f'for a rate that displays as 0% ({s.rate_amount})'
            )
        if s.input in ('net_revenue', 'gross_revenue'):
            max_rate = 5e-7
            max_calc = round(bal * max_rate, 2) + 0.005
            if round(calc, 2) > max_calc:
                return (
                    f'calculated_amount {calc} exceeds upper bound {max_calc} '
                    f'for a rate that displays as 0% ({s.rate_amount}) '
                    f'with selected_balance={bal} (input={s.input} does not cascade)'
                )
        return None

    # Effective balance the engine used: calc = effective_bal * rate
    effective_bal = calc / displayed_pct

    # Tolerance must account for two rounding sources:
    # 1. calculated_amount is rounded to cents, so effective_bal has
    #    uncertainty of ±0.005 / rate.
    # 2. writer._pct rounds the displayed rate to 4 decimal places of
    #    percent (6 dp as a ratio, i.e. ±5e-7).  This shifts the implied
    #    effective_bal by up to bal * (5e-7 / rate).
    cent_tol = 0.005 / displayed_pct
    rate_display_tol = bal * 5e-7 / displayed_pct
    rounding_tol = cent_tol + rate_display_tol
    if effective_bal < -rounding_tol:
        return (
            f'calculated_amount {calc} / rate {s.rate_amount} implies '
            f'negative effective balance {effective_bal:.2f}'
        )

    # Revenue inputs don't cascade, so effective_bal should match the snapshot.
    if (
        s.input in ('net_revenue', 'gross_revenue')
        and effective_bal > bal + rounding_tol
    ):
        return (
            f'calculated_amount {calc} / rate {s.rate_amount} implies '
            f'effective balance {effective_bal:.2f} > selected_balance {bal} '
            f'(input={s.input} does not cascade)'
        )

    return None


def analyze_row(
    s: SummaryRow, side: str, amount: object, projected_balance: object
) -> str:
    """Verify a single FROM or TO adjustment row against the summary data."""
    amt = _to_float(amount)
    proj = _to_float(projected_balance)
    issues: list[str] = []
    skips: list[str] = []
    calc = s.calculated_amount

    if calc is None:
        skips.append('calculated_amount missing in Summary')
    elif s.rate_type == 'flat_rate':
        exp_calc = expected_amount(s)
        if exp_calc is None:
            skips.append(
                'cannot compute expected amount (missing rate or balance data)'
            )
        elif round(calc, 2) != round(exp_calc, 2):
            issues.append(f'calculated_amount {calc} != expected {exp_calc}')
    else:
        pct_issue = _verify_percent_amount(s)
        if pct_issue is not None:
            issues.append(pct_issue)

    if side == 'FROM':
        if amt is None:
            skips.append('FROM amount missing')
        elif calc is None:
            skips.append('cannot check FROM amount (calculated_amount missing)')
        elif round(amt, 2) != round(-calc, 2):
            issues.append(f'FROM amount {amt} != expected {-calc}')

        base = s.from_input_balance()
        if proj is None:
            skips.append('projected_balance missing')
        elif base is None:
            skips.append(f'cannot check projected_balance ({s.input} balance missing)')
        elif calc is None:
            skips.append('cannot check projected_balance (calculated_amount missing)')
        else:
            exp_proj = round(base - calc, 2)
            if round(proj, 2) != exp_proj:
                issues.append(
                    f'projected_balance {projected_balance} != expected {exp_proj} '
                    f'({s.input}={base} - calc={calc})'
                )
    elif side == 'TO':
        if amt is None:
            skips.append('TO amount missing')
        elif calc is None:
            skips.append('cannot check TO amount (calculated_amount missing)')
        elif round(amt, 2) != round(calc, 2):
            issues.append(f'TO amount {amt} != expected {calc}')

        base = s.to_input_balance()
        if proj is None:
            skips.append('projected_balance missing')
        elif base is None:
            skips.append(
                f'cannot check projected_balance (to_{s.input} balance missing)'
            )
        elif calc is None:
            skips.append('cannot check projected_balance (calculated_amount missing)')
        else:
            exp_proj = round(base + calc, 2)
            if round(proj, 2) != exp_proj:
                issues.append(
                    f'projected_balance {proj} != expected {exp_proj} '
                    f'(to_{s.input}={base} + calc={calc})'
                )
    else:
        issues.append(f'unexpected Side value: {side!r}')

    pct = _rate_as_float(s.rate_amount)
    rate_desc = (
        s.rate_amount
        if isinstance(s.rate_amount, str)
        else (
            f'${s.rate_amount}'
            if s.rate_type == 'flat_rate'
            else f'{pct * 100 if pct is not None else "?"}%'
        )
    )
    skip_suffix = f'; SKIP: {"; ".join(skips)}' if skips else ''
    if issues:
        return f'ISSUE: {"; ".join(issues)}{skip_suffix}'
    if skips:
        return f'SKIP: {"; ".join(skips)}'
    return (
        f'CORRECT: {s.rate_type} {rate_desc} of {s.input}={s.selected_balance}; '
        f'calc={calc}; {side} amt={amt}, proj_bal={proj}'
    )


def _build_cascade_prefix_sums(
    summary: dict[int, SummaryRow],
) -> dict[int, float]:
    """Precompute per-ET-ID the net amount already transferred from its from_contract.

    Returns a dict mapping et_id → net prior outflow from that record's
    from_contract_id, mirroring the engine's cascade logic in O(n log n).
    """
    from collections import defaultdict

    running: defaultdict[int, float] = defaultdict(float)
    result: dict[int, float] = {}
    for et_id in sorted(summary):
        s = summary[et_id]
        result[et_id] = (
            running[s.from_contract_id] if s.from_contract_id is not None else 0.0
        )
        calc = s.calculated_amount or 0.0
        if s.from_contract_id is not None:
            running[s.from_contract_id] += calc
        if s.to_contract_id is not None:
            running[s.to_contract_id] -= calc
    return result


def _blocked_reason(s: SummaryRow, prior_transferred: float) -> str:
    """Determine why a record with calculated_amount=0 was blocked."""
    # Engine-flagged errors (e.g. allocation validation, missing Snowflake data)
    if s.error:
        return f'engine error: {s.error}'

    bal = s.selected_balance
    if bal is None:
        return 'SKIP: selected_balance missing, cannot verify block reason'

    if s.rate_type == 'flat_rate':
        rate = _to_float(s.rate_amount)
        if rate is None:
            return 'flat_rate blocked: rate_amount missing/NA, engine treats as 0.0'
        if round(rate, 2) == 0:
            return f'flat_rate rounds to zero: {s.rate_amount} -> {round(rate, 2)}'
        # _can_apply_flat_rate gates on closing_balance (not selected_balance),
        # and cascades prior transfers regardless of the input field.
        if not s.negative and s.closing_balance:
            effective_cb = s.closing_balance - prior_transferred
            if effective_cb < 0:
                return (
                    f'flat_rate blocked: closing_balance {effective_cb:.2f} < 0 '
                    f'(snapshot={s.closing_balance}, prior={prior_transferred:.2f})'
                )
        return 'UNEXPECTED: flat_rate should have produced a non-zero amount'

    # percent: closing-balance gate (negative=False blocks on negative closing balance)
    if not s.negative and s.closing_balance is not None and s.closing_balance < 0:
        return (
            f'percent blocked: closing_balance={s.closing_balance} < 0, negative=false'
        )

    # For closing_balance input, the engine cascades: subtract prior outflows.
    effective_bal = bal
    if s.input == 'closing_balance':
        effective_bal = bal - prior_transferred
        if effective_bal <= 0:
            return (
                f'percent blocked: cascade-depleted balance {effective_bal:.2f} '
                f'(snapshot={bal}, prior_transfers={prior_transferred:.2f}) is non-positive'
            )

    if effective_bal <= 0:
        return f'percent blocked: {s.input}={bal} is non-positive'
    pct = _rate_as_float(s.rate_amount)
    if pct is None:
        return 'percent blocked: rate_amount missing/NA, engine returns 0.0'
    product = round(effective_bal * pct, 2)
    if product == 0:
        return (
            f'percent rounds to zero: {effective_bal} x {pct} = {effective_bal * pct}'
        )
    return f'UNEXPECTED: should have calculated {product}'


def analyze_blocked(summary: dict[int, SummaryRow]) -> list[str]:
    """Report on records that were correctly blocked (calculated_amount=0)."""
    cascade = _build_cascade_prefix_sums(summary)
    lines = []
    for et_id in sorted(summary):
        s = summary[et_id]
        if s.calculated_amount is None or round(s.calculated_amount, 2) != 0:
            continue
        reason = _blocked_reason(s, cascade[et_id])
        status = 'CORRECT' if not reason.startswith('UNEXPECTED') else 'ERROR'
        lines.append(f'ET{et_id}: {reason} -> {status}')
    return lines


def main() -> None:
    """Load the output workbook, verify all rows, and write the Analysis column."""
    path = sys.argv[1] if len(sys.argv) > 1 else DEFAULT_OUTPUT
    wb = openpyxl.load_workbook(path)

    if 'Summary' not in wb.sheetnames or 'Adjustments' not in wb.sheetnames:
        print(
            "ERROR: Workbook must have 'Summary' and 'Adjustments' sheets",
            file=sys.stderr,
        )
        sys.exit(1)

    summary = load_summary(wb['Summary'])
    ws_adj = wb['Adjustments']

    headers = [cell.value for cell in ws_adj[1]]
    col = {h: i for i, h in enumerate(headers)}

    required = {'Earnings Transfer ID', 'Side', 'Amount', 'Projected Balance'}
    missing = required - col.keys()
    if missing:
        print(
            f'ERROR: Adjustments sheet missing required headers: {", ".join(sorted(missing))}',
            file=sys.stderr,
        )
        sys.exit(1)

    analysis_col = (col['Analysis'] + 1) if 'Analysis' in col else len(headers) + 1
    ws_adj.cell(row=1, column=analysis_col, value='Analysis')

    row_issue_count = 0
    checked = 0
    transfer_issues: dict[
        tuple[int, str], str
    ] = {}  # (et_id, issue_kind) -> first message
    for row_idx in range(2, ws_adj.max_row + 1):
        et_id = ws_adj.cell(row=row_idx, column=col['Earnings Transfer ID'] + 1).value
        if et_id is None:
            continue
        side = ws_adj.cell(row=row_idx, column=col['Side'] + 1).value
        amount = ws_adj.cell(row=row_idx, column=col['Amount'] + 1).value
        proj_bal = ws_adj.cell(row=row_idx, column=col['Projected Balance'] + 1).value

        s = summary.get(et_id)
        if not s:
            ws_adj.cell(
                row=row_idx,
                column=analysis_col,
                value=f'ISSUE: ET {et_id} missing from Summary',
            )
            row_issue_count += 1
            transfer_issues.setdefault(
                (et_id, 'missing_from_summary'), f'ET {et_id} missing from Summary'
            )
            continue

        analysis = analyze_row(s, side, amount, proj_bal)
        ws_adj.cell(row=row_idx, column=analysis_col, value=analysis)
        checked += 1
        if analysis.startswith('ISSUE'):
            row_issue_count += 1
            # Extract issue kind from "ISSUE: <kind>: ..." or use full message
            kind = analysis.split(':')[1].strip() if ':' in analysis[6:] else analysis
            transfer_issues.setdefault((et_id, kind), analysis)
            print(f'  !! Row {row_idx}: ET{et_id} {side} — {analysis}')

    # Report blocked transfers
    blocked = analyze_blocked(summary)
    blocked_errors = [line for line in blocked if 'UNEXPECTED' in line]

    unique_transfers = len({et_id for (et_id, _) in transfer_issues})
    print(f'\nAdjustment rows checked: {checked}')
    print(f'Row-level issues: {row_issue_count}')
    print(
        f'Unique issues: {len(transfer_issues)} across {unique_transfers} transfer(s)'
    )
    print(f'Blocked transfers (calc=0): {len(blocked)}')
    if blocked_errors:
        print(f'Unexpected blocked: {len(blocked_errors)}')
        for line in blocked_errors:
            print(f'  !! {line}')
    else:
        print('All blocked transfers correctly gated.')

    wb.save(path)
    print(f'\nAnalysis column written to {path}')

    if transfer_issues or blocked_errors:
        sys.exit(1)


if __name__ == '__main__':
    main()