"""
engine/pipeline.py — rule-set execution: source resolution, stage chaining,
result serialization.

Stage inputs are either uploaded files ("template:<name>") or a prior stage's
output ("stage:<slug>:matched|unmatched_a|unmatched_b") — the reference
implementation's rule chaining, made explicit and tenant-isolated (everything
lives in per-session DataFrames, no shared DB tables).
"""

import json

import pandas as pd

from .loader import load_source
from .matcher import run_stage, MATCH_TIER_COL


class PipelineError(ValueError):
    """Raised for rule-set/source resolution problems."""


def run(rule_set: dict, templates: dict, files: dict, context: dict) -> dict:
    """
    rule_set:  decoded rule-set definition (recon_rule_schema.json)
    templates: {template_name: decoded template definition}
    files:     {template_name: base64 file content}
    context:   {"location_map": {...}, ...} tenant/session bindings

    Returns a JSON-serializable dict: per-stage rows + counts + amounts +
    tier breakdown, plus an overall summary.
    """
    loaded: dict[str, pd.DataFrame] = {}
    outputs: dict[str, dict] = {}
    stage_results = []

    for stage in rule_set["stages"]:
        if not stage.get("enabled", True):
            continue
        slug = stage["stage"]

        df_a = _resolve_source(stage["side_a"]["source"], loaded, outputs, templates, files, context)
        df_b = _resolve_source(stage["side_b"]["source"], loaded, outputs, templates, files, context)

        result = run_stage(df_a, df_b, stage, context)
        outputs[slug] = result
        stage_results.append(_serialize_stage(stage, result))

    return {
        "stages": stage_results,
        "summary": _overall_summary(stage_results),
        "warnings": _date_warnings(loaded, context),
    }


def _date_warnings(loaded: dict, context: dict) -> list:
    """OQ-05: detect the dominant transaction date in each uploaded file and
    flag files whose data disagrees with the session's recon date (wrong date
    picked, or wrong file uploaded)."""
    recon_date = context.get("recon_date")
    warnings = []

    for name, df in loaded.items():
        if "txn_date" not in df.columns:
            continue
        dates = pd.to_datetime(df["txn_date"], errors="coerce").dt.date.dropna()
        if dates.empty:
            continue
        dominant = dates.mode().iloc[0]
        share = (dates == dominant).mean()

        if recon_date:
            try:
                expected = pd.to_datetime(recon_date).date()
            except Exception:
                continue
            if dominant != expected:
                warnings.append(
                    f"File '{name}': most transactions are dated {dominant} "
                    f"but the selected recon date is {expected} — check the "
                    f"date or the uploaded file."
                )
        elif share < 0.9:
            warnings.append(
                f"File '{name}': transactions span multiple dates "
                f"(most common: {dominant}) — no recon date was selected."
            )

    return warnings


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

def _resolve_source(ref: str, loaded, outputs, templates, files, context) -> pd.DataFrame:
    if ref.startswith("template:"):
        name = ref.removeprefix("template:")
        if name not in loaded:
            if name not in templates:
                raise PipelineError(f"Rule set references unknown template '{name}'")
            if name not in files:
                raise PipelineError(
                    f"No file uploaded for template '{name}' "
                    f"(uploaded: {sorted(files.keys())})"
                )
            loaded[name] = load_source(files[name], templates[name], context)
        return loaded[name]

    if ref.startswith("stage:"):
        _, slug, which = ref.split(":", 2)
        if slug not in outputs:
            raise PipelineError(
                f"Stage output '{ref}' referenced before stage '{slug}' ran — "
                f"stages execute in order"
            )
        frame = outputs[slug].get(which)
        if frame is None:
            raise PipelineError(f"Unknown stage output kind '{which}' in '{ref}'")
        return frame

    raise PipelineError(f"Unsupported source ref: '{ref}'")


def _serialize_stage(stage: dict, result: dict) -> dict:
    matched = result["matched"]
    unmatched_a = result["unmatched_a"]
    unmatched_b = result["unmatched_b"]
    output_cfg = stage.get("output", {})

    return {
        "stage": stage["stage"],
        "display_name": stage.get("display_name", stage["stage"]),
        "matched_rows": _records(matched),
        "matched_b_rows": _records(result["matched_b"]),
        "unmatched_a_rows": _records(unmatched_a),
        "unmatched_b_rows": _records(unmatched_b),
        "matched_count": len(matched),
        "unmatched_a_count": len(unmatched_a),
        "unmatched_b_count": len(unmatched_b),
        "total_count": len(matched) + len(unmatched_a),
        "matched_amount": _safe_sum(matched, "amount"),
        "unmatched_a_amount": _safe_sum(unmatched_a, "amount"),
        "unmatched_b_amount": _safe_sum(unmatched_b, "amount"),
        "total_amount": _safe_sum(matched, "amount") + _safe_sum(unmatched_a, "amount"),
        "tier_counts": result["tier_counts"],
        "location_summary": _location_summary(matched, unmatched_a),
        "sheet_matched": output_cfg.get("matched_sheet"),
        "sheet_unmatched": output_cfg.get("unmatched_sheet"),
        "sheet_unmatched_b": output_cfg.get("unmatched_sheet_b"),
    }


def _overall_summary(stage_results: list) -> dict:
    return {
        "stage_count": len(stage_results),
        "matched_count": sum(s["matched_count"] for s in stage_results),
        "unmatched_a_count": sum(s["unmatched_a_count"] for s in stage_results),
        "unmatched_b_count": sum(s["unmatched_b_count"] for s in stage_results),
        "matched_amount": sum(s["matched_amount"] for s in stage_results),
        "unmatched_a_amount": sum(s["unmatched_a_amount"] for s in stage_results),
    }


def _location_summary(matched: pd.DataFrame, unmatched_a: pd.DataFrame) -> list:
    if "location" not in matched.columns and "location" not in unmatched_a.columns:
        return []

    rows = []
    combined = pd.concat(
        [matched.assign(_status="matched"), unmatched_a.assign(_status="unmatched")],
        ignore_index=True,
    )
    if "location" not in combined.columns or combined.empty:
        return []

    amounts = pd.to_numeric(combined.get("amount"), errors="coerce").fillna(0)
    combined = combined.assign(_amount=amounts)

    for location, group in combined.groupby("location", dropna=False):
        m = group[group["_status"] == "matched"]
        u = group[group["_status"] == "unmatched"]
        rows.append({
            "location": None if pd.isna(location) else location,
            "total_count": len(group),
            "total_amount": float(group["_amount"].sum()),
            "matched_count": len(m),
            "matched_amount": float(m["_amount"].sum()),
            "unmatched_count": len(u),
            "unmatched_amount": float(u["_amount"].sum()),
        })
    return rows


def _records(df: pd.DataFrame) -> list:
    """JSON-safe row dicts: NaN/NaT → null, timestamps → ISO strings."""
    if df.empty:
        return []
    return json.loads(df.to_json(orient="records", date_format="iso"))


def _safe_sum(df: pd.DataFrame, col: str) -> float:
    if col not in df.columns:
        return 0.0
    try:
        return float(pd.to_numeric(df[col], errors="coerce").fillna(0).sum())
    except Exception:
        return 0.0