"""Join filtered integration failures with CI bisect attribution and cluster.

CI's `new_failures_with_bad_commit_grouped_by_authors.json` shape is:
  { author_or_null: { model: { "single-gpu" | "multi-gpu": [ {test, bad_commit,
    status, pr_number, parent, failure_at_*_commit, job_link}, ... ] } } }

A status of `"git bisect found the bad commit."` is the only one we trust.
Anything starting with `"flaky:"` we tag as flaky and segregate.
Anything missing entirely becomes `unpinned`.
"""

from __future__ import annotations

from collections import defaultdict


_GOOD_STATUS = "git bisect found the bad commit."


def _index_attribution(new_failures: dict) -> dict[tuple[str, str, str], dict]:
    """Flatten `{author -> {model -> {gpu -> [records]}}}` to
    `{(model, gpu, test) -> record}`. Adds `author` to each record."""
    out: dict[tuple[str, str, str], dict] = {}
    if not new_failures:
        return out
    for author, by_model in (new_failures or {}).items():
        if not isinstance(by_model, dict):
            continue
        for model, by_gpu in by_model.items():
            if not isinstance(by_gpu, dict):
                continue
            for gpu_label, items in by_gpu.items():
                gpu = gpu_label.replace("-gpu", "")
                for rec in items or []:
                    test = rec.get("test", "")
                    enriched = {**rec, "author": author if author != "null" else None}
                    out[(model, gpu, test)] = enriched
    return out


def cluster_failures(
    filtered: list[dict],
    new_failures_latest: dict | None,
    classifier,
) -> dict:
    """Produce the triage report data structure.

    Inputs:
      filtered            ── output of filter.intersect_across_days
      new_failures_latest ── most-recent day's new_failures JSON (may be None)
      classifier          ── classify.classify

    Returns: dict with keys
      `clusters`    {bad_commit: {meta..., failures: [...]}}, sorted by size desc
      `flaky`       [failure, ...] (CI marked status="flaky:...")
      `unpinned`    [failure, ...] (no CI attribution found)
      `totals`      {total, clusters, flaky, unpinned}
    """
    attr = _index_attribution(new_failures_latest or {})

    clusters: dict[str, dict] = {}
    flaky: list[dict] = []
    unpinned: list[dict] = []

    for f in filtered:
        key = (f["model"], f["gpu"], f["test"])
        rec = attr.get(key)
        f = {
            **f,
            "failure_mode": classifier(f.get("latest_trace") or f.get("trace") or ""),
        }
        if rec is None:
            unpinned.append(f)
            continue
        status = rec.get("status") or ""
        if status.startswith("flaky"):
            flaky.append({**f, "status": status, "author": rec.get("author")})
            continue
        if status != _GOOD_STATUS:
            unpinned.append({**f, "status": status, "author": rec.get("author")})
            continue
        bc = rec.get("bad_commit")
        if not bc:
            unpinned.append({**f, "author": rec.get("author")})
            continue
        c = clusters.setdefault(
            bc,
            {
                "bad_commit": bc,
                "pr_number": rec.get("pr_number"),
                "author": rec.get("author"),
                "merged_by": rec.get("merged_by"),
                "parent": rec.get("parent"),
                "job_link": rec.get("job_link"),
                "failure_excerpt": (rec.get("failure_at_bad_commit") or "").strip(),
                "failures": [],
            },
        )
        c["failures"].append(f)

    # Sort clusters by failure count desc, then by author for stability.
    clusters_sorted = dict(
        sorted(clusters.items(), key=lambda kv: (-len(kv[1]["failures"]), kv[1].get("author") or ""))
    )

    return {
        "clusters": clusters_sorted,
        "flaky": flaky,
        "unpinned": unpinned,
        "totals": {
            "total": len(filtered),
            "in_clusters": sum(len(c["failures"]) for c in clusters_sorted.values()),
            "clusters": len(clusters_sorted),
            "flaky": len(flaky),
            "unpinned": len(unpinned),
        },
    }


# ── big-model exclusion ────────────────────────────────────────────────
# Skip from local bisect anything whose model name appears here. Total params
# >100B is the user-locked threshold. We key by `model` (the path component
# under `tests/models/...`) since that's the only stable identifier in the CI
# failure record.
BIG_MODEL_SKIP: set[str] = {
    # Llama-4 Maverick (402B) — Scout-17B-16E (109B) STAYS, we have it on disk
    "llama4_maverick",  # if present
    # DeepSeek-V4-Flash-Base, ~600B
    "deepseek_v4_flash_base",
    # Cohere Command-A-Plus, ~111B, gated
    "cohere2_moe",  # command-a-plus is the canonical checkpoint there
    # Qwen3-235B-A22B / Qwen3-Next variants
    "qwen3_next",
    "qwen3_235b",
    # Other known-massive MoE bases
    "deepseek_r1",
    "dbrx",  # 132B
}


def is_big_model(failure: dict) -> bool:
    return failure["model"] in BIG_MODEL_SKIP


if __name__ == "__main__":
    import sys

    from fetch import fetch_last_n
    from filter import per_day_integration_failures, intersect_across_days
    from classify import classify

    cache = sys.argv[1] if len(sys.argv) > 1 else None
    daily = fetch_last_n(7, cache_dir=cache)
    per_day = per_day_integration_failures(daily)
    kept = intersect_across_days(per_day, min_days=5)
    # Latest day's bisect attribution
    latest_date = max(daily.keys())
    nf_latest = daily[latest_date].get("new_failures")
    report = cluster_failures(kept, nf_latest, classify)

    t = report["totals"]
    print(
        f"total={t['total']}  clusters={t['clusters']}  in_clusters={t['in_clusters']}  "
        f"flaky={t['flaky']}  unpinned={t['unpinned']}"
    )
    print("\nTop 8 clusters by size:")
    for bc, c in list(report["clusters"].items())[:8]:
        models = ", ".join(sorted({f["model"] for f in c["failures"]})[:5])
        more = "" if len({f["model"] for f in c["failures"]}) <= 5 else " …"
        print(
            f"  {bc[:10]}  PR #{c['pr_number'] or '—':<5}  "
            f"{c['author'] or '?':<22}  {len(c['failures']):>3} failures  "
            f"[{models}{more}]"
        )
    print(f"\nUnpinned: {len(report['unpinned'])}")
    if report["unpinned"]:
        from collections import Counter
        modes = Counter(f["failure_mode"] for f in report["unpinned"])
        for m, n in modes.most_common():
            print(f"  {m}: {n}")