from __future__ import annotations

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


SCRIPT = Path(__file__).with_name("fixed_cohort_rank.py")
SPEC = importlib.util.spec_from_file_location("fixed_cohort_rank", 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 _write_trace(path: Path) -> None:
    lengths = [100, 1500, 3000, 6000, 12000, 24000]
    rows = []
    for index in range(120):
        rows.append(
            {
                "timestamp": index * 0.5,
                "sampling_u": (index + 1) / 121,
                "input_length": lengths[index % len(lengths)],
                "output_length": 10,
                "prompt": f"raw prompt {index}",
            }
        )
    path.write_text("".join(json.dumps(row) + "\n" for row in rows))


def test_prepare_protocol_keeps_identical_cohort_and_scales_only_arrivals(
    tmp_path: Path,
) -> None:
    trace = tmp_path / "trace.jsonl"
    _write_trace(trace)
    manifest_path = MODULE.prepare_protocol(
        argparse.Namespace(
            trace=trace,
            expected_trace_sha256="",
            output_root=tmp_path / "protocol",
            cohort_size=24,
            seed=7,
            rate=[0.2, 0.4],
        )
    )
    manifest = json.loads(manifest_path.read_text())
    assert manifest["selection"]["cohort_bin_quotas"] == [4, 4, 4, 4, 4, 4]
    assert manifest["load_contract"]["binary_search"] is False
    assert manifest["load_contract"]["monotonicity_assumed"] is False

    materialized = []
    for record in manifest["rates"].values():
        with Path(record["path"]).open(newline="") as source:
            rows = list(csv.DictReader(source))
        materialized.append(rows)
        duration = float(rows[-1]["arrived_at"]) - float(rows[0]["arrived_at"])
        assert len(rows) / duration == record["offered_request_rate"]
    assert [row["source_row_index"] for row in materialized[0]] == [
        row["source_row_index"] for row in materialized[1]
    ]
    assert [row["num_prefill_tokens"] for row in materialized[0]] == [
        row["num_prefill_tokens"] for row in materialized[1]
    ]
    assert MODULE.audit_protocol(manifest_path)["status"] == "passed"


def test_score_requests_applies_all_slos_to_same_outcomes() -> None:
    requests = [
        {"input_tokens": 100, "success": True, "ttft_ms": 900.0},
        {"input_tokens": 16000, "success": True, "ttft_ms": 2500.0},
        {"input_tokens": 24000, "success": False, "ttft_ms": None},
    ]
    score = MODULE.score_requests(requests)
    assert score["scores"]["linear_8k_primary"]["passed_request_count"] == 2
    assert score["scores"]["legacy_step_1s_2s"]["passed_request_count"] == 1
    assert score["ttft_ms"]["count"] == 2


def test_rank_surface_records_nonmonotone_feasibility_without_binary_search() -> None:
    config_results = [
        {
            "config": {"name": "a", "tp": 4},
            "loads": [
                {
                    "offered_request_rate": 0.2,
                    "scores": {"linear_8k_primary": {"feasible": False, "slo_pass_rate": 0.9}},
                },
                {
                    "offered_request_rate": 0.4,
                    "scores": {"linear_8k_primary": {"feasible": True, "slo_pass_rate": 0.95}},
                },
            ],
        },
        {
            "config": {"name": "b", "tp": 8},
            "loads": [
                {
                    "offered_request_rate": 0.2,
                    "scores": {"linear_8k_primary": {"feasible": True, "slo_pass_rate": 1.0}},
                },
                {
                    "offered_request_rate": 0.4,
                    "scores": {"linear_8k_primary": {"feasible": False, "slo_pass_rate": 0.9}},
                },
            ],
        },
    ]
    ranking = MODULE.rank_surface(config_results, slo_name="linear_8k_primary")
    a = next(item for item in ranking if item["config"]["name"] == "a")
    assert a["maximum_tested_feasible_request_rate"] == 0.4
    assert a["monotonicity_violations"] == [[0.2, 0.4]]


def test_average_ranks_use_midrank_for_capacity_ties() -> None:
    assert MODULE._average_ranks({"a": 0.2, "b": 0.2, "c": 0.1}) == {
        "a": 1.5,
        "b": 1.5,
        "c": 3.0,
    }


def test_community_study_uses_affine_primary_slo(tmp_path: Path) -> None:
    payload = MODULE._community_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["slo"]["ttft_rule"] == {
        "kind": "linear_ms",
        "intercept_ms": 1000.0,
        "per_token_ms": 0.125,
    }
    assert payload["trace"]["request_mode"] == "raw_completion"
    assert payload["trace"]["completion_tokens_override"] == 1


def test_warmup_is_disjoint_and_covers_each_input_bin() -> None:
    lengths = [100, 1500, 3000, 6000, 12000, 24000]
    requests = []
    cohort_ids = set()
    for bin_index, length in enumerate(lengths):
        for suffix in ("cohort", "warmup"):
            row_id = f"{bin_index}-{suffix}"
            if suffix == "cohort":
                cohort_ids.add(row_id)
            requests.append(
                MODULE.TraceRequest(
                    row_id=row_id,
                    arrival_s=10.0,
                    sampling_u=0.5,
                    body={"prompt": row_id},
                    prompt_tokens_hint=length,
                    completion_tokens_hint=1,
                )
            )
    selected = MODULE._select_warmup_requests(requests, cohort_ids=cohort_ids)
    assert len(selected) == len(MODULE.INPUT_BINS) - 1
    assert not cohort_ids & {request.row_id for request in selected}
    assert [MODULE._input_bin(int(request.prompt_tokens_hint)) for request in selected] == list(
        range(len(MODULE.INPUT_BINS) - 1)
    )
    assert {request.arrival_s for request in selected} == {0.0}
