#!/usr/bin/env python3
"""Audit whether a controlled knob change produces identifiable telemetry deltas.

This is a development-only feasibility audit.  It compares adjacent MNS
interventions at an identical TP, offered-load anchor, and request sequence
against same-config primary/confirmation repeat noise.  It does not claim that
the observed response is causal or that it improves an end-to-end tuner.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import sys
from collections import defaultdict
from pathlib import Path
from statistics import median
from typing import Any, Iterable, Mapping


HERE = Path(__file__).resolve().parent
COMMON_STATE_DIR = HERE.parent / "telemetry-residual"
sys.path.insert(0, str(COMMON_STATE_DIR))

from common_state import load_jsonl, summarize_engine  # noqa: E402


SCHEMA = "intervention-response-audit-v0"
HORIZONS_S = (5.0, 10.0)
GATE_FEATURES = (
    "scheduler_steps_per_s",
    "decode_batch_size.mean",
    "prefill_token_fraction",
    "queue_waiting_mean",
    "queue_running_mean",
    "kv_usage_mean",
    "graph_padding_fraction",
)
ALL_FEATURES = (
    "scheduler_steps_per_s",
    "batch_size.mean",
    "batch_tokens.mean",
    "decode_batch_size.mean",
    "prefill_token_fraction",
    "queue_waiting_mean",
    "queue_running_mean",
    "preemptions",
    "kv_usage_mean",
    "kv_usage_max",
    "kv_usage_end_minus_start",
    "graph_none_share",
    "graph_full_share",
    "graph_padding_fraction",
)
EXPECTED_ACTION_PAIRS = 17
MIN_REPEAT_PAIRS = 20
MIN_STABLE_FEATURES = 2
MIN_SIGN_CONSISTENCY = 0.75
MIN_EFFECT_TO_NOISE = 2.0
MIN_ABOVE_NOISE_P95_FRACTION = 0.5


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


def quantile(values: Iterable[float], probability: float) -> float:
    ordered = sorted(float(value) for value in values)
    if not ordered:
        raise ValueError("quantile requires at least one value")
    if not 0.0 <= probability <= 1.0:
        raise ValueError("quantile probability must be in [0, 1]")
    position = probability * (len(ordered) - 1)
    lower = math.floor(position)
    upper = math.ceil(position)
    if lower == upper:
        return ordered[lower]
    weight = position - lower
    return ordered[lower] * (1.0 - weight) + ordered[upper] * weight


def flatten_state(summary: Mapping[str, Any]) -> dict[str, float]:
    common = summary["common"]
    engine = summary["engine_only"]
    state = {
        "scheduler_steps_per_s": float(common["scheduler_steps_per_s"]),
        "batch_size.mean": float(common["batch_size"]["mean"]),
        "batch_tokens.mean": float(common["batch_tokens"]["mean"]),
        "decode_batch_size.mean": float(common["decode_batch_size"]["mean"]),
        "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": float(common["preemptions"]),
        "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"]),
    }
    if set(state) != set(ALL_FEATURES):
        raise ValueError("flattened state does not match the frozen feature set")
    if any(not math.isfinite(value) for value in state.values()):
        raise ValueError("flattened state contains a non-finite value")
    return state


def _trial_role(path: Path) -> str:
    return "confirmation" if path.parent.name.startswith("confirm-") else "primary"


def load_trials(
    raw_root: Path,
    *,
    horizons_s: tuple[float, ...] = HORIZONS_S,
) -> tuple[dict[float, list[dict[str, Any]]], list[dict[str, Any]]]:
    by_horizon = {horizon: [] for horizon in horizons_s}
    stream_provenance = []
    for cell_dir in sorted(path for path in raw_root.iterdir() if path.is_dir()):
        streams = sorted((cell_dir / "opprof").glob("*.jsonl"))
        if len(streams) != 1:
            raise ValueError(f"{cell_dir}: expected exactly one Layer-1 stream")
        stream = streams[0]
        records = load_jsonl(stream)
        stream_provenance.append(
            {
                "path": str(stream),
                "sha256": sha256_file(stream),
                "bytes": stream.stat().st_size,
            }
        )
        result_paths = sorted(cell_dir.glob("anchor-*/result.json"))
        result_paths.extend(sorted(cell_dir.glob("confirm-*-anchor-*/result.json")))
        for result_path in result_paths:
            result = json.loads(result_path.read_text(encoding="utf-8"))
            start_ns = int(result["interval"]["start_mono_ns"])
            elapsed_s = float(result["interval"]["elapsed_s"])
            for horizon_s in horizons_s:
                if elapsed_s < horizon_s:
                    raise ValueError(
                        f"{result_path}: elapsed {elapsed_s} is shorter than {horizon_s}s"
                    )
                state = flatten_state(
                    summarize_engine(
                        records,
                        start_ns=start_ns,
                        end_ns=start_ns + int(horizon_s * 1e9),
                        request_count=int(result["selection"]["count"]),
                    )
                )
                by_horizon[horizon_s].append(
                    {
                        "trial_id": str(result_path.relative_to(raw_root)),
                        "result_sha256": sha256_file(result_path),
                        "role": _trial_role(result_path),
                        "cell": str(result["cell"]),
                        "study_sha256": str(result["study_sha256"]),
                        "tp": int(result["tp"]),
                        "mns": int(result["mns"]),
                        "anchor": float(result["anchor"]),
                        "request_hash": str(
                            result["selection"]["request_id_order_sha256"]
                        ),
                        "request_count": int(result["selection"]["count"]),
                        "early_stopped": bool(result["early_stopped"]),
                        "full_pass_rate": float(result["pass_rate"]),
                        "full_feasible": bool(result["feasible"]),
                        "state": state,
                    }
                )
    return by_horizon, stream_provenance


def _group_key(trial: Mapping[str, Any]) -> tuple[Any, ...]:
    return (
        trial["study_sha256"],
        trial["tp"],
        trial["anchor"],
        trial["request_hash"],
    )


def _delta(source: Mapping[str, Any], target: Mapping[str, Any]) -> dict[str, float]:
    return {
        feature: float(target["state"][feature]) - float(source["state"][feature])
        for feature in ALL_FEATURES
    }


def _pair(source: Mapping[str, Any], target: Mapping[str, Any], kind: str) -> dict[str, Any]:
    if _group_key(source) != _group_key(target):
        raise ValueError("pair endpoints do not share workload identity")
    return {
        "kind": kind,
        "group": {
            "study_sha256": source["study_sha256"],
            "tp": source["tp"],
            "anchor": source["anchor"],
            "request_hash": source["request_hash"],
        },
        "source": {
            "trial_id": source["trial_id"],
            "cell": source["cell"],
            "mns": source["mns"],
            "early_stopped": source["early_stopped"],
            "full_pass_rate": source["full_pass_rate"],
            "full_feasible": source["full_feasible"],
        },
        "target": {
            "trial_id": target["trial_id"],
            "cell": target["cell"],
            "mns": target["mns"],
            "early_stopped": target["early_stopped"],
            "full_pass_rate": target["full_pass_rate"],
            "full_feasible": target["full_feasible"],
        },
        "delta_state": _delta(source, target),
        "descriptive_full_outcome": {
            "delta_pass_rate": target["full_pass_rate"] - source["full_pass_rate"],
            "feasibility_transition": (
                f"{str(source['full_feasible']).lower()}->"
                f"{str(target['full_feasible']).lower()}"
            ),
        },
    }


def build_pairs(
    trials: list[dict[str, Any]],
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
    primary_groups: dict[tuple[Any, ...], list[dict[str, Any]]] = defaultdict(list)
    primary_by_cell_anchor: dict[tuple[Any, ...], dict[str, Any]] = {}
    confirmations = []
    for trial in trials:
        if trial["role"] == "primary":
            primary_groups[_group_key(trial)].append(trial)
            primary_by_cell_anchor[
                (trial["cell"], trial["anchor"], trial["request_hash"])
            ] = trial
        else:
            confirmations.append(trial)

    actions = []
    for group in primary_groups.values():
        ordered = sorted(group, key=lambda item: item["mns"])
        for source, target in zip(ordered, ordered[1:], strict=False):
            if target["mns"] == source["mns"] * 2:
                actions.append(_pair(source, target, "mns_increase"))

    repeats = []
    for confirmation in confirmations:
        key = (
            confirmation["cell"],
            confirmation["anchor"],
            confirmation["request_hash"],
        )
        primary = primary_by_cell_anchor.get(key)
        if primary is None:
            raise ValueError(f"{confirmation['trial_id']}: missing matched primary")
        if primary["mns"] != confirmation["mns"]:
            raise ValueError("repeat endpoints changed MNS")
        repeats.append(_pair(primary, confirmation, "same_config_repeat"))
    return actions, repeats


def response_statistics(
    actions: list[dict[str, Any]],
    repeats: list[dict[str, Any]],
) -> dict[str, Any]:
    statistics = {}
    for feature in ALL_FEATURES:
        action = [float(pair["delta_state"][feature]) for pair in actions]
        noise = [float(pair["delta_state"][feature]) for pair in repeats]
        action_abs = [abs(value) for value in action]
        noise_abs = [abs(value) for value in noise]
        positive = sum(value > 1e-12 for value in action)
        negative = sum(value < -1e-12 for value in action)
        zero = len(action) - positive - negative
        nonzero = positive + negative
        sign_consistency = max(positive, negative) / nonzero if nonzero else 0.0
        action_median = median(action_abs)
        noise_median = median(noise_abs)
        noise_p95 = quantile(noise_abs, 0.95)
        effect_to_noise = (
            action_median / noise_median
            if noise_median > 0
            else (math.inf if action_median > 0 else 0.0)
        )
        above_noise = sum(value > noise_p95 for value in action_abs) / len(action_abs)
        qualifies = (
            feature in GATE_FEATURES
            and sign_consistency >= MIN_SIGN_CONSISTENCY
            and effect_to_noise >= MIN_EFFECT_TO_NOISE
            and above_noise >= MIN_ABOVE_NOISE_P95_FRACTION
        )
        statistics[feature] = {
            "action_delta": numeric(action),
            "repeat_delta": numeric(noise),
            "action_abs_median": action_median,
            "repeat_abs_median": noise_median,
            "repeat_abs_p95": noise_p95,
            "effect_to_repeat_median": (
                effect_to_noise if math.isfinite(effect_to_noise) else None
            ),
            "effect_to_repeat_median_is_infinite": math.isinf(effect_to_noise),
            "action_signs": {
                "positive": positive,
                "negative": negative,
                "zero": zero,
                "consistency": sign_consistency,
            },
            "action_above_repeat_p95_fraction": above_noise,
            "gate_feature": feature in GATE_FEATURES,
            "qualifies": qualifies,
        }
    return statistics


def analyze_horizon(trials: list[dict[str, Any]], horizon_s: float) -> dict[str, Any]:
    actions, repeats = build_pairs(trials)
    feature_statistics = response_statistics(actions, repeats)
    qualifying = sorted(
        feature for feature, item in feature_statistics.items() if item["qualifies"]
    )
    all_values = [
        value
        for trial in trials
        for value in trial["state"].values()
    ]
    action_vectors = {
        tuple(round(float(pair["delta_state"][feature]), 12) for feature in ALL_FEATURES)
        for pair in actions
    }
    pair_invariants = {
        "expected_action_pair_count": len(actions) == EXPECTED_ACTION_PAIRS,
        "sufficient_repeat_pair_count": len(repeats) >= MIN_REPEAT_PAIRS,
        "all_pair_hashes_match": all(
            pair["group"]["request_hash"] for pair in [*actions, *repeats]
        ),
        "all_values_finite": all(math.isfinite(value) for value in all_values),
        "state_vectors_not_all_identical": len(action_vectors) > 1,
        "ratios_bounded": all(
            0.0 <= trial["state"][feature] <= 1.0
            for trial in trials
            for feature in (
                "prefill_token_fraction",
                "kv_usage_mean",
                "kv_usage_max",
                "graph_none_share",
                "graph_full_share",
                "graph_padding_fraction",
            )
        ),
        "nonnegative_counters": all(
            trial["state"][feature] >= 0.0
            for trial in trials
            for feature in (
                "scheduler_steps_per_s",
                "batch_size.mean",
                "batch_tokens.mean",
                "decode_batch_size.mean",
                "queue_waiting_mean",
                "queue_running_mean",
                "preemptions",
            )
        ),
    }
    red_flags = [name for name, passed in pair_invariants.items() if not passed]
    pass_deltas = [
        pair["descriptive_full_outcome"]["delta_pass_rate"] for pair in actions
    ]
    transitions = defaultdict(int)
    for pair in actions:
        transitions[pair["descriptive_full_outcome"]["feasibility_transition"]] += 1
    return {
        "horizon_s": horizon_s,
        "actions": actions,
        "repeats": repeats,
        "feature_statistics": feature_statistics,
        "qualifying_features": qualifying,
        "descriptive_full_outcome": {
            "delta_pass_rate": numeric(pass_deltas),
            "positive": sum(value > 1e-12 for value in pass_deltas),
            "negative": sum(value < -1e-12 for value in pass_deltas),
            "zero": sum(abs(value) <= 1e-12 for value in pass_deltas),
            "feasibility_transitions": dict(sorted(transitions.items())),
            "limitation": (
                "Full outcomes may use different elapsed durations when a trial "
                "early-stopped; they are descriptive and are not a gate input."
            ),
        },
        "sanity": {
            "trials": len(trials),
            "action_pairs": len(actions),
            "repeat_pairs": len(repeats),
            "distinct_action_vectors": len(action_vectors),
            "invariants": pair_invariants,
            "red_flags": red_flags,
        },
    }


def audit(
    *,
    metrics_path: Path,
    raw_root: Path,
    output_path: Path,
) -> dict[str, Any]:
    trials_by_horizon, streams = load_trials(raw_root)
    horizons = {
        str(int(horizon)): analyze_horizon(trials, horizon)
        for horizon, trials in sorted(trials_by_horizon.items())
    }
    red_flags = sorted(
        {
            red_flag
            for horizon in horizons.values()
            for red_flag in horizon["sanity"]["red_flags"]
        }
    )
    stable_features = sorted(
        set.intersection(
            *(set(horizon["qualifying_features"]) for horizon in horizons.values())
        )
    )
    if red_flags:
        decision = "STOP_DATA_INVALID"
    elif len(stable_features) < MIN_STABLE_FEATURES:
        decision = "STOP_NO_IDENTIFIABLE_RESPONSE"
    else:
        decision = "OPEN_MATCHED_PILOT"
    payload = {
        "schema": SCHEMA,
        "status": "COMPLETE",
        "decision": decision,
        "claim_boundary": (
            "Development-only identifiability gate. Passing opens a controlled "
            "real-GPU pilot; it does not establish tuning benefit or causality."
        ),
        "frozen_gate": {
            "horizons_s": list(HORIZONS_S),
            "expected_action_pairs": EXPECTED_ACTION_PAIRS,
            "minimum_repeat_pairs": MIN_REPEAT_PAIRS,
            "minimum_stable_features": MIN_STABLE_FEATURES,
            "minimum_sign_consistency": MIN_SIGN_CONSISTENCY,
            "minimum_effect_to_repeat_median": MIN_EFFECT_TO_NOISE,
            "minimum_action_above_repeat_p95_fraction": (
                MIN_ABOVE_NOISE_P95_FRACTION
            ),
            "gate_features": list(GATE_FEATURES),
        },
        "stable_qualifying_features": stable_features,
        "horizons": horizons,
        "provenance": {
            "analysis_script": str(Path(__file__).resolve()),
            "analysis_script_sha256": sha256_file(Path(__file__).resolve()),
            "phase6_metrics": str(metrics_path.resolve()),
            "phase6_metrics_sha256": sha256_file(metrics_path),
            "raw_root": str(raw_root.resolve()),
            "streams": streams,
        },
        "sanity": {
            "stream_count": len(streams),
            "stream_bytes": numeric(item["bytes"] for item in streams),
            "red_flags": red_flags,
        },
    }
    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
    return payload


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--metrics", type=Path, required=True)
    parser.add_argument("--raw-root", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    payload = audit(
        metrics_path=args.metrics,
        raw_root=args.raw_root,
        output_path=args.output,
    )
    print(
        json.dumps(
            {
                "decision": payload["decision"],
                "stable_qualifying_features": payload[
                    "stable_qualifying_features"
                ],
                "sanity": payload["sanity"],
            },
            indent=2,
            sort_keys=True,
        )
    )


if __name__ == "__main__":
    main()
