"""Extract integration-test failures and intersect across days.

An "integration test" here = a failure whose pytest path uses a class whose name
ends with `IntegrationTest` or `IntegrationTests`. This matches the canonical
shape at `tests/models/llama4/test_modeling_llama4.py::Llama4IntegrationTest` and
`tests/models/deepseek_v4/...::DeepseekV4IntegrationTest`.
"""

from __future__ import annotations

from collections import defaultdict
from typing import Iterable


INTEGRATION_SUFFIXES = ("IntegrationTest", "IntegrationTests")


def is_integration_test(test_path: str) -> bool:
    """`tests/models/foo/test_modeling_foo.py::FooIntegrationTest::test_x` → True."""
    if "::" not in test_path:
        return False
    cls = test_path.split("::")[1]
    return cls.endswith(INTEGRATION_SUFFIXES)


def model_name_from_key(key: str) -> str:
    """`models_align` → `align`. (CI keys model_results entries this way.)"""
    return key.removeprefix("models_")


def model_name_from_test_path(test_path: str) -> str:
    """`tests/models/foo/test_modeling_foo.py::...` → `foo`."""
    parts = test_path.split("/")
    if len(parts) < 3 or parts[0] != "tests" or parts[1] != "models":
        return ""
    return parts[2]


def iter_failures(model_results: dict) -> Iterable[dict]:
    """Yield one record per (model, gpu, test) failure from a single day's
    `model_results.json`. Each record is:
        {
          "model": "align",
          "gpu": "single" | "multi",
          "test": "tests/models/.../...::Class::test_x",
          "trace": "...",
        }
    Only integration-test failures are yielded.
    """
    for key, entry in model_results.items():
        if not isinstance(entry, dict):
            continue
        model = model_name_from_key(key)
        failures = entry.get("failures") or {}
        for gpu, items in failures.items():
            # `failures` is `{"single": [...], "multi": [...]}` — strip the
            # `-gpu` suffix CI sometimes uses to match `new_failures` shape.
            gpu = gpu.replace("-gpu", "")
            for it in items or []:
                test = it.get("line", "")
                if not is_integration_test(test):
                    continue
                yield {
                    "model": model,
                    "gpu": gpu,
                    "test": test,
                    "trace": (it.get("trace") or "").strip(),
                }


def intersect_across_days(per_day: dict[str, list[dict]], min_days: int = 5) -> list[dict]:
    """Keep `(model, gpu, test)` triples seen on ≥ min_days days.

    Returns one record per surviving triple, enriched with `days_seen`,
    `first_seen`, `latest_trace`.
    """
    seen: dict[tuple[str, str, str], dict] = {}
    for date in sorted(per_day):  # ascending so latest_trace ends up newest
        for rec in per_day[date]:
            key = (rec["model"], rec["gpu"], rec["test"])
            existing = seen.get(key)
            if existing is None:
                seen[key] = {
                    **rec,
                    "days_seen": 1,
                    "first_seen": date,
                    "latest_seen": date,
                    "latest_trace": rec["trace"],
                }
            else:
                existing["days_seen"] += 1
                existing["latest_seen"] = date
                existing["latest_trace"] = rec["trace"]
    return [r for r in seen.values() if r["days_seen"] >= min_days]


def per_day_integration_failures(daily: dict[str, dict[str, dict | None]]) -> dict[str, list[dict]]:
    """`daily` is the output of `fetch.fetch_last_n`."""
    out: dict[str, list[dict]] = defaultdict(list)
    for date, payload in daily.items():
        mr = payload.get("model_results") if payload else None
        if not mr:
            continue
        out[date] = list(iter_failures(mr))
    return out


if __name__ == "__main__":
    import sys

    from fetch import fetch_last_n

    daily = fetch_last_n(7, cache_dir=sys.argv[1] if len(sys.argv) > 1 else None)
    per_day = per_day_integration_failures(daily)
    for d, rows in sorted(per_day.items()):
        print(f"{d}: {len(rows)} integration-test failures")
    kept = intersect_across_days(per_day, min_days=5)
    print(f"\nKept after ≥5/7 intersection: {len(kept)} unique failures")
    for r in kept[:6]:
        print(f"  {r['model']:<28} {r['gpu']:<6} {r['days_seen']}/7  {r['test'].split('::')[-1]}")
    if len(kept) > 6:
        print(f"  ... and {len(kept) - 6} more")