#!/usr/bin/env python3
"""Choose measurement horizon and next intervention from a completed source run."""

from __future__ import annotations

import argparse
import hashlib
import importlib.util
import json
import math
import os
import statistics
import sys
from pathlib import Path
from typing import Any, Mapping, Sequence

import numpy as np


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


SCHEMA = "active-intervention-prospective-decision-v0"


def load_module(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


MODEL = load_module("active_intervention_prospective_model", HERE / "model.py")
EXTRACT = load_module(
    "active_intervention_prospective_extract", HERE / "extract_training.py"
)


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 numeric(values: Sequence[float]) -> dict[str, Any]:
    finite = [float(value) for value in values]
    if not finite or any(not math.isfinite(value) for value in finite):
        raise ValueError("numeric summary requires finite values")
    return {
        "n": len(finite),
        "min": min(finite),
        "max": max(finite),
        "distinct_n": len(set(finite)),
    }


def load_engine_records(source_root: Path) -> tuple[list[dict[str, Any]], Path]:
    streams = sorted((source_root / "opprof").glob("*.jsonl"))
    if len(streams) != 1:
        raise ValueError(f"expected one source engine stream, found {len(streams)}")
    records = [
        row for row in EXTRACT.load_jsonl(streams[0]) if "step_index" in row
    ]
    if not records:
        raise ValueError("source engine stream has no Layer-1 records")
    return records, streams[0]


def candidate_example(
    *,
    source_config: Mapping[str, Any],
    target_config: Mapping[str, Any],
    action_id: str,
    offered_rate_per_gpu: float,
    outcome: Mapping[str, float],
    telemetry: Mapping[str, float],
) -> dict[str, Any]:
    return {
        "source": {
            "mns": int(source_config["mns"]),
            "mbbt": int(source_config["mbbt"]),
            "offered_rate_per_gpu": float(offered_rate_per_gpu),
            "outcome": dict(outcome),
            "telemetry": dict(telemetry),
        },
        "action": {
            "id": action_id,
            "target_mns": int(target_config["mns"]),
            "target_mbbt": int(target_config["mbbt"]),
        },
    }


def aggregate_checkpoint(
    *,
    models: Sequence[Any],
    examples_by_action: Mapping[str, Sequence[Mapping[str, Any]]],
    include_telemetry: bool,
    confidence_z: float,
    minimum_margin: float,
) -> dict[str, Any]:
    rows = []
    for action_id, examples in sorted(examples_by_action.items()):
        raw = []
        for example in examples:
            source = example["source"]
            action = example["action"]
            noop = (
                int(source["mns"]) == int(action["target_mns"])
                and int(source["mbbt"]) == int(action["target_mbbt"])
            )
            if noop:
                raw.extend(0.0 for _model in models)
                continue
            names, values = MODEL.feature_vector(
                example, include_telemetry=include_telemetry
            )
            if any(model.feature_names != tuple(names) for model in models):
                raise ValueError("prospective feature schema does not match frozen model")
            raw.extend(model.predict(values) for model in models)
        clipped = np.clip(np.asarray(raw, dtype=np.float64), -1.0, 1.0)
        prediction = {
            "mean": float(clipped.mean()),
            "std": float(clipped.std(ddof=0)),
            "min": float(clipped.min()),
            "max": float(clipped.max()),
            "distinct_n": len(set(float(value) for value in clipped)),
            "sample_n": int(clipped.size),
        }
        rows.append(
            {
                "action_id": action_id,
                "prediction": prediction,
                "lower": prediction["mean"] - confidence_z * prediction["std"],
                "upper": prediction["mean"] + confidence_z * prediction["std"],
            }
        )
    rows.sort(key=lambda row: (-row["prediction"]["mean"], row["action_id"]))
    best, second = rows[:2]
    margin = float(best["prediction"]["mean"] - second["prediction"]["mean"])
    confident = bool(
        margin >= minimum_margin and best["lower"] > second["upper"]
    )
    return {
        "selected_action": best["action_id"],
        "confident": confident,
        "predicted_margin": margin,
        "candidates": rows,
    }


def apply_measurement_and_acquisition(checkpoints: list[dict[str, Any]]) -> dict[str, Any]:
    selected = checkpoints[-1]
    stop_reason = "full_measurement_fallback"
    for previous, current in zip(checkpoints, checkpoints[1:], strict=False):
        if (
            previous["confident"]
            and current["confident"]
            and previous["selected_action"] == current["selected_action"]
        ):
            selected = current
            stop_reason = "two_consecutive_confident_checkpoints"
            break
    candidates = selected["candidates"]
    mean_best = candidates[0]
    non_noop = [row for row in candidates if row["action_id"] != "noop"]
    if selected["confident"]:
        chosen = mean_best
        decision_kind = "exploit"
    else:
        positive_ucb = [row for row in non_noop if float(row["upper"]) > 0.0]
        if positive_ucb:
            chosen = max(
                positive_ucb,
                key=lambda row: (float(row["upper"]), row["action_id"]),
            )
            decision_kind = "diagnostic_ucb"
        else:
            chosen = next(row for row in candidates if row["action_id"] == "noop")
            decision_kind = "abstain_no_positive_ucb"
    remaining = [row for row in candidates if row["action_id"] != chosen["action_id"]]
    remaining.sort(key=lambda row: (-float(row["upper"]), row["action_id"]))
    order = [chosen["action_id"], *(row["action_id"] for row in remaining)]
    return {
        "selected_phase": selected["phase"],
        "selected_cutoff_s": selected["cutoff_s"],
        "measurement_stop_reason": stop_reason,
        "decision_kind": decision_kind,
        "selected_action": chosen["action_id"],
        "intervention_order": order,
        "selected_checkpoint": selected,
        "checkpoints": checkpoints,
    }


def build_decision(
    *, manifest_path: Path, policy_path: Path, run_root: Path
) -> dict[str, Any]:
    manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
    policy = json.loads(policy_path.read_text(encoding="utf-8"))
    if manifest.get("schema") != "active-intervention-prospective-manifest-v0":
        raise ValueError("unexpected prospective manifest schema")
    if policy.get("schema") != "active-intervention-policy-v0":
        raise ValueError("unexpected frozen policy schema")
    if sha256_file(policy_path) != manifest["policy"]["sha256"]:
        raise ValueError("frozen policy hash changed after manifest preparation")
    configs = {str(item["id"]): item for item in manifest["configs"]}
    source_id = str(manifest["source_config_id"])
    source_config = configs[source_id]
    source_root = run_root / "sessions" / source_id
    engine_records, stream_path = load_engine_records(source_root)
    phases = [f"{fraction:.2f}" for fraction in manifest["checkpoints"]["fractions"]]
    confidence_z = float(policy["measurement_policy"]["confidence_z"])
    minimum_margin = float(policy["measurement_policy"]["minimum_margin"])

    examples: dict[str, dict[str, dict[str, Mapping[str, Any]]]] = {}
    source_measurements: dict[str, dict[str, Any]] = {}
    source_normalized = []
    telemetry_values = []
    for repetition in sorted(int(key) for key in manifest["repetitions"]):
        item = manifest["repetitions"][str(repetition)]
        result_root = source_root / f"rep{repetition}"
        result = json.loads((result_root / "result.json").read_text(encoding="utf-8"))
        if result["selection"]["request_id_order_sha256"] != item["selection"][
            "request_id_order_sha256"
        ]:
            raise ValueError(f"source request hash mismatch: rep{repetition}")
        requests = EXTRACT.load_jsonl(result_root / "requests.jsonl")
        offered_rate = float(item["selection"]["offered_req_s_per_gpu"])
        offered_total = offered_rate * int(manifest["engine"]["tp"])
        source_normalized.append(
            float(result["slo_pass_count"])
            / float(manifest["engine"]["duration_s"])
            / offered_total
        )
        start_ns = int(result["interval"]["start_mono_ns"])
        examples[str(repetition)] = {}
        source_measurements[str(repetition)] = {
            "result": str(result_root / "result.json"),
            "result_sha256": sha256_file(result_root / "result.json"),
            "request_sha256": sha256_file(result_root / "requests.jsonl"),
            "phases": {},
        }
        for phase, cutoff_s in zip(
            phases, manifest["checkpoints"]["seconds"], strict=True
        ):
            outcome = EXTRACT.prefix_outcome(
                requests, cutoff_s=float(cutoff_s), offered_total=offered_total
            )
            admitted_count = sum(
                float(request["arrival_s"]) <= float(cutoff_s)
                for request in requests
            )
            state = summarize_engine(
                engine_records,
                start_ns=start_ns,
                end_ns=start_ns + round(float(cutoff_s) * 1e9),
                request_count=admitted_count,
            )
            if not all(state["sanity"]["invariants"].values()):
                raise ValueError(
                    f"source engine state invariant failed: rep{repetition} {phase}"
                )
            telemetry = EXTRACT.telemetry_record(state)
            telemetry_values.extend(float(value) for value in telemetry.values())
            source_measurements[str(repetition)]["phases"][phase] = {
                "cutoff_s": float(cutoff_s),
                "outcome": outcome,
                "telemetry": telemetry,
                "engine_sanity": state["sanity"],
            }
            examples[str(repetition)][phase] = {
                action_id: candidate_example(
                    source_config=source_config,
                    target_config=configs[str(target_id)],
                    action_id=action_id,
                    offered_rate_per_gpu=offered_rate,
                    outcome=outcome,
                    telemetry=telemetry,
                )
                for action_id, target_id in manifest["actions"].items()
            }

    decisions = {}
    for mode, include_telemetry in (("outcome_only", False), ("telemetry", True)):
        checkpoints = []
        for phase, cutoff_s in zip(
            phases, manifest["checkpoints"]["seconds"], strict=True
        ):
            models = MODEL.models_from_json(policy["phases"][phase][mode]["models"])
            examples_by_action = {
                action_id: [
                    examples[str(repetition)][phase][action_id]
                    for repetition in sorted(int(key) for key in manifest["repetitions"])
                ]
                for action_id in manifest["actions"]
            }
            checkpoint = aggregate_checkpoint(
                models=models,
                examples_by_action=examples_by_action,
                include_telemetry=include_telemetry,
                confidence_z=confidence_z,
                minimum_margin=minimum_margin,
            )
            checkpoints.append(
                {"phase": phase, "cutoff_s": float(cutoff_s), **checkpoint}
            )
        decisions[mode] = apply_measurement_and_acquisition(checkpoints)

    ceiling = float(manifest["gates"]["source_ceiling_normalized_goodput"])
    source_median = float(statistics.median(source_normalized))
    status = "STOP_SOURCE_CEILING" if source_median >= ceiling else "SELECTED"
    phase_admission_monotonic = all(
        all(
            left <= right + 1e-12
            for left, right in zip(values, values[1:], strict=False)
        )
        for repetition in source_measurements.values()
        for values in (
            [
                float(repetition["phases"][phase]["outcome"]["admitted_fraction"])
                for phase in phases
            ],
        )
    )
    telemetry_ratio_keys = {
        "prefill_token_fraction",
        "kv_usage_mean",
        "kv_usage_max",
        "graph_none_share",
        "graph_full_share",
        "graph_padding_fraction",
    }
    telemetry_records = [
        measurement["telemetry"]
        for repetition in source_measurements.values()
        for measurement in repetition["phases"].values()
    ]
    invariants = {
        "three_source_repetitions": len(source_normalized) == 3,
        "source_goodput_nonnegative": all(value >= 0.0 for value in source_normalized),
        "source_goodput_bounded": all(
            value <= 1.0 + 1e-12 for value in source_normalized
        ),
        "four_actions": set(manifest["actions"]) == {"noop", "mns", "mbbt", "joint"},
        "four_checkpoints": len(phases) == 4,
        "finite_telemetry": all(math.isfinite(value) for value in telemetry_values),
        "nonnegative_telemetry": all(
            float(value) >= 0.0
            for record in telemetry_records
            for key, value in record.items()
            if key != "kv_usage_end_minus_start"
        ),
        "telemetry_ratios_bounded": all(
            0.0 <= float(record[key]) <= 1.0 + 1e-12
            for record in telemetry_records
            for key in telemetry_ratio_keys
        ),
        "telemetry_not_all_identical": len(set(telemetry_values)) > 1,
        "phase_admission_monotonic": phase_admission_monotonic,
        "orders_are_permutations": all(
            set(decisions[mode]["intervention_order"]) == set(manifest["actions"])
            for mode in decisions
        ),
    }
    red_flags = [name for name, passed in invariants.items() if not passed]
    if red_flags:
        status = "STOP_SANITY"
    return {
        "schema": SCHEMA,
        "status": status,
        "manifest": str(manifest_path),
        "manifest_sha256": sha256_file(manifest_path),
        "policy": str(policy_path),
        "policy_sha256": sha256_file(policy_path),
        "source_stream": str(stream_path),
        "source_stream_sha256": sha256_file(stream_path),
        "source_measurements": source_measurements,
        "source_normalized_goodput": {
            "values": source_normalized,
            "median": source_median,
            **numeric(source_normalized),
        },
        "decisions": decisions,
        "sanity": {
            "invariants": invariants,
            "red_flags": red_flags,
            "telemetry_values": numeric(telemetry_values),
        },
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--manifest", type=Path, required=True)
    parser.add_argument("--policy", 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()
    decision = build_decision(
        manifest_path=args.manifest, policy_path=args.policy, run_root=args.run_root
    )
    atomic_json(args.output, decision)
    print(
        json.dumps(
            {
                "status": decision["status"],
                "source_normalized_goodput": decision["source_normalized_goodput"],
                "outcome_only": {
                    key: decision["decisions"]["outcome_only"][key]
                    for key in ("selected_cutoff_s", "decision_kind", "selected_action")
                },
                "telemetry": {
                    key: decision["decisions"]["telemetry"][key]
                    for key in ("selected_cutoff_s", "decision_kind", "selected_action")
                },
            },
            sort_keys=True,
        )
    )


if __name__ == "__main__":
    main()
