#!/usr/bin/env python3
"""Classify retained failures without replacing or rerunning any outputs."""

from __future__ import annotations

import argparse
import json
from collections import Counter
from pathlib import Path


MODELS = ("qwen38", "qwen36", "gemma4")


def load(path: Path):
    return json.loads(path.read_text(encoding="utf-8"))


def classify(base: Path) -> dict:
    result: dict[str, dict] = {}

    rows = [load(path) for path in sorted((base / "reasoning").glob("*/*.json"))]
    counts = Counter()
    for row in rows:
        if row.get("passed"):
            counts["pass"] += 1
        elif not (row.get("output") or "").strip():
            counts["no_visible_final"] += 1
        elif row.get("answer_correct") and not row.get("format_compliant"):
            counts["format_only_failure"] += 1
        else:
            counts["wrong_answer"] += 1
    result["reasoning"] = {"total": len(rows), "counts": dict(counts)}

    rows = [load(path) for path in sorted((base / "documents").glob("*.json"))]
    counts = Counter()
    for row in rows:
        if row.get("passed"):
            counts["pass"] += 1
        elif row.get("error"):
            counts["runtime_error"] += 1
        elif not (row.get("output") or "").strip():
            counts["no_visible_final"] += 1
        else:
            counts["answer_or_citation_mismatch"] += 1
    result["documents"] = {"total": len(rows), "counts": dict(counts)}

    rows = [load(path) for path in sorted((base / "tools-single").glob("*/*.json"))]
    counts = Counter()
    for row in rows:
        if row.get("passed"):
            counts["pass"] += 1
        elif not row.get("calls"):
            counts["no_tool_call"] += 1
        else:
            counts["wrong_tool_or_arguments"] += 1
    result["tools_single"] = {"total": len(rows), "counts": dict(counts)}

    rows = [load(path) for path in sorted((base / "tools-multi").glob("*/*.json"))]
    counts = Counter()
    for row in rows:
        if row.get("passed"):
            counts["pass"] += 1
        elif int(row.get("completed_steps") or 0) < 2:
            counts["incomplete_workflow"] += 1
        else:
            counts["sequence_or_final_mismatch"] += 1
    result["tools_multi"] = {"total": len(rows), "counts": dict(counts)}

    rows = [load(path) for path in sorted((base / "coding").glob("*/*.json"))]
    counts = Counter()
    for row in rows:
        if row.get("passed"):
            counts["pass"] += 1
        elif int(row.get("invalid_calls") or 0):
            counts["invalid_tool_call_and_test_failure"] += 1
        else:
            counts["tests_failed"] += 1
    result["coding"] = {"total": len(rows), "counts": dict(counts)}

    for category in ("instruction", "writing", "multilingual", "vision"):
        rows = [load(path) for path in sorted((base / category).glob("*.json"))]
        result[category] = {
            "total": len(rows), "counts": {
                "pass": sum(bool(row.get("passed")) for row in rows),
                "mechanical_or_objective_failure": sum(not bool(row.get("passed")) for row in rows),
            },
        }
    return result


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--results-root", type=Path, required=True)
    parser.add_argument("--out", type=Path, required=True)
    args = parser.parse_args()
    result = {
        "method": "Classification of retained outputs only; no inference reruns or replacements.",
        "models": {model: classify(args.results_root / model) for model in MODELS},
    }
    args.out.parent.mkdir(parents=True, exist_ok=True)
    args.out.write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8")
    print(json.dumps({model: result["models"][model] for model in MODELS}, indent=2))


if __name__ == "__main__":
    main()
