"""Auto-correct expected_exception in case.json based on actual test output.

Many imported cases predict the wrong exception type (e.g., AssertionError when
the actual error is NameError). This script runs each case's baseline test,
captures the actual exception, updates case.json, and verifies the oracle patch
still fixes it.

Usage::

    PYTHONPATH=. uv run python evals/importers/fixup_expected_exceptions.py [--bug-source debugbench]
"""

from __future__ import annotations

import argparse
import json
import shutil
import tempfile
from pathlib import Path

from evals.run_eval import (
    _REPO_ROOT,
    Condition,
    apply_diff,
    load_cases,
    run_pytest,
)
from evals.validate_cases import (
    expected_exception_in_output,
    extract_exception_candidates,
)


def fixup_cases(
    cases_dir: Path,
    *,
    bug_source: str | None = None,
    timeout_sec: int = 30,
    output_format: str = "json_compact",
) -> dict[str, str]:
    """Auto-correct expected_exception for cases that fail only on exception mismatch.

    Returns dict mapping case_id -> status.
    """
    cases = load_cases(cases_dir)
    results: dict[str, str] = {}

    for spec in cases:
        case_id = spec.case_id
        case_path = spec.path

        # Filter by bug_source if specified.
        if bug_source and spec.bug_source != bug_source:
            continue

        # Read full case.json for modification.
        meta = json.loads((case_path / "case.json").read_text(encoding="utf-8"))
        expected_exc = spec.expected_exception
        python_args = spec.python_args
        case_timeout = spec.timeout_sec

        with tempfile.TemporaryDirectory(prefix=f"llmdebug-fixup-{case_id}-") as tmp:
            workdir = Path(tmp)
            shutil.copytree(
                case_path,
                workdir,
                dirs_exist_ok=True,
                ignore=shutil.ignore_patterns("__pycache__", "*.pyc", "*.pyo"),
            )

            # Run baseline.
            baseline = run_pytest(
                workdir, python_args, Condition.TRACEBACK_ONLY, case_timeout, output_format
            )

            if baseline.timed_out:
                results[case_id] = "skip: timed out"
                continue

            if baseline.returncode == 0:
                results[case_id] = "skip: baseline passed"
                continue

            combined = (baseline.stdout or "") + "\n" + (baseline.stderr or "")

            # Check if current expected_exception is already correct.
            if expected_exception_in_output(expected_exc, combined):
                # Already correct — verify oracle patch works.
                patch_text = (case_path / "oracle.patch").read_text(encoding="utf-8")
                applied, _ = apply_diff(workdir, patch_text)
                if applied:
                    patched = run_pytest(
                        workdir,
                        python_args,
                        Condition.TRACEBACK_ONLY,
                        case_timeout,
                        output_format,
                    )
                    if patched.returncode == 0:
                        results[case_id] = "ok: already correct"
                        continue
                    results[case_id] = "skip: oracle didn't fix"
                    continue
                results[case_id] = "skip: oracle patch failed"
                continue

            # Find actual exception.
            candidates = extract_exception_candidates(combined)
            if not candidates:
                results[case_id] = "skip: no exception found"
                continue

            actual_exc = candidates[0]  # First (most likely) candidate

            # Verify oracle patch fixes the issue with the new exception.
            patch_text = (case_path / "oracle.patch").read_text(encoding="utf-8")
            applied, _ = apply_diff(workdir, patch_text)
            if not applied:
                results[case_id] = "skip: oracle patch failed"
                continue

            patched = run_pytest(
                workdir, python_args, Condition.TRACEBACK_ONLY, case_timeout, output_format
            )
            if patched.returncode != 0:
                results[case_id] = "skip: oracle didn't fix"
                continue

            # Update case.json with corrected expected_exception.
            meta["expected_exception"] = actual_exc
            (case_path / "case.json").write_text(
                json.dumps(meta, indent=2, ensure_ascii=False) + "\n",
                encoding="utf-8",
            )
            results[case_id] = f"fixed: {expected_exc} -> {actual_exc}"
            print(f"  [fixed] {case_id}: {expected_exc} -> {actual_exc}")

    return results


def main() -> None:
    parser = argparse.ArgumentParser(description="Auto-correct expected_exception in case.json.")
    parser.add_argument("--cases-dir", default="evals/cases")
    parser.add_argument("--bug-source", default=None)
    args = parser.parse_args()

    cases_dir = Path(args.cases_dir)
    if not cases_dir.is_absolute():
        cases_dir = _REPO_ROOT / cases_dir

    results = fixup_cases(cases_dir, bug_source=args.bug_source)

    fixed = sum(1 for v in results.values() if v.startswith("fixed"))
    ok = sum(1 for v in results.values() if v.startswith("ok"))
    skipped = sum(1 for v in results.values() if v.startswith("skip"))

    print(f"\nResults: {fixed} fixed, {ok} already correct, {skipped} skipped")

    # Delete cases that were skipped (unfixable).
    deleted = 0
    cases_dir_path = Path(args.cases_dir)
    if not cases_dir_path.is_absolute():
        cases_dir_path = _REPO_ROOT / cases_dir_path
    for case_id, status in results.items():
        if status.startswith("skip"):
            case_dir = cases_dir_path / case_id
            if case_dir.exists():
                shutil.rmtree(case_dir, ignore_errors=True)
                deleted += 1

    if deleted:
        print(f"Deleted {deleted} unfixable cases.")


if __name__ == "__main__":
    main()
