#!/usr/bin/env python3
"""Extract paired source/action examples from the accepted action-aware run."""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import os
import sys
from pathlib import Path
from statistics import fmean
from typing import Any, Mapping


PHASES = ("0.25", "0.50", "0.75", "1.00")
HERE = Path(__file__).resolve().parent
COMMON_STATE = HERE.parent / "telemetry-residual"
sys.path.insert(0, str(COMMON_STATE))

from common_state import summarize_engine  # noqa: E402


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as source:
        for chunk in iter(lambda: source.read(1 << 20), b""):
            digest.update(chunk)
    return digest.hexdigest()


def atomic_json(path: Path, payload: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + ".tmp")
    temporary.write_text(
        json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )
    os.replace(temporary, path)


def load_jsonl(path: Path) -> list[dict[str, Any]]:
    records = []
    with path.open(encoding="utf-8") as source:
        for line_number, line in enumerate(source, 1):
            if not line.strip():
                continue
            try:
                records.append(json.loads(line))
            except json.JSONDecodeError as error:
                raise ValueError(f"{path}:{line_number}: invalid JSON") from error
    if not records:
        raise ValueError(f"{path}: no request records")
    return records


def prefix_outcome(
    requests: list[Mapping[str, Any]], *, cutoff_s: float, offered_total: float
) -> dict[str, float]:
    admitted = [request for request in requests if float(request["arrival_s"]) <= cutoff_s]
    completed = [
        request
        for request in requests
        if request.get("completed_elapsed_s") is not None
        and float(request["completed_elapsed_s"]) <= cutoff_s
    ]
    if not admitted:
        raise ValueError("prefix has no admitted requests")
    admitted_ids = {str(request["request_id"]) for request in admitted}
    if any(str(request["request_id"]) not in admitted_ids for request in completed):
        raise ValueError("prefix completion precedes admission")
    passed = sum(bool(request["slo_pass"]) for request in completed)
    ttft = [float(request["ttft_ms"]) for request in completed]
    tpot = [float(request["tpot_ms"]) for request in completed]
    total = len(requests)
    return {
        "normalized_slo_goodput": passed / cutoff_s / offered_total,
        "admitted_fraction": len(admitted) / total,
        "completed_over_admitted": len(completed) / len(admitted),
        "completed_pass_rate": passed / max(1, len(completed)),
        "completed_fail_fraction_of_total": (len(completed) - passed) / total,
        "outstanding_over_admitted": (len(admitted) - len(completed)) / len(admitted),
        "ttft_max_over_slo_max": max(ttft, default=0.0) / 6000.0,
        "ttft_mean_over_slo_max": fmean(ttft) / 6000.0 if ttft else 0.0,
        "tpot_max_over_slo": max(tpot, default=0.0) / 50.0,
        "tpot_mean_over_slo": fmean(tpot) / 50.0 if tpot else 0.0,
        "admitted_input_tokens_mean_over_limit": fmean(
            float(request["raw_input_tokens"]) for request in admitted
        )
        / 8192.0,
    }


def telemetry_record(state: Mapping[str, Any]) -> dict[str, float]:
    common = state["common"]
    engine = state["engine_only"]
    executed_steps = int(state["sanity"]["executed_steps"])
    if executed_steps <= 0:
        raise ValueError("telemetry phase contains no executed engine steps")
    return {
        "scheduler_steps_per_s": float(common["scheduler_steps_per_s"]),
        "batch_size_mean": float(common["batch_size"]["mean"]),
        "batch_size_cv": float(common["batch_size"]["cv"]),
        "batch_tokens_mean": float(common["batch_tokens"]["mean"]),
        "batch_tokens_cv": float(common["batch_tokens"]["cv"]),
        "decode_batch_size_mean": float(common["decode_batch_size"]["mean"]),
        "decode_batch_size_cv": float(common["decode_batch_size"]["cv"]),
        "prefill_token_fraction": float(common["prefill_token_fraction"]),
        "queue_waiting_mean": float(common["queue_waiting_mean"]),
        "queue_running_mean": float(common["queue_running_mean"]),
        "preemptions_per_step": float(common["preemptions"]) / executed_steps,
        "kv_usage_mean": float(engine["kv_usage_mean"]),
        "kv_usage_max": float(engine["kv_usage_max"]),
        "kv_usage_end_minus_start": float(engine["kv_usage_end_minus_start"]),
        "graph_none_share": float(engine["graph_none_share"]),
        "graph_full_share": float(engine["graph_full_share"]),
        "graph_padding_fraction": float(engine["graph_padding_fraction"]),
    }


def load_stream(path: Path, *, expected_sha256: str) -> list[dict[str, Any]]:
    if sha256_file(path) != expected_sha256:
        raise ValueError(f"engine stream hash mismatch: {path}")
    decoded = load_jsonl(path)
    records = [row for row in decoded if "step_index" in row]
    if not records:
        raise ValueError(f"engine stream has no Layer-1 records: {path}")
    return records


def build_dataset(
    *, audit_path: Path, manifest_path: Path, run_root: Path
) -> dict[str, Any]:
    audit = json.loads(audit_path.read_text(encoding="utf-8"))
    manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
    if audit.get("schema") != "action-aware-constraint-pilot-audit-v0":
        raise ValueError("unexpected action-aware audit schema")
    if audit["sanity"]["red_flags"]:
        raise ValueError(f"action-aware audit red flags: {audit['sanity']['red_flags']}")
    configs = {str(item["id"]): item for item in manifest["configs"]}
    runs = {
        (str(run["config_id"]), int(run["repetition"])): run
        for run in audit["runs"]
    }
    source_ids = {str(regime["source"]) for regime in manifest["regimes"].values()}
    stream_entries = {
        str(item["config_id"]): item
        for item in audit["streams"]
        if str(item["config_id"]) in source_ids
    }
    if set(stream_entries) != source_ids:
        raise ValueError("audit is missing a source config engine stream")
    streams = {
        config_id: load_stream(
            Path(item["stream"]), expected_sha256=str(item["stream_sha256"])
        )
        for config_id, item in stream_entries.items()
    }
    examples = []
    request_hashes = []
    for regime_name, regime in sorted(manifest["regimes"].items()):
        source_id = str(regime["source"])
        for repetition in sorted(int(value) for value in manifest["repetitions"]):
            source_run = runs[(source_id, repetition)]
            source_config = configs[source_id]
            request_path = run_root / "sessions" / source_id / f"rep{repetition}" / "requests.jsonl"
            requests = load_jsonl(request_path)
            request_hashes.append(sha256_file(request_path))
            offered_rate_per_gpu = float(
                manifest["repetitions"][str(repetition)]["selection"][
                    "offered_req_s_per_gpu"
                ]
            )
            offered_total = offered_rate_per_gpu * int(manifest["engine"]["tp"])
            source_goodput = float(source_run["outcome"]["slo_goodput_req_s"])
            source_normalized = min(1.0, source_goodput / offered_total)
            decision_id = f"{regime_name}-rep{repetition}"
            for phase in PHASES:
                cutoff_s = float(manifest["engine"]["duration_s"]) * float(phase)
                outcome = prefix_outcome(
                    requests, cutoff_s=cutoff_s, offered_total=offered_total
                )
                admitted_count = sum(
                    float(request["arrival_s"]) <= cutoff_s for request in requests
                )
                start_ns = int(source_run["state"]["interval"]["start_ns"])
                phase_state = summarize_engine(
                    streams[source_id],
                    start_ns=start_ns,
                    end_ns=start_ns + round(cutoff_s * 1e9),
                    request_count=admitted_count,
                )
                if not all(phase_state["sanity"]["invariants"].values()):
                    raise ValueError(
                        f"engine state invariant failed: {decision_id} phase {phase}"
                    )
                telemetry = telemetry_record(phase_state)
                actions = {"noop": source_id, **regime["actions"]}
                for action_name, target_id in sorted(actions.items()):
                    target_run = runs[(str(target_id), repetition)]
                    target_config = configs[str(target_id)]
                    target_goodput = float(target_run["outcome"]["slo_goodput_req_s"])
                    normalized = target_goodput / offered_total
                    if not 0.0 <= normalized <= 1.0 + 1e-12:
                        raise ValueError("target normalized goodput is outside [0, 1]")
                    examples.append(
                        {
                            "phase": phase,
                            "cutoff_s": cutoff_s,
                            "decision_id": decision_id,
                            "regime": regime_name,
                            "repetition": repetition,
                            "source": {
                                "config_id": source_id,
                                "mns": int(source_config["mns"]),
                                "mbbt": int(source_config["mbbt"]),
                                "offered_rate_per_gpu": offered_rate_per_gpu,
                                "outcome": outcome,
                                "telemetry": telemetry,
                            },
                            "action": {
                                "id": action_name,
                                "target_config_id": str(target_id),
                                "target_mns": int(target_config["mns"]),
                                "target_mbbt": int(target_config["mbbt"]),
                            },
                            "target_slo_goodput_req_s": target_goodput,
                            "target_normalized_goodput": min(1.0, normalized),
                            "source_normalized_goodput": source_normalized,
                            "target_delta_normalized_goodput": min(1.0, normalized)
                            - source_normalized,
                        }
                    )
    invariants = {
        "expected_examples": len(examples) == len(PHASES) * 2 * 3 * 3,
        "four_phases": sorted({example["phase"] for example in examples})
        == sorted(PHASES),
        "six_decisions": len({example["decision_id"] for example in examples}) == 6,
        "three_actions_per_decision_phase": all(
            sum(
                item["decision_id"] == decision
                and item["phase"] == phase
                for item in examples
            )
            == 3
            for decision in {item["decision_id"] for item in examples}
            for phase in PHASES
        ),
        "targets_not_all_identical": len(
            {example["target_normalized_goodput"] for example in examples}
        )
        > 1,
        "bounded_prefix_ratios": all(
            0.0 <= float(value) <= 1.0
            for example in examples
            for key, value in example["source"]["outcome"].items()
            if key
            in {
                "admitted_fraction",
                "completed_over_admitted",
                "completed_pass_rate",
                "completed_fail_fraction_of_total",
                "outstanding_over_admitted",
            }
        ),
        "direct_telemetry_without_binding_labels": all(
            not any(token in key for token in ("exclusive", "unresolved", "both"))
            for example in examples
            for key in example["source"]["telemetry"]
        ),
        "treatment_effects_bounded": all(
            -1.0 <= float(example["target_delta_normalized_goodput"]) <= 1.0
            for example in examples
        ),
    }
    red_flags = [name for name, passed in invariants.items() if not passed]
    if red_flags:
        raise RuntimeError(f"training dataset sanity failed: {red_flags}")
    return {
        "schema": "active-intervention-training-v0",
        "status": "VALID",
        "provenance": {
            "audit": str(audit_path),
            "audit_sha256": sha256_file(audit_path),
            "manifest": str(manifest_path),
            "manifest_sha256": sha256_file(manifest_path),
            "run_root": str(run_root),
            "source_request_sha256": sorted(set(request_hashes)),
            "source_stream_sha256": sorted(
                str(item["stream_sha256"]) for item in stream_entries.values()
            ),
        },
        "examples": examples,
        "sanity": {"invariants": invariants, "red_flags": red_flags},
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--audit", type=Path, required=True)
    parser.add_argument("--manifest", type=Path, required=True)
    parser.add_argument("--run-root", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    dataset = build_dataset(
        audit_path=args.audit,
        manifest_path=args.manifest,
        run_root=args.run_root,
    )
    atomic_json(args.output, dataset)
    print(
        json.dumps(
            {
                "status": dataset["status"],
                "examples": len(dataset["examples"]),
                "sanity": dataset["sanity"],
            },
            sort_keys=True,
        )
    )


if __name__ == "__main__":
    main()