from __future__ import annotations

import importlib.util
import json
import sys
from pathlib import Path


SCRIPT = Path(__file__).with_name("community_prefill_grid.py")
SPEC = importlib.util.spec_from_file_location("community_prefill_grid", SCRIPT)
MODULE = importlib.util.module_from_spec(SPEC)
assert SPEC.loader is not None
sys.modules[SPEC.name] = MODULE
SPEC.loader.exec_module(MODULE)


def test_grid_matches_frozen_frontier_surface() -> None:
    assert len(MODULE.GRID) == 8
    assert {item.name for item in MODULE.GRID} == {
        f"tp{tp}_mns{mns}_mbt{mbt}"
        for tp in (4, 8)
        for mns in (64, 128)
        for mbt in (8192, 16384)
    }
    assert {item.num_gpu_blocks for item in MODULE.GRID if item.tp == 4} == {26101}
    assert {item.num_gpu_blocks for item in MODULE.GRID if item.tp == 8} == {62351}


def test_study_uses_raw_prompt_and_fixed_serving_contract(tmp_path: Path) -> None:
    payload = MODULE.study_payload(
        tp=8,
        repo=tmp_path,
        python=tmp_path / "python",
        vllm=tmp_path / "vllm",
        model=tmp_path / "model",
        trace=tmp_path / "trace.jsonl",
        windows=tmp_path / "windows.json",
        port=18918,
    )
    assert payload["trace"]["request_mode"] == "raw_completion"
    assert payload["trace"]["completion_tokens_override"] == 1
    assert payload["trace"]["replay_time_scale"] == 1.0
    assert payload["engine"]["base_flags"]["enable-expert-parallel"] is True
    assert payload["engine"]["base_flags"]["disable-custom-all-reduce"] is True
    assert payload["engine"]["base_flags"]["num-gpu-blocks-override"] == 62351
    assert payload["engine"]["base_flags"]["enable-prefix-caching"] is False


def test_capacity_interval_replays_binary_decisions() -> None:
    probes = [
        {"threshold": 0.0625, "feasible": False},
        {"threshold": 0.03125, "feasible": False},
        {"threshold": 0.015625, "feasible": True},
        {"threshold": 0.0234375, "feasible": False},
        {"threshold": 0.01953125, "feasible": True},
        {"threshold": 0.021484375, "feasible": True},
    ]
    assert MODULE.capacity_interval(probes) == [0.021484375, 0.0234375]


def test_manifest_json_round_trip(tmp_path: Path) -> None:
    path = tmp_path / "x.json"
    MODULE.write_json(path, {"schema": MODULE.SCHEMA})
    assert json.loads(path.read_text()) == {"schema": MODULE.SCHEMA}