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

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


HERE = Path(__file__).resolve().parent


def load(name: str, path: Path):
    spec = importlib.util.spec_from_file_location(name, path)
    module = importlib.util.module_from_spec(spec)
    assert spec.loader is not None
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


def write_json(path: Path, payload) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(payload) + "\n", encoding="utf-8")


def write_jsonl(path: Path, rows) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(
        "".join(json.dumps(row) + "\n" for row in rows), encoding="utf-8"
    )


def engine_record(index: int, timestamp_ns: int) -> dict:
    alternate = index % 2
    return {
        "step_index": index,
        "submit_mono_ns": timestamp_ns,
        "model_executed": True,
        "scheduled_requests": 1 + alternate,
        "decode_batch_size": alternate,
        "prefill_tokens": 8 + alternate,
        "decode_tokens": alternate,
        "preemptions": 0,
        "queues": {"waiting": alternate, "running": 1 + alternate},
        "kv": {"usage": 0.1 + 0.01 * alternate},
        "cudagraph": {
            "runtime_mode": "FULL" if alternate else "NONE",
            "bucket_tokens": 16,
            "padding_tokens": alternate,
        },
        "dropped_records_before": 0,
    }


def main() -> None:
    extractor = load("active_intervention_extract_test", HERE / "extract_training.py")
    trainer = load("active_intervention_train_test", HERE / "train_policy.py")
    with tempfile.TemporaryDirectory() as temporary:
        root = Path(temporary)
        run_root = root / "runs"
        configs = [
            {"id": "a_base", "mns": 16, "mbbt": 8192},
            {"id": "a_mns", "mns": 64, "mbbt": 8192},
            {"id": "a_mbbt", "mns": 16, "mbbt": 16384},
            {"id": "b_base", "mns": 64, "mbbt": 2048},
            {"id": "b_mns", "mns": 128, "mbbt": 2048},
            {"id": "b_mbbt", "mns": 64, "mbbt": 8192},
        ]
        manifest = {
            "engine": {"duration_s": 300.0, "tp": 4},
            "configs": configs,
            "repetitions": {
                str(rep): {"selection": {"offered_req_s_per_gpu": 0.01}}
                for rep in (1, 2, 3)
            },
            "regimes": {
                "A": {
                    "source": "a_base",
                    "actions": {"mns": "a_mns", "mbbt": "a_mbbt"},
                },
                "B": {
                    "source": "b_base",
                    "actions": {"mns": "b_mns", "mbbt": "b_mbbt"},
                },
            },
        }
        manifest_path = root / "manifest.json"
        write_json(manifest_path, manifest)

        streams = []
        source_starts: dict[tuple[str, int], int] = {}
        for source_index, source_id in enumerate(("a_base", "b_base")):
            rows = []
            index = 0
            for repetition in (1, 2, 3):
                start_ns = int((source_index * 2000 + repetition * 400) * 1e9)
                source_starts[(source_id, repetition)] = start_ns
                for second in (1, 30, 76, 105, 151, 180, 226, 255):
                    rows.append(engine_record(index, start_ns + int(second * 1e9)))
                    index += 1
            stream_path = root / f"{source_id}-stream.jsonl"
            write_jsonl(stream_path, rows)
            streams.append(
                {
                    "config_id": source_id,
                    "stream": str(stream_path),
                    "stream_sha256": extractor.sha256_file(stream_path),
                }
            )

        request_rows = [
            {
                "request_id": f"r{index}",
                "arrival_s": arrival,
                "completed_elapsed_s": arrival + 10,
                "slo_pass": index != 3,
                "ttft_ms": 1000 + index * 100,
                "tpot_ms": 20 + index,
                "raw_input_tokens": 1000 + index * 100,
            }
            for index, arrival in enumerate((5.0, 80.0, 155.0, 230.0), 1)
        ]
        for source_id in ("a_base", "b_base"):
            for repetition in (1, 2, 3):
                write_jsonl(
                    run_root
                    / "sessions"
                    / source_id
                    / f"rep{repetition}"
                    / "requests.jsonl",
                    request_rows,
                )

        goodput = {
            "a_base": 0.020,
            "a_mns": 0.036,
            "a_mbbt": 0.028,
            "b_base": 0.032,
            "b_mns": 0.030,
            "b_mbbt": 0.038,
        }
        runs = []
        for config in configs:
            for repetition in (1, 2, 3):
                item = {
                    "config_id": config["id"],
                    "repetition": repetition,
                    "outcome": {
                        "slo_goodput_req_s": goodput[config["id"]]
                        + repetition * 0.0001
                    },
                }
                if config["id"] in ("a_base", "b_base"):
                    start_ns = source_starts[(config["id"], repetition)]
                    item["state"] = {
                        "interval": {
                            "start_ns": start_ns,
                            "end_ns": start_ns + int(300 * 1e9),
                        }
                    }
                runs.append(item)
        audit = {
            "schema": "action-aware-constraint-pilot-audit-v0",
            "sanity": {"red_flags": []},
            "streams": streams,
            "runs": runs,
        }
        audit_path = root / "audit.json"
        write_json(audit_path, audit)

        dataset = extractor.build_dataset(
            audit_path=audit_path, manifest_path=manifest_path, run_root=run_root
        )
        assert dataset["status"] == "VALID"
        assert len(dataset["examples"]) == 72
        assert not dataset["sanity"]["red_flags"]
        assert all(
            "exclusive" not in feature
            for example in dataset["examples"]
            for feature in example["source"]["telemetry"]
        )
        dataset_path = root / "dataset.json"
        write_json(dataset_path, dataset)
        policy = trainer.build_policy(dataset_path)
        assert policy["status"] in {
            "RETROSPECTIVE_INCREMENTAL_SIGNAL",
            "NO_RETROSPECTIVE_INCREMENTAL_SIGNAL",
        }
        assert not policy["sanity"]["red_flags"]
    print("active intervention pipeline: PASS")


if __name__ == "__main__":
    main()