#!/usr/bin/env python3

from __future__ import annotations

import importlib.util
import sys
from pathlib import Path


ROOT = Path(__file__).parent


def load(name: str):
    path = ROOT / name
    spec = importlib.util.spec_from_file_location(path.stem, path)
    assert spec and spec.loader
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


def test_percentile() -> None:
    client = load("qwen30_prefill_client.py")
    assert client.percentile([4.0, 1.0, 3.0, 2.0], 0.50) == 2.0
    assert client.percentile([4.0, 1.0, 3.0, 2.0], 0.95) == 4.0
    assert client.percentile([], 0.95) is None
    assert client.ttft_slo_ms(512) == 1064.0
    assert client.ttft_slo_ms(2048) == 1256.0


def test_grid_and_trace(tmp_path: Path) -> None:
    surface = load("run_frontier_qwen30_prefill_surface.py")
    assert len(surface.GRID) == 12
    assert {config.tp for config in surface.GRID} == {1, 2, 4}
    trace = tmp_path / "trace.csv"
    surface.write_trace(trace, requests=3, rate=4.0)
    lines = trace.read_text().splitlines()
    assert len(lines) == 4
    assert lines[1].split(",")[:3] == ["0.000000000000", "2048", "1"]
    assert lines[3].split(",")[:3] == ["0.500000000000", "2048", "1"]


def test_kendall_tau_b() -> None:
    analysis = load("analyze_qwen30_prefill_fidelity.py")
    assert analysis.kendall_tau_b([1, 2, 3], [1, 2, 3])["kendall_tau_b"] == 1
    assert analysis.kendall_tau_b([1, 2, 3], [3, 2, 1])["kendall_tau_b"] == -1


def test_configure_cc_command(tmp_path: Path) -> None:
    surface = load("run_frontier_qwen30_prefill_surface.py")
    base = ["python", "--cc_backend_config_type", "analytical", "--other", "x"]
    analytical = surface.configure_cc_command(
        base, backend="analytical", allreduce_csv=None, cache=tmp_path
    )
    assert analytical == base
    profile = tmp_path / "all_reduce.csv"
    profile.write_text("header\n")
    vidur = surface.configure_cc_command(
        base, backend="vidur", allreduce_csv=profile, cache=tmp_path / "cache"
    )
    assert vidur[2] == "vidur"
    assert "--vidur_cc_backend_config_all_reduce_input_file" in vidur
    assert str(profile) in vidur