"""
engine/matcher.py — tiered one-to-one consumption matching (D5).

A stage matches side A against side B through an ordered tier cascade:
rows matched in a tier are CONSUMED — removed from later tiers — and each
side-B row is used at most once. Every matched A row records which tier
matched it (`match_tier`), giving finance an audit trail of exact vs
fallback vs tolerant matches.

This deliberately replaces the reference implementation's LEFT JOIN +
re-evaluation approach, which allowed one B row to match many A rows
(cartesian duplication patched with dedup afterwards).
"""

import pandas as pd

MATCH_TIER_COL = "match_tier"
_B_CARRY_SUFFIX = "_b"


def run_stage(df_a: pd.DataFrame, df_b: pd.DataFrame, stage: dict, context: dict) -> dict:
    """Execute one stage. Returns dict with keys:
    matched, matched_b, unmatched_a, unmatched_b (DataFrames) and tier_counts.

    match_mode "one_to_one" (default): each A row pairs with at most one B row.
    match_mode "many_to_one" (R7): MANY A rows whose amounts SUM to one B row
    (batch settlements — e.g. one UPI settlement line covering N transactions).
    """
    df_a = apply_filters(df_a, stage.get("side_a", {}).get("filters"))
    df_b = apply_filters(df_b, stage.get("side_b", {}).get("filters"))

    many_to_one = stage.get("match_mode", "one_to_one") == "many_to_one"

    remaining_a = df_a.index
    available_b = df_b.index
    pair_frames = []
    tier_counts = {}

    for tier in stage["tiers"]:
        if not tier.get("enabled", True):
            continue
        if remaining_a.empty or available_b.empty:
            break

        if many_to_one:
            pairs = _match_tier_many_to_one(df_a.loc[remaining_a], df_b.loc[available_b], tier)
        else:
            pairs = _match_tier(df_a.loc[remaining_a], df_b.loc[available_b], tier)
        tier_counts[tier["name"]] = len(pairs)
        if pairs.empty:
            continue

        pairs[MATCH_TIER_COL] = tier["name"]
        pair_frames.append(pairs)
        remaining_a = remaining_a.difference(pairs["a_idx"])
        available_b = available_b.difference(pairs["b_idx"])

    if pair_frames:
        all_pairs = pd.concat(pair_frames, ignore_index=True)
        matched = df_a.loc[all_pairs["a_idx"]].copy()
        matched[MATCH_TIER_COL] = all_pairs[MATCH_TIER_COL].values
        matched = _attach_carry_fields(matched, df_b, all_pairs, stage)
        # Full side-B rows aligned 1:1 with `matched` — persisted as
        # side_b_payload in recon_matches (D8).
        matched_b = df_b.loc[all_pairs["b_idx"]].copy()
        matched_b[MATCH_TIER_COL] = all_pairs[MATCH_TIER_COL].values
    else:
        matched = df_a.iloc[0:0].copy()
        matched[MATCH_TIER_COL] = pd.Series(dtype="object")
        matched_b = df_b.iloc[0:0].copy()
        matched_b[MATCH_TIER_COL] = pd.Series(dtype="object")

    return {
        "matched": matched,
        "matched_b": matched_b,
        "unmatched_a": df_a.loc[remaining_a].copy(),
        "unmatched_b": df_b.loc[available_b].copy(),
        "tier_counts": tier_counts,
    }


def apply_filters(df: pd.DataFrame, filters: list | None) -> pd.DataFrame:
    if not filters:
        return df
    mask = pd.Series(True, index=df.index)
    for f in filters:
        field, op = f["field"], f["op"]
        if field not in df.columns:
            # A filter on an absent optional field excludes nothing for
            # null-checks and everything for value checks would be surprising —
            # treat missing column as all-null.
            col = pd.Series(pd.NA, index=df.index, dtype="object")
        else:
            col = df[field]

        col_str = col.astype("string").str.strip()

        if op == "eq":
            mask &= col_str == str(f["value"])
        elif op == "ne":
            mask &= col_str != str(f["value"])
        elif op == "in":
            mask &= col_str.isin([str(v) for v in f["values"]])
        elif op == "not_in":
            mask &= ~col_str.isin([str(v) for v in f["values"]])
        elif op == "contains":
            mask &= col_str.str.contains(str(f["value"]), case=False, na=False)
        elif op == "not_contains":
            mask &= ~col_str.str.contains(str(f["value"]), case=False, na=False)
        elif op == "not_null":
            mask &= col.notna() & (col_str != "")
        elif op == "is_null":
            mask &= col.isna() | (col_str == "")
        else:
            raise ValueError(f"Unknown filter op: '{op}'")
    return df[mask]


# ── private ──────────────────────────────────────────────────────────────────

def _match_tier(df_a: pd.DataFrame, df_b: pd.DataFrame, tier: dict) -> pd.DataFrame:
    """Return DataFrame with columns a_idx, b_idx — one row per matched pair."""
    if tier.get("amount_tolerance") or tier.get("date_window"):
        return _match_tolerant(df_a, df_b, tier)
    return _match_exact(df_a, df_b, tier)


def _match_exact(df_a: pd.DataFrame, df_b: pd.DataFrame, tier: dict) -> pd.DataFrame:
    keys = tier["keys"]
    a_keys = _key_frame(df_a, [k[0] for k in keys])
    b_keys = _key_frame(df_b, [k[1] for k in keys])
    if a_keys.empty or b_keys.empty:
        return pd.DataFrame(columns=["a_idx", "b_idx"])

    key_cols = [f"k{i}" for i in range(len(keys))]
    a_keys.columns = key_cols
    b_keys.columns = key_cols

    # One-to-one: the n-th A row with a given key combination pairs with the
    # n-th B row with the same combination — no cartesian duplication.
    a_keys["_occ"] = a_keys.groupby(key_cols).cumcount()
    b_keys["_occ"] = b_keys.groupby(key_cols).cumcount()

    merged = pd.merge(
        a_keys.reset_index(names="a_idx"),
        b_keys.reset_index(names="b_idx"),
        on=key_cols + ["_occ"],
        how="inner",
    )
    return merged[["a_idx", "b_idx"]]


def _match_tolerant(df_a: pd.DataFrame, df_b: pd.DataFrame, tier: dict) -> pd.DataFrame:
    """Exact keys produce candidates; tolerance/date-window filters them;
    greedy assignment (smallest amount difference first) makes it one-to-one."""
    keys = tier["keys"]
    a_keys = _key_frame(df_a, [k[0] for k in keys])
    b_keys = _key_frame(df_b, [k[1] for k in keys])
    if a_keys.empty or b_keys.empty:
        return pd.DataFrame(columns=["a_idx", "b_idx"])

    key_cols = [f"k{i}" for i in range(len(keys))]
    a_keys.columns = key_cols
    b_keys.columns = key_cols

    candidates = pd.merge(
        a_keys.reset_index(names="a_idx"),
        b_keys.reset_index(names="b_idx"),
        on=key_cols,
        how="inner",
    )[["a_idx", "b_idx"]]
    if candidates.empty:
        return candidates

    candidates["_diff"] = 0.0

    tol = tier.get("amount_tolerance")
    if tol:
        fa, fb = tol["fields"]
        amt_a = pd.to_numeric(df_a[fa], errors="coerce").reindex(candidates["a_idx"]).values
        amt_b = pd.to_numeric(df_b[fb], errors="coerce").reindex(candidates["b_idx"]).values
        diff = pd.Series(amt_a - amt_b, index=candidates.index).abs()
        if tol["type"] == "abs":
            limit = tol["value"]
            candidates = candidates[diff <= limit]
        else:  # pct — value is a percentage of side A's amount
            limit = pd.Series(amt_a, index=candidates.index).abs() * tol["value"] / 100.0
            candidates = candidates[diff <= limit]
        candidates["_diff"] = diff.loc[candidates.index]

    window = tier.get("date_window")
    if window and not candidates.empty:
        fa, fb = window["fields"]
        date_a = pd.to_datetime(df_a[fa], errors="coerce").reindex(candidates["a_idx"]).values
        date_b = pd.to_datetime(df_b[fb], errors="coerce").reindex(candidates["b_idx"]).values
        day_diff = pd.Series((date_a - date_b), index=candidates.index).abs()
        candidates = candidates[day_diff <= pd.Timedelta(days=window["days"])]

    if candidates.empty:
        return pd.DataFrame(columns=["a_idx", "b_idx"])

    # Greedy one-to-one: best (smallest difference) pairs win.
    candidates = candidates.sort_values(["_diff", "a_idx"])
    taken_a, taken_b, chosen = set(), set(), []
    for row in candidates.itertuples(index=False):
        if row.a_idx in taken_a or row.b_idx in taken_b:
            continue
        taken_a.add(row.a_idx)
        taken_b.add(row.b_idx)
        chosen.append((row.a_idx, row.b_idx))
    return pd.DataFrame(chosen, columns=["a_idx", "b_idx"])


def _match_tier_many_to_one(df_a: pd.DataFrame, df_b: pd.DataFrame, tier: dict) -> pd.DataFrame:
    """Batch-settlement matching: group A rows by the tier keys, sum their
    amounts, and match each GROUP to one B row with the same keys whose amount
    equals the group total (within amount_tolerance if configured).

    Returned pairs repeat the B row index for every A row in a matched group,
    keeping the 1:1 alignment downstream code expects."""
    keys = tier["keys"]
    a_keys = _key_frame(df_a, [k[0] for k in keys])
    b_keys = _key_frame(df_b, [k[1] for k in keys])
    if a_keys.empty or b_keys.empty:
        return pd.DataFrame(columns=["a_idx", "b_idx"])

    key_cols = [f"k{i}" for i in range(len(keys))]
    a_keys.columns = key_cols
    b_keys.columns = key_cols

    tol = tier.get("amount_tolerance") or {}
    field_a, field_b = tol.get("fields", ["amount", "amount"])

    a_amounts = pd.to_numeric(df_a[field_a], errors="coerce").reindex(a_keys.index).fillna(0)
    group_sums = (
        a_keys.assign(group_amount=a_amounts, member_idx=a_keys.index)
        .groupby(key_cols)
        .agg(group_total=("group_amount", "sum"), members=("member_idx", list))
        .reset_index()
    )

    b_amounts = pd.to_numeric(df_b[field_b], errors="coerce").reindex(b_keys.index)
    b_candidates = b_keys.assign(b_amount=b_amounts).reset_index(names="b_idx")

    candidates = pd.merge(group_sums, b_candidates, on=key_cols, how="inner")
    if candidates.empty:
        return pd.DataFrame(columns=["a_idx", "b_idx"])

    diff = (candidates["group_total"] - candidates["b_amount"]).abs()
    if tol.get("type") == "pct":
        limit = candidates["group_total"].abs() * tol["value"] / 100.0
    elif tol.get("type") == "abs":
        limit = tol["value"]
    else:
        limit = 0.005  # exact to the paisa
    candidates = candidates.assign(diff=diff)[diff <= limit]
    if candidates.empty:
        return pd.DataFrame(columns=["a_idx", "b_idx"])

    # Greedy one group ↔ one B row, best (smallest difference) first.
    candidates = candidates.sort_values(["diff", "b_idx"])
    taken_groups, taken_b, pairs = set(), set(), []
    for _, row in candidates.iterrows():
        group_key = tuple(row[c] for c in key_cols)
        if group_key in taken_groups or row["b_idx"] in taken_b:
            continue
        taken_groups.add(group_key)
        taken_b.add(row["b_idx"])
        for a_idx in row["members"]:
            pairs.append((a_idx, row["b_idx"]))

    return pd.DataFrame(pairs, columns=["a_idx", "b_idx"])


def _key_frame(df: pd.DataFrame, fields: list) -> pd.DataFrame:
    """Build comparable key columns; rows with any missing key are excluded
    (a null key never matches anything)."""
    missing = [f for f in fields if f not in df.columns]
    if missing:
        raise ValueError(f"Match keys reference missing canonical fields: {missing}")

    out = pd.DataFrame(index=df.index)
    for i, field in enumerate(fields):
        out[f"k{i}"] = _key_series(df[field])
    return out.dropna()


def _key_series(series: pd.Series) -> pd.Series:
    """Canonical comparable representation. Every key becomes a STRING so the
    two sides always merge regardless of source dtype:
    - datetimes → 'YYYY-MM-DD'
    - numeric-looking values → integer paise ('500', 500, '500.00' → '50000'),
      avoiding float-equality traps; values ≥ 1e13 (full card/account numbers)
      stay as digit strings to dodge float precision loss
    - everything else → stripped string; empties/nan → missing (never match)
    """
    if pd.api.types.is_datetime64_any_dtype(series):
        s = series.dt.strftime("%Y-%m-%d").astype("string")
        return s.where(series.notna(), pd.NA)

    s = (
        series.astype("string")
        .str.strip()
        .str.replace(r"\.0$", "", regex=True)
    )
    s = s.where(s.notna() & (s != "") & (s.str.lower() != "nan"), pd.NA)

    num = pd.to_numeric(s, errors="coerce")
    as_paise = num.notna() & (num.abs() < 1e13)
    if as_paise.any():
        s = s.copy()
        s[as_paise] = (num[as_paise] * 100).round().astype("int64").astype(str)
    return s


def _attach_carry_fields(matched: pd.DataFrame, df_b: pd.DataFrame,
                         pairs: pd.DataFrame, stage: dict) -> pd.DataFrame:
    carry = stage.get("output", {}).get("carry_fields", [])
    for field in carry:
        if field not in df_b.columns:
            continue
        name = field if field not in matched.columns else field + _B_CARRY_SUFFIX
        matched[name] = df_b[field].reindex(pairs["b_idx"]).values
    return matched