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}


def test_objective_fidelity_separates_raw_and_per_gpu_rankings() -> None:
    frontier = [
        {
            "config": {"name": "tp8", "tp": 8},
            "maximum_tested_feasible_request_rate": 0.4,
            "maximum_tested_feasible_request_rate_per_gpu": 0.05,
        },
        {
            "config": {"name": "tp4", "tp": 4},
            "maximum_tested_feasible_request_rate": 0.3,
            "maximum_tested_feasible_request_rate_per_gpu": 0.075,
        },
    ]
    real = [
        {
            "config": {"name": "tp8", "tp": 8},
            "maximum_tested_feasible_request_rate": 0.2,
            "maximum_tested_feasible_request_rate_per_gpu": 0.025,
        },
        {
            "config": {"name": "tp4", "tp": 4},
            "maximum_tested_feasible_request_rate": 0.3,
            "maximum_tested_feasible_request_rate_per_gpu": 0.075,
        },
    ]
    raw = MODULE._objective_fidelity(
        frontier,
        real,
        field="maximum_tested_feasible_request_rate",
    )
    efficiency = MODULE._objective_fidelity(
        frontier,
        real,
        field="maximum_tested_feasible_request_rate_per_gpu",
    )
    assert raw["frontier_top1_set"] == ["tp8"]
    assert raw["real_top1_set"] == ["tp4"]
    assert abs(raw["top1_regret_fraction_worst_tie_break"] - 1 / 3) < 1e-12
    assert efficiency["frontier_top1_set"] == efficiency["real_top1_set"] == ["tp4"]


def test_objective_fidelity_marks_null_capacity_unrankable() -> None:
    ranking = [
        {
            "config": {"name": "a", "tp": 4},
            "maximum_tested_feasible_request_rate": None,
        }
    ]
    result = MODULE._objective_fidelity(
        ranking,
        ranking,
        field="maximum_tested_feasible_request_rate",
    )
    assert result["rankable"] is False


def test_boundary_repeat_rates_keep_every_transition_and_censored_fallback() -> None:
    def result(feasibility: list[bool]) -> dict[str, object]:
        return {
            "config": {"name": "test"},
            "loads": [
                {
                    "offered_request_rate": rate,
                    "scores": {MODULE.PRIMARY_SLO: {"feasible": feasible}},
                }
                for rate, feasible in zip((0.1, 0.2, 0.3, 0.4), feasibility)
            ],
        }

    assert MODULE._boundary_repeat_rates(result([True, True, False, False])) == [
        0.2,
        0.3,
    ]
    assert MODULE._boundary_repeat_rates(result([True, False, True, False])) == [
        0.1,
        0.2,
        0.3,
        0.4,
    ]
    assert MODULE._boundary_repeat_rates(result([True, True, True, True])) == [
        0.3,
        0.4,
    ]
    assert MODULE._boundary_repeat_rates(result([False, False, False, False])) == [
        0.1,
        0.2,
    ]

    sensitivity = result([True, True, False, False])
    for load, feasible in zip(sensitivity["loads"], [True, True, True, False]):
        load["scores"]["linear_6k"] = {"feasible": feasible}
    assert MODULE._boundary_repeat_rates(
        sensitivity, slo_name="linear_6k"
    ) == [0.3, 0.4]
    assert MODULE._slo_has_feasible_load(
        sensitivity, slo_name="linear_6k"
    )
    assert not MODULE._slo_has_feasible_load(
        result([False, False, False, False]),
        slo_name=MODULE.PRIMARY_SLO,
    )


def test_request_slo_classification_exposes_aggregate_error_cancellation() -> None:
    def results(ttfts: list[float]) -> list[dict[str, object]]:
        return [
            {
                "config": {"name": "a"},
                "loads": [
                    {
                        "offered_request_rate": 0.1,
                        "requests": [
                            {
                                "request_id": str(index),
                                "input_tokens": 0,
                                "success": True,
                                "ttft_ms": ttft,
                            }
                            for index, ttft in enumerate(ttfts)
                        ],
                    }
                ],
            }
        ]

    classification = MODULE._request_slo_classification(
        results([900.0, 1100.0]), results([1100.0, 900.0])
    )[MODULE.PRIMARY_SLO]
    assert classification["frontier_failure_requests"] == 1
    assert classification["real_failure_requests"] == 1
    assert classification["failure_set_jaccard"] == 0.0
    assert classification["label_accuracy"] == 0.0
    assert classification["false_slo_pass_requests"] == 1
    assert classification["false_slo_fail_requests"] == 1


def test_frontier_freeze_validation_uses_manifest_rates(tmp_path: Path) -> None:
    protocol_path = tmp_path / "protocol.json"
    protocol = {
        "rates": {
            "low": {"offered_request_rate": 0.2},
            "high": {"offered_request_rate": 0.4},
        }
    }
    protocol_path.write_text(json.dumps(protocol))
    freeze_path = tmp_path / "freeze.json"
    freeze_path.write_text(
        json.dumps(
            {
                "schema": MODULE.FRONTIER_SCHEMA,
                "status": "frozen_before_community_run",
                "protocol_manifest_sha256": MODULE.sha256(protocol_path),
                "config_results": [
                    {
                        "config": {"name": config.name},
                        "loads": [
                            {"offered_request_rate": 0.2},
                            {"offered_request_rate": 0.4},
                        ],
                    }
                    for config in MODULE.frontier_grid.GRID
                ],
            }
        )
    )
    validated = MODULE._validate_frontier_freeze(freeze_path, protocol_path)
    assert validated["protocol_manifest_sha256"] == MODULE.sha256(protocol_path)


def test_compare_trials_separates_real_label_flips_from_stable_frontier_errors(
    tmp_path: Path,
) -> None:
    config = {"name": "a", "tp": 4}

    def config_result(ttfts: list[float]) -> dict[str, object]:
        requests = [
            {
                "request_id": str(index),
                "input_tokens": 0,
                "success": True,
                "ttft_ms": ttft,
            }
            for index, ttft in enumerate(ttfts)
        ]
        load = {
            "offered_request_rate": 0.1,
            "requests": requests,
            **MODULE.score_requests(requests),
        }
        return {"config": config, "loads": [load]}

    frontier_result = config_result([900.0, 1100.0])
    first_result = config_result([900.0, 900.0])
    repeat_result = config_result([900.0, 1100.0])

    def freeze(path: Path, schema: str, result: dict[str, object]) -> Path:
        payload = {
            "schema": schema,
            "status": "completed",
            "config_results": [result],
            "rankings": {
                name: MODULE.rank_surface([result], slo_name=name)
                for name in MODULE.SLO_VARIANTS
            },
        }
        path.write_text(json.dumps(payload))
        return path

    frontier_path = freeze(tmp_path / "frontier.json", MODULE.FRONTIER_SCHEMA, frontier_result)
    first_path = freeze(tmp_path / "first.json", MODULE.COMMUNITY_SCHEMA, first_result)
    repeat_path = freeze(tmp_path / "repeat.json", MODULE.COMMUNITY_SCHEMA, repeat_result)
    output = tmp_path / "stability.json"
    MODULE.compare_trials(
        argparse.Namespace(
            frontier_freeze=frontier_path,
            first_community_freeze=first_path,
            repeat_community_freeze=repeat_path,
            output=output,
        )
    )
    result = json.loads(output.read_text())
    primary = result["variants"][MODULE.PRIMARY_SLO]
    assert primary["shared_requests"] == 2
    assert primary["real_trial_label_flips"] == 1
    assert primary["frontier_wrong_both_stable_real_label"] == 0
    assert primary["real_trial_load_label_accuracy"] == 0.0
    assert result["primary_original_grid_capacity_stability"]["stable_configs"] == 0
