from __future__ import annotations

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


SCRIPT = Path(__file__).with_name("frontier_prefill_grid.py")
SPEC = importlib.util.spec_from_file_location("frontier_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_binary_search_lattice_covers_six_probe_tree() -> None:
    anchors = MODULE.binary_search_lattice()
    assert len(anchors) == 63
    assert anchors[0] == 0.001953125
    assert anchors[-1] == 0.123046875
    assert 0.0625 in anchors


def test_prepare_preserves_stable_arrival_order(tmp_path: Path) -> None:
    trace = tmp_path / "trace.jsonl"
    rows = [
        {"timestamp": 1.0, "sampling_u": 0.02, "input_length": 9000, "id": "a"},
        {"timestamp": 0.5, "sampling_u": 0.03, "input_length": 100, "id": "b"},
        {"timestamp": 1.0, "sampling_u": 0.01, "input_length": 200, "id": "c"},
        {"timestamp": 0.2, "sampling_u": 0.01, "input_length": 40000, "id": "x"},
    ]
    trace.write_text("".join(json.dumps(row) + "\n" for row in rows))
    manifest_path = MODULE.prepare_traces(trace, tmp_path / "out", "")
    manifest = json.loads(manifest_path.read_text())
    record = manifest["anchors"][MODULE.anchor_key(0.0625)]
    with Path(record["path"]).open(newline="") as source:
        selected = list(csv.DictReader(source))
    assert [row["source_request_id"] for row in selected] == ["b", "a", "c"]
    assert [float(row["slo_ttft_ms"]) for row in selected] == [1000, 2000, 1000]


def test_score_uses_trace_slo_and_request_id_mapping(tmp_path: Path) -> None:
    trace = tmp_path / "trace.csv"
    with trace.open("w", newline="") as output:
        writer = csv.DictWriter(
            output,
            fieldnames=["arrived_at", "num_prefill_tokens", "num_decode_tokens", "slo_ttft_ms"],
        )
        writer.writeheader()
        writer.writerows(
            [
                {"arrived_at": 0, "num_prefill_tokens": 10, "num_decode_tokens": 1, "slo_ttft_ms": 1000},
                {"arrived_at": 1, "num_prefill_tokens": 9000, "num_decode_tokens": 1, "slo_ttft_ms": 2000},
            ]
        )
    metrics = tmp_path / "request_metrics.csv"
    with metrics.open("w", newline="") as output:
        writer = csv.DictWriter(output, fieldnames=["Request Id", "ttft"])
        writer.writeheader()
        writer.writerows([{"Request Id": 1, "ttft": 1900}, {"Request Id": 0, "ttft": 1100}])
    score = MODULE.score_request_metrics(trace, metrics)
    assert score["passed_request_count"] == 1
    assert score["slo_pass_rate"] == 0.5
    assert score["feasible"] is False


def test_build_command_preserves_trace_arrivals(tmp_path: Path) -> None:
    config = MODULE.GRID[0]
    command = MODULE.build_command(
        python=tmp_path / "python",
        frontier_source=tmp_path / "frontier",
        profile_root=tmp_path / "profiles",
        profile_paths={
            "linear": tmp_path / "linear.csv",
            "attention": tmp_path / "attention.csv",
            "moe": tmp_path / "moe.csv",
            "all_reduce": tmp_path / "all_reduce.csv",
        },
        trace=tmp_path / "trace.csv",
        config=config,
        probe_dir=tmp_path / "probe",
        run_id="arrival-contract",
        cache_root=tmp_path / "cache",
    )

    assert command[command.index("--simulation_mode") + 1] == "offline"
    assert "--offline_use_generated_request_arrivals" in command