"""
Recovers the premium breakdown the original policy transform dropped.

`transform_policies.py` modeled prima neta, derecho de póliza and comisión and
nothing else, which lost two things from every migrated policy:

  1. RECARGO — the financing surcharge on a policy paid in more than one
     exhibición, plus the whole second money row (`p_neta_2`, `recargo_2`,
     `d_pol_2`, `com_2`) that a semestral policy carries because each payment
     is priced separately. These were not deleted, they were swept into
     `policies.coveragesJson` as loose strings alongside the real coverages —
     unqueryable, and mislabeled as coverage amounts.

  2. FORMA PAGO — dropped outright. The column was marked "consumed" by the
     transform's coverage sweep but never written to any column, so it exists
     nowhere in the platform database. It is the field that decides whether a
     recargo is legitimate on a row at all, so it cannot be inferred back from
     the money.

The transform has been fixed in the same commit, so a full `run_all.py` now
produces all of this directly. This script exists for a database that must not
be re-imported: it reads the same staged Parquet and patches in place.

What it does NOT do: invent IVA or the printed TOTAL. Those were never columns
in the home tables — they were unbound calculated controls on the Access form —
so there is genuinely nothing to recover, and both stay null until a human
edits the policy. The app computes them (apps/api/src/policies/premium.ts).

Idempotent, and never overwrites a non-null value: a figure a human has since
corrected in the app wins over the legacy one.

    ./.venv/bin/python backfill_policy_premium_breakdown.py --env dev
"""

from __future__ import annotations

import json
import sys
from decimal import Decimal, InvalidOperation
from pathlib import Path

import pandas as pd

from dbenv import connect
from sync import parse_mode

STG = Path(__file__).parent / "output" / "stg_seguros"
LEGACY_DB = "SEGUROS 16_be"
NULL = "∅"

# Legacy column -> what it is, per source table. Only the tables that actually
# carry a breakdown appear; the auto tables have no recargo and no second row.
HOME_TABLES = ("mult", "incendio", "m_empr")

# Keys the transform used to dump into coveragesJson that are now real columns.
# Stripped once migrated so the blob stops pretending they are coverages.
MIGRATED_COVERAGE_KEYS = (
    "recargo", "recargo_2", "p_neta_2", "d_pol_2", "com_2",
)

_FREQ = {
    "ANNUAL": "ANNUAL",
    "ANUAL": "ANNUAL",
    "SEMESTRAL": "SEMIANNUAL",
    "TRIMESTRAL": "QUARTERLY",
    "MENSUAL": "MONTHLY",
    "CONTADO": "SINGLE",
}


def s(v):
    if v is None or pd.isna(v):
        return None
    v = str(v).strip()
    return None if v in ("", NULL) else v


def dec(v):
    v = s(v)
    if v is None:
        return None
    try:
        return Decimal(v.replace(",", ""))
    except (InvalidOperation, ValueError):
        return None


def freq(v):
    return _FREQ.get((s(v) or "").upper())


def load(name):
    df = pd.read_parquet(STG / f"{name}.parquet").sort_values("_row_num").reset_index(drop=True)
    for c in df.columns:
        if c != "_row_num":
            df[c] = df[c].astype("string").str.strip()
    return df


def main():
    env, _sync_mode = parse_mode()

    # Fails closed rather than reporting a clean run over nothing: an empty
    # staging directory and a policy set with no recargo look identical from
    # the database side, and "0 rows updated" would read as success.
    if not STG.exists():
        print(f"[policy-premium] staged Parquet missing at {STG} — run extract/load first.")
        return 3

    conn = connect(env)
    c = conn.cursor()
    print(f"[policy-premium] target env: {env}")

    # Every policy that came from the insurance ETL, keyed by provenance. The
    # id is needed to reach the installments, coveragesJson to strip the keys.
    c.execute(
        "SELECT legacySourceTable, legacyId, id, coveragesJson "
        "FROM policies WHERE legacySourceDb = %s AND legacyId IS NOT NULL",
        (LEGACY_DB,),
    )
    by_key = {(t, lid): (pid, cov) for t, lid, pid, cov in c.fetchall()}
    print(f"  {len(by_key)} migrated polic(ies) in the target database")

    pol_updates = []       # (surcharge, paymentFrequency, coveragesJson, policyId)
    inst_updates = []      # (netPremium, surcharge, policyFee, commission, policyId, seq)
    seen_tables = 0

    for table in HOME_TABLES + (
        "tabla_autos", "tabla_autos_ampl", "tabla_autos_limit",
        "tabla_autos_ampl_r", "tabla_autos_rc_r", "mca2", "licencias",
    ):
        path = STG / f"{table}.parquet"
        if not path.exists():
            continue
        seen_tables += 1
        df = load(table)
        home = table in HOME_TABLES

        for _, row in df.iterrows():
            key = (table, str(int(row["_row_num"])))
            hit = by_key.get(key)
            if not hit:
                continue
            pid, cov_raw = hit

            surcharge = dec(row.get("recargo")) if home else None
            frequency = freq(row.get("forma_pago"))

            # Strip the now-modeled keys out of the coverage blob. Rewritten
            # only when something actually changes, so a policy whose blob a
            # human has edited is left byte-identical.
            cov_new = None
            if cov_raw:
                try:
                    cov = json.loads(cov_raw) if isinstance(cov_raw, str) else cov_raw
                except (TypeError, ValueError):
                    cov = None
                if isinstance(cov, dict):
                    kept = {k: v for k, v in cov.items() if k not in MIGRATED_COVERAGE_KEYS}
                    if len(kept) != len(cov):
                        cov_new = json.dumps(kept, ensure_ascii=False) if kept else None

            if surcharge is not None or frequency is not None or cov_new is not None:
                pol_updates.append((surcharge, frequency, cov_new, cov_new is not None, pid))

            # Per-payment breakdown. Slot 1 is the unsuffixed money row, slot 2
            # the _2 twin; the auto tables have a single slot and no recargo.
            if home:
                slots = [
                    (1, "p_neta", "recargo", "d_pol", "com"),
                    (2, "p_neta_2", "recargo_2", "d_pol_2", "com_2"),
                ]
            else:
                pn = "prima1" if table == "mca2" else "prima_neta"
                dp = "d_poliza1" if table == "mca2" else "d_poliza"
                slots = [(1, pn, None, dp, None)]

            for seq, pn, rc, dp, cm in slots:
                vals = (
                    dec(row.get(pn)) if pn else None,
                    dec(row.get(rc)) if rc else None,
                    dec(row.get(dp)) if dp else None,
                    dec(row.get(cm)) if cm else None,
                )
                if all(v is None for v in vals):
                    continue
                inst_updates.append((*vals, pid, seq))

    if not seen_tables:
        print(f"[policy-premium] no policy tables staged under {STG} — nothing to do.")
        return 3

    # COALESCE on every target: a column a human has already filled in the app
    # keeps its value, the legacy figure only lands where there is a hole.
    for surcharge, frequency, cov_new, rewrite_cov, pid in pol_updates:
        c.execute(
            "UPDATE policies SET "
            "  surcharge = COALESCE(surcharge, %s), "
            "  paymentFrequency = COALESCE(paymentFrequency, %s), "
            "  coveragesJson = IF(%s, %s, coveragesJson) "
            "WHERE id = %s",
            (surcharge, frequency, 1 if rewrite_cov else 0, cov_new, pid),
        )

    for netp, surch, fee, comm, pid, seq in inst_updates:
        c.execute(
            "UPDATE policy_payment_installments SET "
            "  netPremium = COALESCE(netPremium, %s), "
            "  surcharge  = COALESCE(surcharge, %s), "
            "  policyFee  = COALESCE(policyFee, %s), "
            "  commission = COALESCE(commission, %s) "
            "WHERE policyId = %s AND sequence = %s",
            (netp, surch, fee, comm, pid, seq),
        )

    conn.commit()
    print(f"  policies    : {len(pol_updates)} row(s) touched")
    print(f"  installments: {len(inst_updates)} row(s) touched")

    # --- validation ---------------------------------------------------------
    c.execute("SELECT COUNT(*) FROM policies WHERE surcharge IS NOT NULL AND surcharge <> 0")
    n_surch = c.fetchone()[0]
    c.execute("SELECT COUNT(*) FROM policies WHERE paymentFrequency IS NOT NULL")
    n_freq = c.fetchone()[0]
    c.execute(
        "SELECT COUNT(*) FROM policies "
        "WHERE paymentFrequency IN ('ANNUAL','SINGLE') AND surcharge IS NOT NULL AND surcharge <> 0"
    )
    n_bad = c.fetchone()[0]
    c.execute(
        "SELECT COUNT(*) FROM policies WHERE coveragesJson IS NOT NULL "
        "AND JSON_EXTRACT(coveragesJson, '$.recargo') IS NOT NULL"
    )
    n_left = c.fetchone()[0]

    print(f"  -> policies with a recargo    : {n_surch}")
    print(f"  -> policies with a forma pago : {n_freq}")
    print(f"  -> recargo still in coverages : {n_left}")

    # A surcharge on an annual policy contradicts the rule the capture form
    # enforces, so it is worth surfacing rather than leaving for someone to
    # find in a total. It is a warning, not a failure: the books are the books.
    if n_bad:
        print(f"  !! {n_bad} annual/contado polic(ies) carry a non-zero recargo — review by hand")
    return 0


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