"""Two-source data reconciliation.

Given two DataFrames (typically a bank/credit-card feed and a ledger
export), find which rows on the left correspond to rows on the right
based on amount, date, and optional reference/description fields.

Output buckets:
    matched          — one row per accepted pair, with both originals.
    unmatched_left   — left rows with no acceptable right counterpart.
    unmatched_right  — right rows with no acceptable left counterpart.
    review           — ambiguous cases (a left row had >1 equally good
                       right candidates, or vice versa) surfaced for the
                       user to disambiguate manually.

Matching strategy is a multi-pass greedy one-to-one assignment:
    Pass 1: exact key match (when ``key_columns`` is set on either side)
    Pass 2: exact (amount, date) match
    Pass 3: amount within tolerance AND date within window
    Pass 4: + optional description fuzzy similarity boost

Within each pass, candidate pairs are scored and assigned greedily by
descending score; ties for the same left row that span multiple right
rows (or vice versa) are sent to ``review`` instead of being matched
arbitrarily.

The module is pure: no I/O, no Streamlit, no logging side effects beyond
loguru. Caller drives file reading and result rendering.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Optional

import pandas as pd
from loguru import logger

try:
    from rapidfuzz import fuzz as _rf_fuzz
    _HAS_RAPIDFUZZ = True
except ImportError:  # pragma: no cover — rapidfuzz is in requirements.txt
    _HAS_RAPIDFUZZ = False


# ---------------------------------------------------------------------------
# Options & result
# ---------------------------------------------------------------------------


@dataclass
class ReconcileOptions:
    """Configuration for :func:`reconcile`.

    ``left_amount`` / ``right_amount`` are required: every match needs
    an amount to anchor on. Everything else is optional.
    """

    # Amount columns (required). Values are coerced to float; non-numeric
    # rows are dropped from matching but appear in the unmatched buckets.
    left_amount: str = ""
    right_amount: str = ""

    # Date columns. When both are set, candidates must fall within
    # ``date_tolerance_days``. When unset, date is ignored entirely.
    left_date: Optional[str] = None
    right_date: Optional[str] = None

    # Optional reference / key columns for exact-match Pass 1. List
    # forms must be the same length so the i-th left key pairs with the
    # i-th right key (e.g. ``["check_no"]`` ↔ ``["ref"]``).
    left_keys: list[str] = field(default_factory=list)
    right_keys: list[str] = field(default_factory=list)

    # Description columns for fuzzy similarity boost (optional). Only
    # used when ``desc_min_score`` > 0 AND rapidfuzz is installed.
    left_desc: Optional[str] = None
    right_desc: Optional[str] = None
    desc_min_score: int = 0  # 0–100; 0 disables fuzzy.

    # Tolerances. Defaults are exact match.
    amount_tolerance: float = 0.0   # absolute (e.g. 0.01 for cent rounding)
    date_tolerance_days: int = 0    # ± N calendar days

    # Some bank feeds use opposite sign convention from the ledger
    # (debits positive vs. negative). Flipping this multiplies the
    # right side's amount by -1 before matching.
    invert_right_sign: bool = False


@dataclass
class ReconcileResult:
    """Outcome of a reconcile run.

    All four DataFrames preserve the original columns from each side,
    prefixed with ``left_`` and ``right_`` where applicable, plus a
    small set of bookkeeping columns (``match_pass``, ``amount_diff``,
    ``date_diff_days``, ``desc_score``).
    """

    matched: pd.DataFrame
    unmatched_left: pd.DataFrame
    unmatched_right: pd.DataFrame
    review: pd.DataFrame
    stats: dict[str, int] = field(default_factory=dict)


# ---------------------------------------------------------------------------
# Public entry point
# ---------------------------------------------------------------------------


def reconcile(
    left: pd.DataFrame,
    right: pd.DataFrame,
    options: ReconcileOptions,
) -> ReconcileResult:
    """Reconcile *left* against *right* using *options*.

    Neither input is mutated. The result's frames hold copies of the
    relevant rows from the originals, joined via the bookkeeping
    columns described on :class:`ReconcileResult`.
    """
    _validate_options(left, right, options)

    # Normalize amounts and dates to typed columns we can reason about
    # without re-parsing per pass. The work columns live on copies so
    # callers' inputs are untouched.
    L = _prep_side(left, options, side="left")
    R = _prep_side(right, options, side="right")

    # Track which left/right indices remain unmatched across passes.
    # Seeded from the FULL input frame, not the prepped one — rows
    # dropped during prep (unparseable amount/date) must still surface
    # in the unmatched bucket so users can see they exist. Candidate
    # generators iterate L.index, so prep-dropped rows simply never
    # get claimed.
    left_open: set = set(left.index)
    right_open: set = set(right.index)

    matched_pairs: list[dict] = []
    review_pairs: list[dict] = []

    # Pass 1 — exact key match on user-supplied reference columns.
    if options.left_keys and options.right_keys:
        _run_pass(
            L, R, left_open, right_open, matched_pairs, review_pairs,
            options=options, pass_name="key",
            candidate_fn=_candidates_by_key,
        )

    # Pass 2 — exact (amount, date) match.
    _run_pass(
        L, R, left_open, right_open, matched_pairs, review_pairs,
        options=options, pass_name="exact",
        candidate_fn=_candidates_exact,
    )

    # Pass 3 — tolerance-window match.
    if options.amount_tolerance > 0 or options.date_tolerance_days > 0:
        _run_pass(
            L, R, left_open, right_open, matched_pairs, review_pairs,
            options=options, pass_name="tolerance",
            candidate_fn=_candidates_tolerance,
        )

    # Pass 4 — description fuzzy boost (only over what's left).
    if (
        options.desc_min_score > 0
        and options.left_desc
        and options.right_desc
        and _HAS_RAPIDFUZZ
    ):
        _run_pass(
            L, R, left_open, right_open, matched_pairs, review_pairs,
            options=options, pass_name="fuzzy",
            candidate_fn=_candidates_fuzzy,
        )

    # Build the four output frames from what remains.
    matched_df = _build_matched(left, right, matched_pairs, options)
    review_df = _build_matched(left, right, review_pairs, options, review=True)
    unmatched_left_df = left.loc[sorted(left_open)].copy()
    unmatched_right_df = right.loc[sorted(right_open)].copy()

    stats = {
        "left_rows": len(left),
        "right_rows": len(right),
        "matched": len(matched_pairs),
        "review": len(review_pairs),
        "unmatched_left": len(unmatched_left_df),
        "unmatched_right": len(unmatched_right_df),
    }
    logger.debug("reconcile stats: {}", stats)

    return ReconcileResult(
        matched=matched_df,
        unmatched_left=unmatched_left_df,
        unmatched_right=unmatched_right_df,
        review=review_df,
        stats=stats,
    )


# ---------------------------------------------------------------------------
# Input validation & prep
# ---------------------------------------------------------------------------


def _validate_options(
    left: pd.DataFrame, right: pd.DataFrame, options: ReconcileOptions
) -> None:
    if not options.left_amount or not options.right_amount:
        raise ValueError(
            "Reconcile requires both left_amount and right_amount columns."
        )
    if options.left_amount not in left.columns:
        raise ValueError(
            f"left_amount column {options.left_amount!r} not in left DataFrame."
        )
    if options.right_amount not in right.columns:
        raise ValueError(
            f"right_amount column {options.right_amount!r} not in right DataFrame."
        )
    if bool(options.left_date) != bool(options.right_date):
        raise ValueError(
            "left_date and right_date must both be set or both be None."
        )
    if options.left_date and options.left_date not in left.columns:
        raise ValueError(f"left_date column {options.left_date!r} not in left.")
    if options.right_date and options.right_date not in right.columns:
        raise ValueError(f"right_date column {options.right_date!r} not in right.")
    if len(options.left_keys) != len(options.right_keys):
        raise ValueError(
            "left_keys and right_keys must be the same length "
            f"(got {len(options.left_keys)} vs {len(options.right_keys)})."
        )
    for c in options.left_keys:
        if c not in left.columns:
            raise ValueError(f"left key column {c!r} not in left DataFrame.")
    for c in options.right_keys:
        if c not in right.columns:
            raise ValueError(f"right key column {c!r} not in right DataFrame.")
    if options.amount_tolerance < 0:
        raise ValueError("amount_tolerance must be >= 0.")
    if options.date_tolerance_days < 0:
        raise ValueError("date_tolerance_days must be >= 0.")
    if not (0 <= options.desc_min_score <= 100):
        raise ValueError("desc_min_score must be between 0 and 100.")


def _prep_side(
    df: pd.DataFrame, options: ReconcileOptions, side: str
) -> pd.DataFrame:
    """Return a copy with ``_amt`` and ``_date`` work columns added.

    Rows whose amount cannot be parsed as a number are dropped from the
    matching frame so they fall through to the unmatched bucket on the
    caller side. The same is true for unparseable dates when date
    matching is in use — date is required-when-configured.
    """
    work = df.copy()
    amt_col = options.left_amount if side == "left" else options.right_amount
    date_col = options.left_date if side == "left" else options.right_date

    work["_amt"] = pd.to_numeric(work[amt_col], errors="coerce")
    if side == "right" and options.invert_right_sign:
        work["_amt"] = -work["_amt"]

    if date_col:
        work["_date"] = pd.to_datetime(work[date_col], errors="coerce")
    else:
        work["_date"] = pd.NaT

    # Drop rows that lack the inputs needed to participate. Their
    # original index labels are intentionally preserved on the source
    # frame so they show up in unmatched buckets below.
    bad_amt = work["_amt"].isna()
    bad_date = work["_date"].isna() if date_col else pd.Series(False, index=work.index)
    keep = ~(bad_amt | bad_date)
    if (~keep).any():
        logger.debug(
            "{} side: dropping {} row(s) with unparseable amount/date",
            side, (~keep).sum(),
        )
    return work.loc[keep].copy()


# ---------------------------------------------------------------------------
# Per-pass orchestration
# ---------------------------------------------------------------------------


def _run_pass(
    L: pd.DataFrame,
    R: pd.DataFrame,
    left_open: set,
    right_open: set,
    matched_pairs: list[dict],
    review_pairs: list[dict],
    *,
    options: ReconcileOptions,
    pass_name: str,
    candidate_fn,
) -> None:
    """Run one matching pass over the still-open indices.

    The pass collects (left_idx, right_idx, score, extras) candidates
    from ``candidate_fn``, then greedily assigns by descending score.
    A left row with two right candidates tied at the top score (and
    vice versa) gets routed to the review bucket so we don't pick one
    arbitrarily.
    """
    L_open = L.loc[L.index.intersection(left_open)]
    R_open = R.loc[R.index.intersection(right_open)]
    if L_open.empty or R_open.empty:
        return

    candidates = candidate_fn(L_open, R_open, options)
    if not candidates:
        return

    # Group candidates by left index. For each left row, partition into
    # confident-best (single top score) vs. ambiguous (top score tied).
    by_left: dict = {}
    for cand in candidates:
        by_left.setdefault(cand["left_idx"], []).append(cand)

    # Two-pointer assignment by best-score-first, with reverse-direction
    # ambiguity check so a right row claimed by two equally-good lefts
    # also routes to review.
    by_right_top: dict = {}
    for li, cands in by_left.items():
        cands.sort(key=lambda c: c["score"], reverse=True)
        top = cands[0]["score"]
        leaders = [c for c in cands if c["score"] == top]
        for c in leaders:
            by_right_top.setdefault(c["right_idx"], []).append(c)

    # Sort left rows by their leader's score so high-confidence matches
    # claim their right counterpart first; low-confidence rows lose
    # contention if the right row was already taken.
    left_order = sorted(
        by_left.keys(),
        key=lambda li: -by_left[li][0]["score"],
    )

    for li in left_order:
        if li not in left_open:
            continue
        cands = by_left[li]
        top_score = cands[0]["score"]
        leaders = [c for c in cands if c["score"] == top_score]

        # Filter to still-open right indices.
        leaders = [c for c in leaders if c["right_idx"] in right_open]
        if not leaders:
            continue

        if len(leaders) > 1:
            # Left row is ambiguous on its own side — multiple equally
            # good right candidates remain. Park them all in review.
            for c in leaders:
                review_pairs.append({**c, "pass": pass_name})
            left_open.discard(li)
            for c in leaders:
                right_open.discard(c["right_idx"])
            continue

        pick = leaders[0]
        ri = pick["right_idx"]

        # Mirror check: is the right row contested by another left at
        # the same top score? If so, both lefts go to review and the
        # right row is consumed.
        contenders = [
            c for c in by_right_top.get(ri, [])
            if c["left_idx"] in left_open and c["score"] == pick["score"]
        ]
        if len(contenders) > 1:
            for c in contenders:
                review_pairs.append({**c, "pass": pass_name})
                left_open.discard(c["left_idx"])
            right_open.discard(ri)
            continue

        matched_pairs.append({**pick, "pass": pass_name})
        left_open.discard(li)
        right_open.discard(ri)


# ---------------------------------------------------------------------------
# Candidate generators (one per pass)
# ---------------------------------------------------------------------------


def _candidates_by_key(
    L: pd.DataFrame, R: pd.DataFrame, options: ReconcileOptions
) -> list[dict]:
    """Exact match on the user-supplied key columns + same amount.

    Amount must still tie out; otherwise a shared reference number
    (e.g. a check number reused across years) would over-match. We do
    NOT require date in this pass — the assumption is that a confirmed
    reference like an invoice number is authoritative even when the
    posting date drifts.
    """
    if not options.left_keys:
        return []
    # Build a composite key on each side as a tuple of stringified values.
    L_key = L[options.left_keys].astype(str).agg("|".join, axis=1)
    R_key = R[options.right_keys].astype(str).agg("|".join, axis=1)
    R_by_key: dict = {}
    for ri, k in R_key.items():
        R_by_key.setdefault(k, []).append(ri)

    out: list[dict] = []
    for li, k in L_key.items():
        if k == "" or k == "|".join([""] * len(options.left_keys)):
            continue
        for ri in R_by_key.get(k, []):
            if abs(L.at[li, "_amt"] - R.at[ri, "_amt"]) <= options.amount_tolerance:
                out.append(_score_pair(L, R, li, ri, base_score=1000))
    return out


def _candidates_exact(
    L: pd.DataFrame, R: pd.DataFrame, options: ReconcileOptions
) -> list[dict]:
    """Exact match on amount (and date if configured)."""
    out: list[dict] = []
    has_date = options.left_date is not None
    # Bucket right side by amount for cheap lookup.
    R_by_amt: dict = {}
    for ri, amt in R["_amt"].items():
        R_by_amt.setdefault(amt, []).append(ri)

    for li, amt in L["_amt"].items():
        for ri in R_by_amt.get(amt, []):
            if has_date and L.at[li, "_date"] != R.at[ri, "_date"]:
                continue
            out.append(_score_pair(L, R, li, ri, base_score=900))
    return out


def _candidates_tolerance(
    L: pd.DataFrame, R: pd.DataFrame, options: ReconcileOptions
) -> list[dict]:
    """Amount within tolerance and (if configured) date within window.

    Quadratic in the open set size. For typical reconciliation sizes
    (a month of statements: low thousands of rows) this is fine; if a
    user hands us 100k×100k we'll need a smarter blocking strategy.
    """
    out: list[dict] = []
    has_date = options.left_date is not None
    tol = options.amount_tolerance
    win = pd.Timedelta(days=options.date_tolerance_days) if has_date else None

    R_amts = R["_amt"].to_numpy()
    R_dates = R["_date"].to_numpy() if has_date else None
    R_index = R.index.to_numpy()

    for li in L.index:
        l_amt = L.at[li, "_amt"]
        l_date = L.at[li, "_date"] if has_date else None
        amt_ok = (R_amts >= l_amt - tol) & (R_amts <= l_amt + tol)
        if has_date:
            date_diff = R_dates - l_date.to_datetime64()
            date_ok = (date_diff >= -win.to_timedelta64()) & (
                date_diff <= win.to_timedelta64()
            )
            mask = amt_ok & date_ok
        else:
            mask = amt_ok
        for ri in R_index[mask]:
            out.append(_score_pair(L, R, li, ri, base_score=500))
    return out


def _candidates_fuzzy(
    L: pd.DataFrame, R: pd.DataFrame, options: ReconcileOptions
) -> list[dict]:
    """Tolerance-pass candidates re-scored by description similarity.

    Only kept when the description similarity meets the threshold AND
    the amount is within tolerance. Score blends the two so a strong
    description match outranks a marginal amount match within the same
    pass.
    """
    if not (_HAS_RAPIDFUZZ and options.left_desc and options.right_desc):
        return []
    out: list[dict] = []
    has_date = options.left_date is not None
    tol = options.amount_tolerance
    win = pd.Timedelta(days=options.date_tolerance_days) if has_date else None
    min_score = options.desc_min_score

    L_desc = L[options.left_desc].astype(str)
    R_desc = R[options.right_desc].astype(str)

    for li in L.index:
        l_amt = L.at[li, "_amt"]
        l_date = L.at[li, "_date"] if has_date else None
        l_text = L_desc.at[li]
        for ri in R.index:
            if abs(R.at[ri, "_amt"] - l_amt) > tol:
                continue
            if has_date:
                diff = R.at[ri, "_date"] - l_date
                if abs(diff) > win:
                    continue
            score = int(_rf_fuzz.token_set_ratio(l_text, R_desc.at[ri]))
            if score < min_score:
                continue
            # Base 300 keeps fuzzy below exact/tolerance passes; the
            # 0–100 description score breaks ties within the pass.
            out.append(
                _score_pair(L, R, li, ri, base_score=300 + score, desc_score=score)
            )
    return out


# ---------------------------------------------------------------------------
# Scoring & output assembly
# ---------------------------------------------------------------------------


def _score_pair(
    L: pd.DataFrame,
    R: pd.DataFrame,
    li,
    ri,
    *,
    base_score: int,
    desc_score: int = 0,
) -> dict:
    """Build the candidate record used by the assignment phase."""
    amt_diff = float(L.at[li, "_amt"] - R.at[ri, "_amt"])
    l_date = L.at[li, "_date"]
    r_date = R.at[ri, "_date"]
    if pd.notna(l_date) and pd.notna(r_date):
        date_diff_days = int((l_date - r_date).days)
    else:
        date_diff_days = None
    # Penalize larger differences within the same pass so closer matches
    # win ties. Cap penalty so it can't flip pass ordering.
    penalty = min(abs(amt_diff) * 10, 50)
    if date_diff_days is not None:
        penalty += min(abs(date_diff_days), 50)
    return {
        "left_idx": li,
        "right_idx": ri,
        "score": base_score - penalty,
        "amount_diff": amt_diff,
        "date_diff_days": date_diff_days,
        "desc_score": desc_score,
    }


def _build_matched(
    left: pd.DataFrame,
    right: pd.DataFrame,
    pairs: list[dict],
    options: ReconcileOptions,
    *,
    review: bool = False,
) -> pd.DataFrame:
    """Assemble a matched/review frame: bookkeeping cols + originals."""
    if not pairs:
        cols = ["match_pass", "score", "amount_diff", "date_diff_days", "desc_score"]
        cols += [f"left_{c}" for c in left.columns]
        cols += [f"right_{c}" for c in right.columns]
        return pd.DataFrame(columns=cols)

    rows = []
    for p in pairs:
        li, ri = p["left_idx"], p["right_idx"]
        row = {
            "match_pass": p["pass"],
            "score": p["score"],
            "amount_diff": p["amount_diff"],
            "date_diff_days": p["date_diff_days"],
            "desc_score": p["desc_score"],
        }
        for c in left.columns:
            row[f"left_{c}"] = left.at[li, c]
        for c in right.columns:
            row[f"right_{c}"] = right.at[ri, c]
        rows.append(row)
    out = pd.DataFrame(rows)
    # Stable ordering: review by left_idx so paired rows stay adjacent;
    # matched by score descending so the user sees the strongest pairs
    # first.
    if review:
        out = out.sort_values("score", ascending=False, kind="stable")
    else:
        out = out.sort_values("score", ascending=False, kind="stable")
    return out.reset_index(drop=True)
