import importlib.util
import re
import unittest
from pathlib import Path


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


ROOT = find_repo_root()


def load_script(name: str):
    path = ROOT / "scripts" / name
    spec = importlib.util.spec_from_file_location(path.stem, path)
    module = importlib.util.module_from_spec(spec)
    assert spec.loader is not None
    spec.loader.exec_module(module)
    return module


class BenchmarkUnitTests(unittest.TestCase):
    def test_detect_npu_index_handles_reallocated_physical_card(self):
        detector = load_script("detect_npu_index.py")
        old_mapping = """
NPU ID  Chip ID  Chip Logic ID  Chip Name
7       0        0              Ascend 910B2C
7       1        -              Mcu
"""
        new_mapping = """
NPU ID  Chip ID  Chip Logic ID  Chip Name
0       0        0              Ascend 910B2C
0       1        -              Mcu
"""
        self.assertEqual(detector.parse_mapping(old_mapping), 7)
        self.assertEqual(detector.parse_mapping(new_mapping), 0)
        with self.assertRaisesRegex(RuntimeError, "Ascend NPU mapping"):
            detector.parse_mapping("NPU ID Chip ID\n")

    def test_parse_npu_smi_hbm_tracks_device_across_target_machine_rows(self):
        bench = load_script("bench_one.py")
        output = """
+--------------------------------------------------------------------------------------------------+
| 7 910B2C | Alarm | 70.1 42 0 0 / 0 |
| 0 | 0000:CF:00.0 | 0 0 / 0 3413 / 65536 |
| 6 910B2C | OK | 68.2 40 0 0 / 0 |
| 0 | 0000:AF:00.0 | 0 0 / 0 9876 / 65536 |
+--------------------------------------------------------------------------------------------------+
| NPU | Chip | Process id | Process name | Process memory(MB) |
| 7 | 0 | 65536 | python | 9999 |
"""
        self.assertEqual(
            bench.parse_npu_smi_hbm(output, npu_index=7),
            {"used_mib": 3413, "total_mib": 65536},
        )
        self.assertIsNone(bench.parse_npu_smi_hbm("unexpected output", 7))

    def test_parse_npu_smi_hbm_does_not_cross_device_or_process_rows(self):
        bench = load_script("bench_one.py")
        output = """
| 7 910B2C | Alarm | 70.1 42 0 0 / 0 |
| 0 | not-a-pci-address | 0 0 / 0 1111 / 65536 |
| 6 910B2C | OK | 68.2 40 0 0 / 0 |
| 0 | 0000:AF:00.0 | 0 0 / 0 9876 / 65536 |
| NPU | Chip | Process id | Process name | Process memory(MB) |
| 7 | 0 | 12345 | python | 7777 / 65536 |
"""
        self.assertIsNone(bench.parse_npu_smi_hbm(output, npu_index=7))
        self.assertEqual(
            bench.parse_npu_smi_hbm(output, npu_index=6),
            {"used_mib": 9876, "total_mib": 65536},
        )

    def test_hbm_summary_does_not_replace_failed_before_sample(self):
        bench = load_script("bench_one.py")
        sampler = bench.ResourceSampler(7)
        sampler.samples = [
            {"used_mib": None, "total_mib": None, "raw": "parse failed"},
            {"used_mib": 1024, "total_mib": 65536},
            {"used_mib": 2048, "total_mib": 65536},
        ]
        self.assertEqual(sampler.hbm_summary(), (None, 2048, 2048))

    def test_resource_sampling_uses_phase_boundaries_not_background_npu_smi(self):
        script = self.read_benchmark_source()
        self.assertNotIn("threading.Thread", script)
        self.assertNotIn("while not self._stop", script)
        self.assertEqual(script.count("sampler.checkpoint("), 2)
        self.assertIn('sampler.checkpoint("before_load")', script)
        self.assertIn('sampler.checkpoint("after_generation")', script)

    @staticmethod
    def read_benchmark_source():
        return (ROOT / "scripts" / "bench_one.py").read_text(encoding="utf-8")

    def test_comparison_formulas_handle_zero_missing_and_negative_changes(self):
        summary = load_script("bench_param_matrix.py")
        self.assertIsNone(summary.percent_change(12.0, 0.0))
        self.assertIsNone(summary.percent_change(None, 10.0))
        self.assertEqual(summary.percent_change(8.0, 10.0), -20.0)
        self.assertIsNone(summary.hbm_saved_gib(None, 1024.0))
        self.assertEqual(summary.hbm_saved_gib(4096.0, 1024.0), 3.0)

    def test_aggregation_groups_labels_and_uses_median(self):
        summary = load_script("bench_param_matrix.py")
        records = [
            {"label": "baseline_0_1", "output_tokens_per_second": 10.0,
             "npu_hbm_peak_mib": 4096.0},
            {"label": "group_8_num_1", "output_tokens_per_second": 5.0,
             "npu_hbm_peak_mib": 2048.0},
            {"label": "baseline_0_1", "output_tokens_per_second": 30.0,
             "npu_hbm_peak_mib": 6144.0},
            {"label": "group_8_num_1", "output_tokens_per_second": 25.0,
             "npu_hbm_peak_mib": 3072.0},
            {"label": "baseline_0_1", "output_tokens_per_second": 20.0,
             "npu_hbm_peak_mib": 5120.0},
            {"label": "group_8_num_1", "output_tokens_per_second": 15.0,
             "npu_hbm_peak_mib": 1024.0},
        ]
        rows = summary.aggregate_records(records, "baseline_0_1")
        by_label = {row["label"]: row for row in rows}
        self.assertEqual(by_label["baseline_0_1"]["repeats"], 3)
        self.assertEqual(
            by_label["baseline_0_1"]["output_tokens_per_second"], 20.0
        )
        self.assertEqual(
            by_label["group_8_num_1"]["npu_hbm_peak_mib"], 2048.0
        )
        self.assertEqual(
            by_label["group_8_num_1"]["throughput_change_percent"], -25.0
        )
        self.assertEqual(by_label["group_8_num_1"]["hbm_saved_gib"], 3.0)


class ScriptAcceptanceTests(unittest.TestCase):
    def read(self, name: str) -> str:
        return (ROOT / "scripts" / name).read_text(encoding="utf-8")

    def test_scripts_are_portable_and_have_no_old_smoke_path(self):
        names = (
            "bench_one.py",
            "bench_param_matrix.py",
            "run_benchmark_all.sh",
            "run_reproduce_all.sh",
        )
        for name in names:
            with self.subTest(name=name):
                text = self.read(name)
                self.assertNotIn("/data/vllm-ascend-smoke", text)
        benchmark = self.read("run_benchmark_all.sh")
        self.assertIn('BASH_SOURCE[0]', benchmark)
        self.assertIn(
            "CANN_ENV=${CANN_ENV:-"
            "/usr/local/Ascend/cann-8.5.1/cann-8.5.1/set_env.sh}",
            benchmark,
        )
        self.assertIn(
            "VENV_ACTIVATE=${VENV_ACTIVATE:-"
            "/data/vllm-ascend-v0.18.0/.venv/bin/activate}",
            benchmark,
        )
        self.assertIn('source "${CANN_ENV}"', benchmark)
        self.assertIn('source "${VENV_ACTIVATE}"', benchmark)

    def test_runtime_environment_paths_can_be_overridden_or_disabled(self):
        for name in (
            "run_benchmark_all.sh",
            "run_accuracy_gsm8k_gpqa.sh",
            "run_reproduce_all.sh",
        ):
            with self.subTest(name=name):
                script = self.read(name)
                self.assertIn("CANN_ENV=${CANN_ENV:-", script)
                self.assertIn("VENV_ACTIVATE=${VENV_ACTIVATE:-", script)
                self.assertIn('source "${CANN_ENV}"', script)
                self.assertIn('source "${VENV_ACTIVATE}"', script)
                self.assertIn(
                    "VLLM_ASCEND_CUSTOM_OP_ROOT="
                    "${VLLM_ASCEND_CUSTOM_OP_ROOT-",
                    script,
                )
                self.assertIn(
                    'if [[ -n "${VLLM_ASCEND_CUSTOM_OP_ROOT}" ]]', script
                )

    def test_benchmark_has_three_cases_rotation_and_fresh_processes(self):
        script = self.read("run_benchmark_all.sh")
        for case in ("baseline_0_1:0:1", "group_8_num_1:8:1", "group_8_num_2:8:2"):
            self.assertIn(case, script)
        for variable in (
            "MODEL_PATH", "RESULTS_DIR", "REPEATS", "WARMUP_ROUNDS",
            "MEASURED_ROUNDS", "PROMPT_COUNT", "MAX_MODEL_LEN",
            "MAX_TOKENS", "GPU_MEMORY_UTILIZATION", "SEED",
            "ASCEND_RT_VISIBLE_DEVICES",
        ):
            self.assertIn(variable, script)
        strict_mode = script.index("set -euo pipefail")
        repeat_loop = script.index("for ((repeat=0; repeat<REPEATS; repeat++)); do")
        offset_loop = script.index(
            "for ((offset=0; offset<${#CASES[@]}; offset++)); do"
        )
        python_case = script.index('python "${SCRIPT_DIR}/bench_one.py"')
        first_inner_done = script.index("    done", python_case)
        first_outer_done = script.index("done", first_inner_done + len("    done"))
        self.assertLess(strict_mode, python_case)
        self.assertLess(repeat_loop, offset_loop)
        self.assertLess(offset_loop, python_case)
        self.assertLess(python_case, first_inner_done)
        self.assertLess(first_inner_done, first_outer_done)
        self.assertIn("start=$((repeat % ${#CASES[@]}))", script)
        self.assertIn("index=$(((start + offset) % ${#CASES[@]}))", script)
        self.assertIn(
            'IFS=: read -r label group_size num_in_group <<< "${CASES[index]}"',
            script,
        )
        case_pipeline = re.search(
            r'python "\$\{SCRIPT_DIR\}/bench_one\.py" \\\n'
            r'(?P<arguments>(?:\s+--[^\n]+\\\n)+)'
            r'\s+--npu-index "\$\{NPU_INDEX\}" 2>&1 \| tee -a '
            r'"\$\{RUN_LOG\}"',
            script,
        )
        self.assertIsNotNone(case_pipeline)
        arguments = case_pipeline.group("arguments")
        self.assertIn('--label "${label}"', arguments)
        self.assertIn('--offload-group-size "${group_size}"', arguments)
        self.assertIn('--offload-num-in-group "${num_in_group}"', arguments)
        self.assertIn(
            '--gpu-memory-utilization "${GPU_MEMORY_UTILIZATION}"', arguments
        )
        self.assertNotIn("|| true", case_pipeline.group(0))
        self.assertIn("set -euo pipefail", script[:python_case])
        self.assertNotIn("|| true", script)

    def test_shell_paths_are_derived_from_bash_source(self):
        benchmark = self.read("run_benchmark_all.sh")
        reproduce = self.read("run_reproduce_all.sh")
        exact_script_dir = (
            'SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"'
        )
        self.assertIn(exact_script_dir, benchmark)
        self.assertIn(exact_script_dir, reproduce)
        self.assertIn(
            'REPO_ROOT="$(cd -- "${SCRIPT_DIR}/.." && pwd)"', reproduce
        )

    def test_runtime_visible_device_and_npu_smi_index_are_independent(self):
        benchmark = self.read("run_benchmark_all.sh")
        reproduce = self.read("run_reproduce_all.sh")
        accuracy = self.read("run_accuracy_gsm8k_gpqa.sh")
        health = self.read("check_npu_health.sh")
        for name, script in (
            ("run_benchmark_all.sh", benchmark),
            ("run_reproduce_all.sh", reproduce),
            ("run_accuracy_gsm8k_gpqa.sh", accuracy),
            ("check_npu_health.sh", health),
        ):
            with self.subTest(name=name):
                self.assertIn("detect_npu_index.py", script)
                self.assertNotIn("NPU_INDEX=${NPU_INDEX:-7}", script)
        for script in (benchmark, reproduce, accuracy):
            self.assertIn(
                "export ASCEND_RT_VISIBLE_DEVICES=${ASCEND_RT_VISIBLE_DEVICES:-0}",
                script,
            )
        self.assertIn('--npu-index "${NPU_INDEX}"', benchmark)

    def test_runtime_scripts_export_bundled_custom_operator_paths(self):
        expected_root = (
            "/data/vllm-ascend-v0.18.0/vllm_ascend/"
            "_cann_ops_custom/vendors/vllm-ascend"
        )
        for name in (
            "run_benchmark_all.sh",
            "run_accuracy_gsm8k_gpqa.sh",
            "run_reproduce_all.sh",
        ):
            with self.subTest(name=name):
                script = self.read(name)
                self.assertIn("VLLM_ASCEND_CUSTOM_OP_ROOT", script)
                self.assertIn(expected_root, script)
                self.assertIn("ASCEND_CUSTOM_OPP_PATH", script)
                self.assertIn("LD_LIBRARY_PATH", script)
                self.assertIn("TRITON_CACHE_DIR", script)

    def test_runtime_scripts_fail_fast_on_unhealthy_npu(self):
        health_check = self.read("check_npu_health.sh")
        self.assertIn("npu-smi info -t health", health_check)
        self.assertIn("Health Status", health_check)
        self.assertIn("exit 42", health_check)
        for name in (
            "run_benchmark_all.sh",
            "run_accuracy_gsm8k_gpqa.sh",
        ):
            with self.subTest(name=name):
                script = self.read(name)
                self.assertIn('bash "${SCRIPT_DIR}/check_npu_health.sh"', script)
                self.assertIn("NPU_CHIP_INDEX", script)

    def test_bench_one_records_required_metrics_and_deterministic_settings(self):
        script = self.read("bench_one.py")
        for metric in (
            "load_seconds", "generate_seconds", "request_count", "output_tokens",
            "output_tokens_per_second", "requests_per_second", "round_latencies_seconds",
            "latency_p50_seconds", "latency_p95_seconds", "process_rss_before_mib",
            "process_rss_peak_mib", "process_rss_after_mib", "npu_hbm_before_mib",
            "npu_hbm_peak_mib", "npu_hbm_after_mib", "offload_group_size",
            "offload_num_in_group", "npu_hbm_samples",
        ):
            self.assertIn(metric, script)
        self.assertIn("temperature=0", script)
        self.assertIn("enforce_eager=True", script)

    def test_unified_script_runs_performance_before_accuracy_and_writes_manifest(self):
        script = self.read("run_reproduce_all.sh")
        performance = script.index("run_benchmark_all.sh")
        accuracy = script.index("run_accuracy_gsm8k_gpqa.sh")
        manifest = script.index("reproduction-manifest.json")
        self.assertLess(performance, accuracy)
        self.assertLess(accuracy, manifest)
        for field in (
            "started_at", "finished_at", "commands", "package_versions",
            "source_commit", "model_path", "artifacts",
        ):
            self.assertIn(field, script)
        self.assertIn("import shlex", script)
        self.assertIn(
            'repo_root / "scripts" / "run_benchmark_all.sh"', script
        )
        self.assertIn(
            'repo_root / "scripts" / "run_accuracy_gsm8k_gpqa.sh"', script
        )
        self.assertIn('f"RESULTS_DIR={performance_dir}"', script)
        self.assertIn('f"RESULTS_DIR={accuracy_dir}"', script)
        self.assertIn("shlex.join", script)
        self.assertNotIn("bash scripts/run_benchmark_all.sh", script)
        self.assertNotIn("bash scripts/run_accuracy_gsm8k_gpqa.sh", script)
        for artifact in (
            "accuracy_summary_json",
            "accuracy_summary_csv",
            "accuracy_summary_markdown",
        ):
            self.assertIn(artifact, script)

    def test_v018_runtime_compatibility_install_is_reproducible(self):
        script = self.read("install_v018_runtime_compat.sh")
        self.assertIn("set -euo pipefail", script)
        self.assertIn(
            "/usr/local/Ascend/cann-8.5.1/cann-8.5.1/set_env.sh", script
        )
        self.assertIn("/data/vllm-ascend-v0.18.0/.venv/bin/activate", script)
        self.assertIn("triton_ascend-3.2.0.dev20260322", script)
        self.assertIn("torch-npu==2.9.0", script)
        self.assertIn(
            "https://mirrors.huaweicloud.com/repository/pypi/simple", script
        )
        self.assertIn("vllm-ascend.obs.cn-north-4.myhuaweicloud.com", script)
        self.assertIn("--no-deps", script)
        self.assertIn('realpath "${REPO_ROOT}"', script)
        self.assertIn('pip install --no-build-isolation --no-deps -e', script)
        self.assertIn('version != "3.2.0.dev20260322"', script)
        self.assertIn('torch_npu_version != "2.9.0"', script)

    def test_manifest_records_triton_ascend_runtime_version(self):
        script = self.read("run_reproduce_all.sh")
        self.assertIn('"triton-ascend"', script)


if __name__ == "__main__":
    unittest.main()
