from __future__ import annotations

import importlib
import json
import sys
from collections.abc import Generator
from pathlib import Path
from types import ModuleType

import pytest

pytestmark = [pytest.mark.ci, pytest.mark.ci_npu]


@pytest.fixture
def loader(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Generator[ModuleType]:
    monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2] / "eval"))
    module = importlib.import_module("common.loader")
    monkeypatch.setattr(module, "EVAL_DIR", tmp_path)
    saved = {f"eval_bench_{name}": sys.modules.pop(f"eval_bench_{name}", None) for name in module.SUPPORTED}
    try:
        yield module
    finally:
        for name, previous in saved.items():
            sys.modules.pop(name, None)
            if previous is not None:
                sys.modules[name] = previous


def write_bench(loader: ModuleType, name: str, source: str) -> None:
    folder = loader.EVAL_DIR / name
    folder.mkdir()
    (folder / "bench.py").write_text(source, encoding="utf-8")


@pytest.mark.parametrize("error_type", [RuntimeError, SystemExit, KeyboardInterrupt])
def test_failed_load_can_be_retried(loader: ModuleType, monkeypatch: pytest.MonkeyPatch, error_type: type[BaseException]) -> None:
    write_bench(
        loader,
        "gpqa",
        "import os\n"
        "if os.environ.get('DEEPINFER_TEST_BENCH_READY') != '1':\n"
        f"    raise {error_type.__name__}('initialization failed')\n"
        "READY = True\n",
    )
    monkeypatch.delenv("DEEPINFER_TEST_BENCH_READY", raising=False)
    with pytest.raises(error_type, match="initialization failed"):
        loader.load_bench("gpqa")
    assert "eval_bench_gpqa" not in sys.modules
    monkeypatch.setenv("DEEPINFER_TEST_BENCH_READY", "1")
    bench = loader.load_bench("gpqa")
    assert bench.READY is True
    assert loader.load_bench("gpqa") is bench


def test_retry_after_dependency_becomes_available(loader: ModuleType, monkeypatch: pytest.MonkeyPatch) -> None:
    dependency_name = "_deepinfer_test_bench_dependency"
    monkeypatch.delitem(sys.modules, dependency_name, raising=False)
    write_bench(loader, "hle", f"import {dependency_name}\nREADY = True\n")
    with pytest.raises(ModuleNotFoundError, match=dependency_name):
        loader.load_bench("hle")
    monkeypatch.setitem(sys.modules, dependency_name, ModuleType(dependency_name))
    assert loader.load_bench("hle").READY is True


def test_successful_load_preserves_dataclass_registration_and_cache(loader: ModuleType) -> None:
    write_bench(
        loader,
        "aime",
        "from __future__ import annotations\nfrom dataclasses import dataclass\n@dataclass\nclass Problem:\n    answer: str = '42'\n",
    )
    bench = loader.load_bench("aime")
    assert bench.Problem().answer == "42"
    assert sys.modules[bench.Problem.__module__] is bench
    assert loader.load_bench("aime") is bench


def test_failed_benchmark_does_not_remove_another_cached_module(loader: ModuleType) -> None:
    write_bench(loader, "aime", "READY = True\n")
    write_bench(loader, "gpqa", "raise RuntimeError('broken')\n")
    aime = loader.load_bench("aime")
    with pytest.raises(RuntimeError, match="broken"):
        loader.load_bench("gpqa")
    assert loader.load_bench("aime") is aime


def test_unsupported_benchmark_is_rejected(loader: ModuleType) -> None:
    with pytest.raises(SystemExit, match="benchmark not wired up"):
        loader.load_bench("unsupported")


def test_missing_benchmark_file_is_rejected(loader: ModuleType) -> None:
    with pytest.raises(SystemExit, match="benchmark implementation not found"):
        loader.load_bench("gpqa")


def test_scoring_can_retry_after_benchmark_import_failure(loader: ModuleType, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    score = importlib.import_module("score")
    dependency = "_deepinfer_score_retry_dependency"
    monkeypatch.delitem(sys.modules, dependency, raising=False)
    write_bench(loader, "gpqa", f"import {dependency}\nfrom gpqa.bench import *\n")
    data = tmp_path / "questions.jsonl"
    data.write_text(json.dumps({"task_id": "q1", "question": "Choose B.", "answer": "B"}) + "\n", encoding="utf-8")
    output = tmp_path / "run"
    output.mkdir()
    (output / "meta.json").write_text(json.dumps({"bench": "gpqa", "data_file": str(data), "n_sample": 1}), encoding="utf-8")
    (output / "raw_rank0.jsonl").write_text(
        json.dumps({"task_id": "q1", "sample_idx": 0, "raw_text": "ANSWER: B"}) + "\n", encoding="utf-8"
    )
    with pytest.raises(ModuleNotFoundError, match=dependency):
        score.score_run(str(output), judge="rule", verbose=False)
    assert not (output / "summary.json").exists()
    monkeypatch.setitem(sys.modules, dependency, ModuleType(dependency))
    assert score.score_run(str(output), judge="rule", verbose=False) == 0
    row = json.loads((output / "scored.jsonl").read_text(encoding="utf-8"))
    summary = json.loads((output / "summary.json").read_text(encoding="utf-8"))
    assert row["pred"] == "B"
    assert row["correct"] is True
    assert summary["coverage_complete"] is True
    assert summary["num_rollouts"] == 1