#!/usr/bin/env python3
from __future__ import annotations

import copy
import importlib.util
from pathlib import Path
from types import SimpleNamespace


HERE = Path(__file__).resolve().parent
ROOT = HERE.parents[1]


def load(name: str, filename: str):
    spec = importlib.util.spec_from_file_location(name, HERE / filename)
    module = importlib.util.module_from_spec(spec)
    assert spec.loader is not None
    spec.loader.exec_module(module)
    return module


def record(*, waiting: int, running: int, tokens: int) -> dict:
    return {
        "queues": {"waiting": waiting, "deferred": 0, "running": running},
        "prefill_tokens": tokens,
        "decode_tokens": 0,
        "kv": {"usage": 0.5},
        "preemptions": 0,
    }


def fake_run(
    config: str,
    repetition: int,
    *,
    goodput: float,
    mns_score: float = 0.0,
    mbbt_score: float = 0.0,
    ambiguous: float = 0.0,
) -> dict:
    binding = {
        "mns_exclusive_fraction": mns_score,
        "mbbt_exclusive_fraction": mbbt_score,
        "both_fraction": ambiguous,
        "waiting_unresolved_fraction": 0.0,
        "kv_usage_max": 0.5,
        "preemptions": 0,
    }
    phases = {
        phase: {
            "mns_exclusive_fraction": mns_score,
            "mbbt_exclusive_fraction": mbbt_score,
        }
        for phase in ("0.25", "0.50", "0.75", "1.00")
    }
    return {
        "config_id": config,
        "repetition": repetition,
        "outcome": {"slo_goodput_req_s": goodput},
        "binding": binding,
        "phases": phases,
    }


def main() -> None:
    analysis = load("action_aware_analysis", "analyze_pilot.py")
    summary = analysis.binding_summary(
        [
            record(waiting=1, running=16, tokens=8),
            record(waiting=1, running=8, tokens=32),
            record(waiting=1, running=16, tokens=32),
            record(waiting=1, running=8, tokens=8),
            record(waiting=0, running=8, tokens=8),
        ],
        mns=16,
        mbbt=32,
    )
    assert summary["mns_exclusive_count"] == 1
    assert summary["mbbt_exclusive_count"] == 1
    assert summary["both_count"] == 1
    assert summary["waiting_unresolved_count"] == 1
    assert summary["waiting_count"] == 4

    # A per-step stream may have a submit gap above one second when the
    # preceding model execution itself spans that interval.  Such a gap is
    # covered telemetry, not a dropped-record interval.
    asynchronous = [
        {"submit_mono_ns": 0, "complete_mono_ns": 1_200_000_000},
        {"submit_mono_ns": 1_100_000_000, "complete_mono_ns": 1_300_000_000},
    ]
    coverage, covered = analysis.telemetry_coverage(
        asynchronous, start_ns=0, end_ns=1_100_000_000
    )
    assert coverage["max_internal_submit_gap_s"] == 1.1
    assert coverage["max_uncovered_gap_s"] == 0.0
    assert covered
    missing = copy.deepcopy(asynchronous)
    missing[0]["complete_mono_ns"] = 0
    assert not analysis.telemetry_coverage(
        missing, start_ns=0, end_ns=1_100_000_000
    )[1]

    mechanism = analysis.mechanism_summary(
        [
            {
                "model_executed": True,
                "submit_mono_ns": 0,
                "complete_mono_ns": 2_000_000,
                "prefill_tokens": 8,
                "prefill_requests": 2,
                "chunked_prefill": {
                    "first": 1,
                    "middle": 0,
                    "final": 0,
                    "unsplit": 1,
                    "tokens": 8,
                },
                "prefix": {"local": {"queries": 10, "hits": 2}},
            },
            {
                "model_executed": True,
                "submit_mono_ns": 2_000_000,
                "complete_mono_ns": 3_000_000,
                "prefill_tokens": 0,
                "prefill_requests": 0,
                "chunked_prefill": {
                    "first": 0,
                    "middle": 0,
                    "final": 0,
                    "unsplit": 0,
                    "tokens": 0,
                },
                "prefix": {"local": {"queries": 0, "hits": 0}},
            },
        ]
    )
    assert mechanism["prefill"]["requests_per_step"] == 2.0
    assert mechanism["prefill"]["chunks"]["first"] == 1
    assert mechanism["prefix"]["hit_rate"] == 0.2
    assert all(mechanism["sanity"]["invariants"].values())

    manifest = {
        "repetitions": {str(index): {} for index in (1, 2, 3)},
        "regimes": {
            "A": {
                "source": "a_base",
                "actions": {"mns": "shared", "mbbt": "a_mbbt"},
            },
            "B": {
                "source": "b_base",
                "actions": {"mns": "b_mns", "mbbt": "shared"},
            },
        },
        "gates": {
            "minimum_relative_winner_margin": 0.10,
            "minimum_exclusive_fraction": 0.10,
            "minimum_exclusive_ratio": 5.0,
            "material_kv_usage": 0.90,
        },
    }
    runs = []
    for repetition in (1, 2, 3):
        runs.extend(
            [
                fake_run(
                    "a_base",
                    repetition,
                    goodput=1.0,
                    mns_score=0.8,
                    mbbt_score=0.01,
                ),
                fake_run(
                    "b_base",
                    repetition,
                    goodput=1.0,
                    mns_score=0.01,
                    mbbt_score=0.7,
                ),
                fake_run("shared", repetition, goodput=3.0),
                fake_run("a_mbbt", repetition, goodput=1.5),
                fake_run("b_mns", repetition, goodput=1.2),
            ]
        )
    result = analysis.evaluate_decisions(runs, manifest)
    assert result["decision"] == "STOP_NO_NEW_INSTRUMENTATION_NEEDED"
    assert result["baselines"] == {
        "always_mns_correct": 3,
        "always_mbbt_correct": 3,
        "binding_correct": 6,
        "decision_count": 6,
    }

    ambiguous = copy.deepcopy(runs)
    for run in ambiguous:
        if run["config_id"] == "b_base":
            run["binding"]["both_fraction"] = 0.8
    assert (
        analysis.evaluate_decisions(ambiguous, manifest)["decision"]
        == "OPEN_EXACT_ATTRIBUTION_ABLATION"
    )

    wrong = copy.deepcopy(runs)
    for run in wrong:
        if run["config_id"] == "b_base":
            run["binding"]["mns_exclusive_fraction"] = 0.8
            run["binding"]["mbbt_exclusive_fraction"] = 0.01
            for phase in run["phases"].values():
                phase["mns_exclusive_fraction"] = 0.8
                phase["mbbt_exclusive_fraction"] = 0.01
    assert (
        analysis.evaluate_decisions(wrong, manifest)["decision"]
        == "STOP_BINDING_NOT_PREDICTIVE"
    )

    prepare = load("action_aware_prepare", "prepare_pilot.py")
    frozen = prepare.build(
        ROOT / "runs/intervention-response-v2/pilot-manifest-v3.json"
    )
    assert frozen["status"] == "PASS"
    assert frozen["sanity"]["red_flags"] == []
    assert [config["id"] for config in frozen["configs"]] == [
        "b_base",
        "a_base",
        "shared",
        "b_mns",
        "a_mbbt",
    ]

    controller = load("action_aware_controller", "pilot_controller.py")
    args = SimpleNamespace(
        manifest=Path("/tmp/manifest.json"),
        run_root=Path("/tmp/action-aware"),
        aituner_root=Path("/tmp/aituner"),
        vllm_source=Path("/tmp/vllm"),
        venv=Path("/tmp/venv"),
        model=Path("/tmp/model"),
        client=Path("/tmp/client.py"),
    )
    controller.configure(args, frozen)
    plan = controller.dry_run_plan(args, frozen)
    assert plan["status"] == "PASS"
    assert len(plan["sessions"]) == 5
    assert plan["projected_h20_hours"] == 7.0
    assert "--max-num-batched-tokens 256" in plan["sessions"][0]["commands"]["server"]
    revised = prepare.build(
        ROOT / "runs/intervention-response-v2/pilot-manifest-v3.json",
        token_source_mbbt=2048,
        prior_attempt_h20_hours=0.38598689953486126,
        prior_attempt_artifact="/tmp/operational-stop-v0.json",
    )
    assert revised["schema"] == "action-aware-constraint-pilot-manifest-v1"
    assert revised["configs"][0]["mbbt"] == 2048
    assert revised["configs"][3]["mbbt"] == 2048
    assert revised["budget"]["hard_cap_h20_hours"] < 8.0
    controller.configure(args, revised)
    revised_plan = controller.dry_run_plan(args, revised)
    assert revised_plan["projected_h20_hours"] < revised_plan["hard_cap_h20_hours"]
    assert (
        "--max-num-batched-tokens 2048"
        in revised_plan["sessions"][0]["commands"]["server"]
    )
    accepted_burnin = {
        "kind": "anchor",
        "selection": {"count": 510},
        "interval": {"elapsed_s": 61.25},
        "pass_rate": 0.5,
        "feasible": False,
    }
    assert controller.burnin_gate(
        accepted_burnin, expected_count=510, maximum_elapsed_s=90.0
    )["elapsed_s"] == 61.25
    warmup = copy.deepcopy(accepted_burnin)
    warmup["kind"] = "warmup"
    try:
        controller.burnin_gate(warmup, expected_count=510, maximum_elapsed_s=90.0)
    except RuntimeError as error:
        assert "non-anchor" in str(error)
    else:
        raise AssertionError("warmup incorrectly passed the burnin gate")
    slow = copy.deepcopy(accepted_burnin)
    slow["interval"]["elapsed_s"] = 91.0
    try:
        controller.burnin_gate(slow, expected_count=510, maximum_elapsed_s=90.0)
    except RuntimeError as error:
        assert "throughput gate failed" in str(error)
    else:
        raise AssertionError("slow burnin incorrectly passed the throughput gate")
    print("action-aware constraint pilot: PASS")


if __name__ == "__main__":
    main()
