import importlib.util
import io
import inspect
import json
import os
import random
import re
import sys
import tempfile
import types
import unittest
from contextlib import redirect_stderr
from argparse import Namespace
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import patch

def find_repo_root():
    for parent in Path(__file__).resolve().parents:
        if (parent / "scripts" / "accuracy_gsm8k_gpqa.py").is_file():
            return parent
    raise RuntimeError("unable to locate repository scripts directory")


ROOT = find_repo_root()
SCRIPT_PATH = ROOT / "scripts" / "accuracy_gsm8k_gpqa.py"
SHELL_PATH = ROOT / "scripts" / "run_accuracy_gsm8k_gpqa.sh"
VERIFY_PATH = ROOT / "scripts" / "verify_accuracy_matrix.py"


def load_accuracy_module():
    spec = importlib.util.spec_from_file_location("accuracy_gsm8k_gpqa", SCRIPT_PATH)
    module = importlib.util.module_from_spec(spec)
    assert spec.loader is not None
    spec.loader.exec_module(module)
    return module


def load_verify_module():
    spec = importlib.util.spec_from_file_location("verify_accuracy_matrix", VERIFY_PATH)
    module = importlib.util.module_from_spec(spec)
    assert spec.loader is not None
    spec.loader.exec_module(module)
    return module


class AccuracyReproductionTest(unittest.TestCase):
    def test_extract_gsm8k_answer_extracts_signed_decimal_and_removes_commas(self):
        accuracy = load_accuracy_module()

        self.assertEqual(
            accuracy.extract_gsm8k_answer("work shown #### -1,234.50"), "-1234.50"
        )
        self.assertEqual(accuracy.extract_gsm8k_answer("final 2,000"), "2000")

    def test_gsm8k_scoring_compares_normalized_answers(self):
        accuracy = load_accuracy_module()

        self.assertTrue(accuracy.score_gsm8k("#### 1,200.0", "The answer is 1200"))
        self.assertTrue(accuracy.score_gsm8k("#### -3", "-3.0"))
        self.assertFalse(accuracy.score_gsm8k("#### 12", "13"))
        self.assertFalse(accuracy.score_gsm8k("#### 1200", "12"))

    def test_gpqa_choice_order_is_seeded_per_question_and_preserves_correct_index(self):
        accuracy = load_accuracy_module()
        row = {
            "Question": "Which value is correct?",
            "Correct Answer": "correct",
            "Incorrect Answer 1": "wrong one",
            "Incorrect Answer 2": "wrong two",
            "Incorrect Answer 3": "wrong three",
        }

        random.seed(1)
        first = accuracy.build_gpqa_choices(row, "seed-7")
        random.seed(98765)
        second = accuracy.build_gpqa_choices(row, "seed-7")

        self.assertEqual(first, second)
        self.assertEqual(first["choices"][first["correct_index"]], "correct")

    def test_gpqa_normalization_supports_both_row_schemas(self):
        for row in (
            {
                "Question": "Title case?",
                "Correct Answer": "yes",
                "Incorrect Answer 1": "no 1",
                "Incorrect Answer 2": "no 2",
                "Incorrect Answer 3": "no 3",
            },
            {
                "question": "Lower case?",
                "correct_answer": "yes",
                "incorrect_answer_1": "no 1",
                "incorrect_answer_2": "no 2",
                "incorrect_answer_3": "no 3",
            },
        ):
            normalized = load_accuracy_module().normalize_gpqa_row(row)

            self.assertIn(normalized["question"], {"Title case?", "Lower case?"})
            self.assertEqual(normalized["correct_answer"], "yes")
            self.assertEqual(len(normalized["incorrect_answers"]), 3)

    def test_gpqa_parser_requires_explicit_standalone_answer(self):
        accuracy = load_accuracy_module()

        self.assertEqual(accuracy.parse_gpqa_answer("Answer: B"), "B")
        self.assertEqual(accuracy.parse_gpqa_answer("(d)"), "D")
        self.assertIsNone(accuracy.parse_gpqa_answer("I considered B, but cannot decide."))

    def test_cli_rejects_invalid_offload_combinations_and_limits(self):
        accuracy = load_accuracy_module()

        for arguments in (
            ["--offload-group-size", "-1"],
            ["--offload-group-size", "0", "--offload-num-in-group", "2"],
            ["--offload-group-size", "2", "--offload-num-in-group", "3"],
            ["--limit", "-1"],
            ["--gpu-memory-utilization", "0"],
            ["--gpu-memory-utilization", "1.1"],
        ):
            with self.assertRaises(SystemExit):
                with redirect_stderr(io.StringIO()):
                    accuracy.parse_args(["--model", "test-model", *arguments])

    def test_result_metadata_records_hash_and_claimed_parameters(self):
        accuracy = load_accuracy_module()
        args = Namespace(
            model="test-model",
            dataset="gsm8k",
            seed="stable-seed",
            gpu_memory_utilization=0.5,
            max_tokens=17,
            temperature=0,
            offload_group_size=8,
            offload_num_in_group=2,
        )

        result = accuracy.build_result_metadata(
            args=args,
            dataset_source={"name": "modelscope/gsm8k", "revision": "main", "split": "test"},
            dataset_sha256="abc123",
            correct=3,
            total=4,
            elapsed_seconds=1.25,
            output_path=Path("/tmp/examples.jsonl"),
        )

        self.assertEqual(result["dataset_sha256"], "abc123")
        self.assertEqual(result["offload_settings"], {"group_size": 8, "num_in_group": 2})
        self.assertEqual(
            result["sampling_settings"],
            {"temperature": 0, "max_tokens": 17, "enforce_eager": True},
        )
        self.assertEqual(
            result["engine_settings"], {"gpu_memory_utilization": 0.5}
        )
        self.assertEqual(result["accuracy"], 0.75)

    def test_generation_passes_offload_parameters_directly_to_llm(self):
        accuracy = load_accuracy_module()
        calls = {}

        class FakeLLM:
            def __init__(self, **kwargs):
                calls["llm_kwargs"] = kwargs

            def generate(self, prompts, sampling):
                calls["prompts"] = prompts
                calls["sampling"] = sampling
                return [SimpleNamespace(outputs=[SimpleNamespace(text="answer")])]

        class FakeSamplingParams:
            def __init__(self, **kwargs):
                self.kwargs = kwargs

        fake_vllm = types.ModuleType("vllm")
        fake_vllm.LLM = FakeLLM
        fake_vllm.SamplingParams = FakeSamplingParams
        args = Namespace(
            model="test-model",
            max_model_len=4096,
            gpu_memory_utilization=0.5,
            max_tokens=17,
            offload_group_size=8,
            offload_num_in_group=2,
        )

        with patch.dict(sys.modules, {"vllm": fake_vllm}):
            with patch.dict(os.environ, {}, clear=True):
                self.assertEqual(accuracy.generate_prompts(args, ["prompt"]), ["answer"])
                self.assertNotIn("VLLM_ASCEND_OFFLOAD_GROUP_SIZE", os.environ)
                self.assertNotIn("VLLM_ASCEND_OFFLOAD_NUM_IN_GROUP", os.environ)

        self.assertEqual(
            calls["llm_kwargs"],
            {
                "model": "test-model",
                "enforce_eager": True,
                "max_model_len": 4096,
                "gpu_memory_utilization": 0.5,
                "offload_group_size": 8,
                "offload_num_in_group": 2,
            },
        )
        self.assertEqual(calls["sampling"].kwargs, {"temperature": 0, "max_tokens": 17})

    def test_modelscope_calls_use_separate_subset_revision_and_split(self):
        accuracy = load_accuracy_module()
        self.assertIn("revision", inspect.signature(accuracy._load_modelscope_rows).parameters)
        calls = []

        class FakeMsDataset:
            @classmethod
            def load(
                cls,
                name,
                subset_name=None,
                version=None,
                split=None,
                trust_remote_code=None,
            ):
                kwargs = {
                    "subset_name": subset_name,
                    "version": version,
                    "trust_remote_code": trust_remote_code,
                }
                if split is not None:
                    kwargs["split"] = split
                calls.append((name, kwargs))
                if subset_name == "main":
                    return [{"question": "q", "answer": "#### 1"}]
                if split is None:
                    return SimpleNamespace(info=SimpleNamespace(splits={"train": {}}))
                return [
                    {
                        "Question": "q",
                        "Correct Answer": "a",
                        "Incorrect Answer 1": "b",
                        "Incorrect Answer 2": "c",
                        "Incorrect Answer 3": "d",
                    }
                ]

        fake_modelscope = types.ModuleType("modelscope")
        fake_msdatasets = types.ModuleType("modelscope.msdatasets")
        fake_msdatasets.MsDataset = FakeMsDataset
        with patch.dict(
            sys.modules,
            {"modelscope": fake_modelscope, "modelscope.msdatasets": fake_msdatasets},
        ):
            try:
                gsm_rows, gsm_source = accuracy._load_modelscope_rows(
                    "gsm8k", revision="release-1"
                )
                gpqa_rows, gpqa_source = accuracy._load_modelscope_rows(
                    "gpqa", revision="release-2"
                )
            except TypeError as error:
                self.fail(f"MsDataset.load received an unsupported keyword: {error}")

        self.assertEqual(gsm_rows[0]["question"], "q")
        self.assertEqual(gpqa_rows[0]["Correct Answer"], "a")
        self.assertEqual(
            calls,
            [
                (
                    "modelscope/gsm8k",
                    {
                        "subset_name": "main",
                        "version": "release-1",
                        "split": "test",
                        "trust_remote_code": True,
                    },
                ),
                (
                    "modelscope/gpqa",
                    {
                        "subset_name": "gpqa_diamond",
                        "version": "release-2",
                        "trust_remote_code": True,
                    },
                ),
                (
                    "modelscope/gpqa",
                    {
                        "subset_name": "gpqa_diamond",
                        "version": "release-2",
                        "split": "train",
                        "trust_remote_code": True,
                    },
                ),
            ],
        )
        self.assertEqual(
            gsm_source,
            {
                "name": "modelscope/gsm8k",
                "subset": "main",
                "revision": "release-1",
                "split": "test",
            },
        )
        self.assertEqual(gpqa_source["subset"], "gpqa_diamond")
        self.assertEqual(gpqa_source["revision"], "release-2")
        self.assertEqual(gpqa_source["split"], "train")

    def test_gpqa_discovers_labeled_rows_from_multisplit_mapping(self):
        accuracy = load_accuracy_module()
        calls = []
        labeled_row = {
            "Question": "mapped question",
            "Correct Answer": "a",
            "Incorrect Answer 1": "b",
            "Incorrect Answer 2": "c",
            "Incorrect Answer 3": "d",
        }

        class FakeMsDataset:
            @classmethod
            def load(
                cls,
                name,
                subset_name=None,
                version=None,
                split=None,
                trust_remote_code=None,
            ):
                calls.append(
                    (name, subset_name, version, split, trust_remote_code)
                )
                return {
                    "unlabeled": [{"Question": "no labels"}],
                    "train": [labeled_row],
                }

        fake_modelscope = types.ModuleType("modelscope")
        fake_msdatasets = types.ModuleType("modelscope.msdatasets")
        fake_msdatasets.MsDataset = FakeMsDataset
        with patch.dict(
            sys.modules,
            {"modelscope": fake_modelscope, "modelscope.msdatasets": fake_msdatasets},
        ):
            try:
                rows, source = accuracy._load_modelscope_rows(
                    "gpqa", revision="release-map"
                )
            except (RuntimeError, TypeError, ValueError) as error:
                self.fail(f"GPQA mapping metadata was not supported: {error}")

        self.assertEqual(rows, [labeled_row])
        self.assertEqual(source["split"], "train")
        self.assertEqual(
            calls,
            [("modelscope/gpqa", "gpqa_diamond", "release-map", None, True)],
        )

    def test_materialized_manifest_is_persisted_and_reused(self):
        accuracy = load_accuracy_module()
        self.assertIn("revision", inspect.signature(accuracy.materialize_dataset).parameters)
        rows = [{"question": "How many?", "answer": "#### 5"}]
        source = {
            "name": "modelscope/gsm8k",
            "subset": "main",
            "revision": "release-1",
            "split": "test",
        }

        with tempfile.TemporaryDirectory() as temporary_directory:
            data_dir = Path(temporary_directory)
            with patch.object(accuracy, "_load_modelscope_rows", return_value=(rows, source)):
                first = accuracy.materialize_dataset(data_dir, "gsm8k", revision="release-1")

            manifest_path = data_dir / "gsm8k.manifest.json"
            self.assertTrue(manifest_path.is_file())
            manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
            self.assertEqual(manifest["dataset_source"], source)
            self.assertEqual(manifest["dataset_sha256"], first[2])
            self.assertEqual(manifest["dataset_path"], str(first[0].resolve()))
            self.assertEqual(manifest["row_count"], 1)

            with patch.object(
                accuracy,
                "_load_modelscope_rows",
                side_effect=AssertionError("local reuse must not access ModelScope"),
            ):
                second = accuracy.materialize_dataset(data_dir, "gsm8k", revision="ignored")

            self.assertEqual(second[0], first[0])
            self.assertEqual(second[1], source)
            self.assertEqual(second[2], first[2])

    def test_existing_jsonl_without_manifest_is_validated_and_adopted_as_local(self):
        accuracy = load_accuracy_module()
        with tempfile.TemporaryDirectory() as temporary_directory:
            data_dir = Path(temporary_directory)
            data_path = data_dir / "gsm8k.jsonl"
            data_path.write_text(
                json.dumps({"question": "Local question", "answer": "#### 9"}) + "\n",
                encoding="utf-8",
            )

            try:
                path, source, data_hash = accuracy.materialize_dataset(data_dir, "gsm8k")
            except RuntimeError as error:
                self.fail(f"valid local JSONL without a manifest was rejected: {error}")

            manifest = json.loads(
                (data_dir / "gsm8k.manifest.json").read_text(encoding="utf-8")
            )
            self.assertEqual(path, data_path)
            self.assertEqual(
                source,
                {
                    "name": "local",
                    "subset": "local",
                    "revision": "local",
                    "split": "local",
                },
            )
            self.assertEqual(manifest["dataset_source"], source)
            self.assertEqual(manifest["dataset_path"], str(data_path.resolve()))
            self.assertEqual(manifest["dataset_sha256"], data_hash)
            self.assertEqual(manifest["row_count"], 1)

    def test_local_jsonl_adoption_rejects_invalid_json_and_fields(self):
        accuracy = load_accuracy_module()
        invalid_contents = (
            "not json\n",
            json.dumps({"question": "missing answer"}) + "\n",
            "\n",
        )
        for contents in invalid_contents:
            with self.subTest(contents=contents):
                with tempfile.TemporaryDirectory() as temporary_directory:
                    data_dir = Path(temporary_directory)
                    (data_dir / "gsm8k.jsonl").write_text(contents, encoding="utf-8")
                    with self.assertRaisesRegex(RuntimeError, "JSONL|fields|examples"):
                        accuracy.materialize_dataset(data_dir, "gsm8k")
                    self.assertFalse((data_dir / "gsm8k.manifest.json").exists())

    def test_manifest_reuse_strictly_validates_path_and_row_count(self):
        accuracy = load_accuracy_module()
        source = {
            "name": "local",
            "subset": "local",
            "revision": "local",
            "split": "local",
        }
        corruptions = (
            ("dataset_path", "wrong-path.jsonl", "path"),
            ("row_count", True, "row_count"),
            ("row_count", 0, "row_count"),
            ("row_count", 2, "row_count"),
            ("dataset_sha256", "incorrect-hash", "hash"),
        )
        for key, value, message in corruptions:
            with self.subTest(key=key, value=value):
                with tempfile.TemporaryDirectory() as temporary_directory:
                    data_dir = Path(temporary_directory)
                    data_path = data_dir / "gsm8k.jsonl"
                    data_path.write_text(
                        json.dumps({"question": "q", "answer": "#### 1"}) + "\n",
                        encoding="utf-8",
                    )
                    manifest = {
                        "dataset": "gsm8k",
                        "dataset_path": str(data_path.resolve()),
                        "dataset_sha256": accuracy.dataset_sha256(data_path),
                        "dataset_source": source,
                        "row_count": 1,
                    }
                    manifest[key] = value
                    (data_dir / "gsm8k.manifest.json").write_text(
                        json.dumps(manifest), encoding="utf-8"
                    )

                    with self.assertRaisesRegex(RuntimeError, message):
                        accuracy.materialize_dataset(data_dir, "gsm8k")

    def test_cli_has_real_configurable_dataset_revision_default(self):
        accuracy = load_accuracy_module()
        try:
            with redirect_stderr(io.StringIO()):
                configured = accuracy.parse_args(
                    ["--model", "m", "--dataset-revision", "release"]
                )
        except SystemExit:
            self.fail("CLI does not accept --dataset-revision")
        self.assertEqual(configured.dataset_revision, "release")
        self.assertEqual(accuracy.parse_args(["--model", "m"]).dataset_revision, "master")

    def test_shell_script_runs_three_comparable_cases_with_required_environment(self):
        script = SHELL_PATH.read_text(encoding="utf-8")

        self.assertEqual(script.count('run_case "'), 3)
        for case in ("baseline_0_1", "group_8_num_1", "group_8_num_2"):
            self.assertIn(case, script)
        for variable in (
            "MODEL_PATH",
            "DATA_DIR",
            "RESULTS_DIR",
            "ACCURACY_LIMIT",
            "SEED",
            "MAX_MODEL_LEN",
            "GPU_MEMORY_UTILIZATION",
            "GSM8K_MAX_TOKENS",
            "GPQA_MAX_TOKENS",
            "MAX_ACCURACY_DROP",
            "ASCEND_RT_VISIBLE_DEVICES",
            "VLLM_TARGET_DEVICE=ascend",
            "VLLM_USE_MODELSCOPE=True",
        ):
            self.assertIn(variable, script)
        device = re.search(
            r"^export ASCEND_RT_VISIBLE_DEVICES=\$\{ASCEND_RT_VISIBLE_DEVICES:-(\d+)\}$",
            script,
            re.MULTILINE,
        )
        self.assertIsNotNone(device)
        self.assertEqual(device.group(1), "0")
        default_device = device.group(1)
        resolve_device = lambda override: override or default_device
        self.assertEqual(resolve_device("3"), "3")
        self.assertEqual(resolve_device(""), "0")

        assignments = dict(
            re.findall(r"^(SCRIPT_DIR|REPO_ROOT)=(.+)$", script, re.MULTILINE)
        )
        self.assertIn('${BASH_SOURCE[0]}', assignments["SCRIPT_DIR"])
        self.assertIn('dirname -- "${BASH_SOURCE[0]}"', assignments["SCRIPT_DIR"])
        self.assertIn('${SCRIPT_DIR}/..', assignments["REPO_ROOT"])
        script_dir_index = script.index("SCRIPT_DIR=")
        repo_root_index = script.index("REPO_ROOT=")
        python_path = 'python "${REPO_ROOT}/scripts/accuracy_gsm8k_gpqa.py"'
        python_index = script.index(python_path)
        self.assertLess(script_dir_index, repo_root_index)
        self.assertLess(repo_root_index, python_index)
        self.assertEqual(script.count(python_path), 2)
        self.assertNotIn("python scripts/accuracy_gsm8k_gpqa.py", script)
        self.assertEqual(
            re.findall(r'^run_case "([^"]+)" (\d+) (\d+)$', script, re.MULTILINE),
            [
                ("baseline_0_1", "0", "1"),
                ("group_8_num_1", "8", "1"),
                ("group_8_num_2", "8", "2"),
            ],
        )
        self.assertIn('python "${REPO_ROOT}/scripts/verify_accuracy_matrix.py"', script)
        self.assertIn('--max-accuracy-drop "$MAX_ACCURACY_DROP"', script)


class AccuracyMatrixVerificationTest(unittest.TestCase):
    @staticmethod
    def write_result(root, case, dataset, accuracy, settings):
        case_dir = Path(root) / case
        case_dir.mkdir(parents=True, exist_ok=True)
        total = 10
        payload = {
            "dataset": dataset,
            "dataset_sha256": f"{dataset}-hash",
            "correct": round(accuracy * total),
            "total": total,
            "accuracy": accuracy,
            "offload_settings": settings,
        }
        (case_dir / f"{dataset}-result.json").write_text(
            json.dumps(payload), encoding="utf-8"
        )

    def populate_matrix(self, root, candidate_accuracy=0.8):
        settings = {
            "baseline_0_1": {"group_size": 0, "num_in_group": 1},
            "group_8_num_1": {"group_size": 8, "num_in_group": 1},
            "group_8_num_2": {"group_size": 8, "num_in_group": 2},
        }
        for dataset in ("gsm8k", "gpqa"):
            self.write_result(root, "baseline_0_1", dataset, 0.8, settings["baseline_0_1"])
            self.write_result(root, "group_8_num_1", dataset, candidate_accuracy, settings["group_8_num_1"])
            self.write_result(root, "group_8_num_2", dataset, candidate_accuracy, settings["group_8_num_2"])

    def test_verifier_accepts_equal_accuracy_and_writes_all_summaries(self):
        verify = load_verify_module()
        with tempfile.TemporaryDirectory() as temporary_directory:
            self.populate_matrix(temporary_directory)
            summary = verify.verify_matrix(Path(temporary_directory), 0.0)
            verify.write_summaries(summary, Path(temporary_directory))

            self.assertTrue(summary["all_passed"])
            self.assertEqual(len(summary["rows"]), 6)
            for name in (
                "accuracy-summary.json",
                "accuracy-summary.csv",
                "accuracy-summary.md",
            ):
                self.assertTrue((Path(temporary_directory) / name).is_file())

    def test_verifier_rejects_accuracy_drop_and_metadata_mismatch(self):
        verify = load_verify_module()
        with tempfile.TemporaryDirectory() as temporary_directory:
            self.populate_matrix(temporary_directory, candidate_accuracy=0.7)
            with self.assertRaisesRegex(RuntimeError, "accuracy drop"):
                verify.verify_matrix(Path(temporary_directory), 0.0)

        with tempfile.TemporaryDirectory() as temporary_directory:
            self.populate_matrix(temporary_directory)
            path = Path(temporary_directory) / "group_8_num_2" / "gpqa-result.json"
            payload = json.loads(path.read_text(encoding="utf-8"))
            payload["dataset_sha256"] = "wrong"
            path.write_text(json.dumps(payload), encoding="utf-8")
            with self.assertRaisesRegex(RuntimeError, "dataset hash"):
                verify.verify_matrix(Path(temporary_directory), 0.0)
