#!/usr/bin/env python3
"""Audit source-only constraint signals against crossed real interventions."""

from __future__ import annotations

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


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 = "action-aware-constraint-pilot-audit-v0"


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


def distribution(values: Iterable[float]) -> dict[str, Any]:
    finite = [float(value) for value in values]
    summary = numeric(finite)
    return {
        **summary,
        "mean": statistics.fmean(finite),
        "p50": quantile(finite, 0.50),
        "p95": quantile(finite, 0.95),
        "p99": quantile(finite, 0.99),
    }


def quantile(values: Iterable[float], probability: float) -> float:
    ordered = sorted(float(value) for value in values)
    if not ordered:
        raise ValueError("quantile requires values")
    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 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):
            try:
                records.append(json.loads(line))
            except json.JSONDecodeError as error:
                raise ValueError(f"{path}:{line_number}: invalid JSON") from error
    return records


def binding_summary(
    records: list[Mapping[str, Any]], *, mns: int, mbbt: int
) -> dict[str, Any]:
    if not records:
        raise ValueError("binding summary requires scheduler records")
    counts = {
        "mns_exclusive": 0,
        "mbbt_exclusive": 0,
        "both": 0,
        "waiting_unresolved": 0,
        "waiting": 0,
    }
    running_utilization = []
    token_utilization = []
    kv_usage = []
    preemptions = 0
    for record in records:
        waiting = int(record["queues"]["waiting"]) + int(
            record["queues"]["deferred"]
        )
        running = int(record["queues"]["running"])
        scheduled_tokens = int(record["prefill_tokens"]) + int(
            record["decode_tokens"]
        )
        if running > mns:
            raise ValueError("running requests exceed configured MNS")
        if scheduled_tokens > mbbt:
            raise ValueError("scheduled tokens exceed configured MBBT")
        mns_hit = waiting > 0 and running == mns
        mbbt_hit = waiting > 0 and scheduled_tokens == mbbt
        if waiting > 0:
            counts["waiting"] += 1
            if mns_hit and mbbt_hit:
                counts["both"] += 1
            elif mns_hit:
                counts["mns_exclusive"] += 1
            elif mbbt_hit:
                counts["mbbt_exclusive"] += 1
            else:
                counts["waiting_unresolved"] += 1
        running_utilization.append(running / mns)
        token_utilization.append(scheduled_tokens / mbbt)
        kv_usage.append(float(record["kv"]["usage"]))
        preemptions += int(record["preemptions"])
    count = len(records)
    return {
        "records": count,
        **{f"{name}_count": value for name, value in counts.items()},
        **{f"{name}_fraction": value / count for name, value in counts.items()},
        "running_utilization_mean": statistics.fmean(running_utilization),
        "running_utilization_max": max(running_utilization),
        "token_utilization_mean": statistics.fmean(token_utilization),
        "token_utilization_max": max(token_utilization),
        "kv_usage_mean": statistics.fmean(kv_usage),
        "kv_usage_max": max(kv_usage),
        "preemptions": preemptions,
    }


def telemetry_coverage(
    records: list[Mapping[str, Any]], *, start_ns: int, end_ns: int
) -> tuple[dict[str, float], bool]:
    if not records:
        raise ValueError("telemetry coverage requires records")
    submit_gaps = [
        (int(right["submit_mono_ns"]) - int(left["submit_mono_ns"])) / 1e9
        for left, right in zip(records, records[1:], strict=False)
    ]
    uncovered_gaps = [
        max(
            0,
            int(right["submit_mono_ns"]) - int(left["complete_mono_ns"]),
        )
        / 1e9
        for left, right in zip(records, records[1:], strict=False)
    ]
    coverage = {
        "start_gap_s": (int(records[0]["submit_mono_ns"]) - start_ns) / 1e9,
        "end_gap_s": (end_ns - int(records[-1]["submit_mono_ns"])) / 1e9,
        "max_internal_submit_gap_s": max(submit_gaps, default=0.0),
        "max_uncovered_gap_s": max(uncovered_gaps, default=0.0),
    }
    covered = (
        0.0 <= coverage["start_gap_s"] <= 1.0
        and 0.0 <= coverage["end_gap_s"] <= 1.0
        and 0.0 <= coverage["max_uncovered_gap_s"] <= 1.0
    )
    return coverage, covered


def mechanism_summary(records: list[Mapping[str, Any]]) -> dict[str, Any]:
    executed = [record for record in records if bool(record["model_executed"])]
    if not executed:
        raise ValueError("mechanism summary requires executed steps")
    prefill = [record for record in executed if int(record["prefill_tokens"]) > 0]
    decode_only = [
        record for record in executed if int(record["prefill_tokens"]) == 0
    ]
    if not prefill or not decode_only:
        raise ValueError("mechanism summary requires prefill and decode-only steps")

    def durations_ms(selected: list[Mapping[str, Any]]) -> list[float]:
        values = [
            (int(record["complete_mono_ns"]) - int(record["submit_mono_ns"]))
            / 1e6
            for record in selected
        ]
        if any(value < 0.0 for value in values):
            raise ValueError("engine step duration must be non-negative")
        return values

    chunk_keys = ("first", "middle", "final", "unsplit", "tokens")
    chunks = {
        key: sum(int(record["chunked_prefill"][key]) for record in executed)
        for key in chunk_keys
    }
    prefill_tokens = [int(record["prefill_tokens"]) for record in prefill]
    prefill_requests = sum(int(record["prefill_requests"]) for record in prefill)
    prefix_queries = sum(
        int(record["prefix"]["local"]["queries"]) for record in executed
    )
    prefix_hits = sum(
        int(record["prefix"]["local"]["hits"]) for record in executed
    )
    invariants = {
        "nonnegative_counts": all(
            value >= 0
            for value in (
                *chunks.values(),
                prefill_requests,
                prefix_queries,
                prefix_hits,
            )
        ),
        "chunk_tokens_match_prefill_tokens": chunks["tokens"]
        == sum(prefill_tokens),
        "prefix_hits_bounded": 0 <= prefix_hits <= prefix_queries,
    }
    return {
        "executed_steps": len(executed),
        "step_duration_ms": distribution(durations_ms(executed)),
        "prefill_steps": len(prefill),
        "prefill_step_duration_ms": distribution(durations_ms(prefill)),
        "decode_only_steps": len(decode_only),
        "decode_only_step_duration_ms": distribution(durations_ms(decode_only)),
        "prefill": {
            "requests": prefill_requests,
            "requests_per_step": prefill_requests / len(prefill),
            "tokens": sum(prefill_tokens),
            "tokens_per_step": distribution(prefill_tokens),
            "chunks": chunks,
        },
        "prefix": {
            "queries": prefix_queries,
            "hits": prefix_hits,
            "hit_rate": prefix_hits / prefix_queries if prefix_queries else 0.0,
        },
        "sanity": {"invariants": invariants},
    }


def request_summary(path: Path, expected_count: int) -> dict[str, Any]:
    rows = load_jsonl(path)
    if len(rows) != expected_count:
        raise ValueError(f"request row count mismatch: {path}")
    ttft = [float(row["ttft_ms"]) for row in rows if row["ttft_ms"] is not None]
    tpot = [float(row["tpot_ms"]) for row in rows if row["tpot_ms"] is not None]
    if not ttft or not tpot:
        raise ValueError(f"missing request latency values: {path}")
    return {
        "ttft_ms": {f"p{int(p * 100)}": quantile(ttft, p) for p in (0.5, 0.95, 0.99)},
        "tpot_ms": {f"p{int(p * 100)}": quantile(tpot, p) for p in (0.5, 0.95, 0.99)},
    }


def load_stream(session_root: Path) -> tuple[list[dict[str, Any]], dict[str, Any]]:
    streams = sorted((session_root / "opprof").glob("*.jsonl"))
    sidecars = sorted((session_root / "opprof").glob("*.jsonl.footer.json"))
    if len(streams) != 1 or len(sidecars) != 1:
        raise ValueError(f"expected one OpProf stream and sidecar: {session_root}")
    decoded = load_jsonl(streams[0])
    records = [row for row in decoded if "step_index" in row]
    footers = [row for row in decoded if row.get("record_type") == "footer"]
    sidecar = json.loads(sidecars[0].read_text(encoding="utf-8"))
    indexes = [int(row["step_index"]) for row in records]
    invariants = {
        "one_footer_last": len(footers) == 1 and decoded[-1] is footers[0],
        "sidecar_final": sidecar.get("final") is True,
        "zero_drops": sidecar.get("dropped_records") == 0,
        "written_matches_records": sidecar.get("written_records") == len(records),
        "contiguous_step_indexes": indexes == list(range(len(indexes))),
        "monotonic_timestamps": all(
            int(right["submit_mono_ns"]) >= int(left["submit_mono_ns"])
            for left, right in zip(records, records[1:], strict=False)
        ),
    }
    return records, {
        "stream": str(streams[0]),
        "stream_sha256": sha256_file(streams[0]),
        "records": len(records),
        "invariants": invariants,
    }


def analyze_run(
    *,
    run_root: Path,
    config: Mapping[str, Any],
    repetition: int,
    expected: Mapping[str, Any],
    stream_records: list[Mapping[str, Any]],
    duration_s: float,
    phase_fractions: list[float],
) -> dict[str, Any]:
    result_root = run_root / "sessions" / str(config["id"]) / f"rep{repetition}"
    result_path = result_root / "result.json"
    result = json.loads(result_path.read_text(encoding="utf-8"))
    selection = result["selection"]
    invariants = {
        "result_schema": result.get("schema") == "action-aware-pilot-result-v0",
        "config_id": result.get("config_id") == config["id"],
        "tp": int(result.get("tp", -1)) == 4,
        "mns": int(result.get("mns", -1)) == int(config["mns"]),
        "mbbt": int(result.get("mbbt", -1)) == int(config["mbbt"]),
        "uncensored": not bool(result.get("early_stopped", True)),
        "slo_early_stop_disabled": result.get("slo_early_stop_disabled") is True,
        "selection_count": int(selection["count"]) == int(expected["selected_count"]),
        "request_accounting": int(result["observed_count"])
        == int(expected["selected_count"]),
        "request_hash": selection["request_id_order_sha256"]
        == expected["request_id_order_sha256"],
        "arrival_hash": selection["arrival_order_sha256"]
        == expected["arrival_order_sha256"],
        "length_hash": selection["raw_length_order_sha256"]
        == expected["input_length_order_sha256"],
    }
    start_ns = int(result["interval"]["start_mono_ns"])
    arrival_end_ns = start_ns + round(duration_s * 1e9)
    full_records = [
        record
        for record in stream_records
        if start_ns <= int(record["submit_mono_ns"]) <= arrival_end_ns
    ]
    if not full_records:
        raise ValueError(f"no telemetry records in measured window: {result_path}")
    coverage, invariants["telemetry_coverage"] = telemetry_coverage(
        full_records, start_ns=start_ns, end_ns=arrival_end_ns
    )
    binding = binding_summary(
        full_records, mns=int(config["mns"]), mbbt=int(config["mbbt"])
    )
    mechanism = mechanism_summary(full_records)
    invariants["mechanism_summary"] = all(
        mechanism["sanity"]["invariants"].values()
    )
    phases = {}
    for fraction in phase_fractions:
        phase_end = start_ns + round(duration_s * fraction * 1e9)
        phase_records = [
            record
            for record in full_records
            if int(record["submit_mono_ns"]) <= phase_end
        ]
        phases[f"{fraction:.2f}"] = binding_summary(
            phase_records, mns=int(config["mns"]), mbbt=int(config["mbbt"])
        )
    state = summarize_engine(
        full_records,
        start_ns=start_ns,
        end_ns=arrival_end_ns,
        request_count=int(result["observed_count"]),
    )
    latency = request_summary(
        result_root / "requests.jsonl", int(result["observed_count"])
    )
    return {
        "config_id": config["id"],
        "mns": int(config["mns"]),
        "mbbt": int(config["mbbt"]),
        "repetition": repetition,
        "result_path": str(result_path),
        "result_sha256": sha256_file(result_path),
        "selection": {
            "count": int(selection["count"]),
            "request_id_order_sha256": selection["request_id_order_sha256"],
            "arrival_order_sha256": selection["arrival_order_sha256"],
            "raw_length_order_sha256": selection["raw_length_order_sha256"],
        },
        "outcome": {
            "pass_rate": float(result["pass_rate"]),
            "feasible": bool(result["feasible"]),
            "slo_pass_count": int(result["slo_pass_count"]),
            "slo_goodput_req_s": int(result["slo_pass_count"]) / duration_s,
            "elapsed_s": float(result["interval"]["elapsed_s"]),
            **latency,
        },
        "binding": binding,
        "mechanism": mechanism,
        "phases": phases,
        "state": state,
        "coverage": coverage,
        "invariants": invariants,
    }


def median(values: Iterable[float]) -> float:
    return float(statistics.median(float(value) for value in values))


def evaluate_decisions(
    runs: list[Mapping[str, Any]], manifest: Mapping[str, Any]
) -> dict[str, Any]:
    by_key = {
        (str(run["config_id"]), int(run["repetition"])): run for run in runs
    }
    repetitions = sorted(int(key) for key in manifest["repetitions"])
    regime_results = {}
    all_predictions = []
    crossed_pass = True
    binding_pass = True
    material_ambiguity = False
    for regime_name, regime in manifest["regimes"].items():
        rows = []
        source_runs = []
        for repetition in repetitions:
            source = by_key[(str(regime["source"]), repetition)]
            mns_target = by_key[(str(regime["actions"]["mns"]), repetition)]
            mbbt_target = by_key[(str(regime["actions"]["mbbt"]), repetition)]
            source_runs.append(source)
            source_goodput = float(source["outcome"]["slo_goodput_req_s"])
            mns_goodput = float(mns_target["outcome"]["slo_goodput_req_s"])
            mbbt_goodput = float(mbbt_target["outcome"]["slo_goodput_req_s"])
            observed = (
                "mns"
                if mns_goodput > mbbt_goodput
                else "mbbt"
                if mbbt_goodput > mns_goodput
                else "tie"
            )
            mns_score = float(source["binding"]["mns_exclusive_fraction"])
            mbbt_score = float(source["binding"]["mbbt_exclusive_fraction"])
            predicted = (
                "mns"
                if mns_score > mbbt_score
                else "mbbt"
                if mbbt_score > mns_score
                else "tie"
            )
            phase_predictions = {}
            for phase, summary in source["phases"].items():
                left = float(summary["mns_exclusive_fraction"])
                right = float(summary["mbbt_exclusive_fraction"])
                phase_predictions[phase] = (
                    "mns" if left > right else "mbbt" if right > left else "tie"
                )
            margin = (
                abs(mns_goodput - mbbt_goodput) / source_goodput
                if source_goodput > 0
                else None
            )
            row = {
                "repetition": repetition,
                "source_goodput_req_s": source_goodput,
                "mns_target_goodput_req_s": mns_goodput,
                "mbbt_target_goodput_req_s": mbbt_goodput,
                "observed_winner": observed,
                "predicted_winner": predicted,
                "prediction_correct": predicted == observed,
                "relative_winner_margin_over_source": margin,
                "mns_exclusive_fraction": mns_score,
                "mbbt_exclusive_fraction": mbbt_score,
                "phase_predictions": phase_predictions,
                "phase_stable": all(value == predicted for value in phase_predictions.values()),
            }
            rows.append(row)
            all_predictions.append(row)

        expected_winner = "mns" if regime_name == "A" else "mbbt"
        minimum_margin = float(manifest["gates"]["minimum_relative_winner_margin"])
        regime_crossed = all(
            row["observed_winner"] == expected_winner
            and row["relative_winner_margin_over_source"] is not None
            and row["relative_winner_margin_over_source"] >= minimum_margin
            for row in rows
        )
        crossed_pass &= regime_crossed
        winning_key = f"{expected_winner}_exclusive_fraction"
        losing_key = (
            "mbbt_exclusive_fraction" if expected_winner == "mns" else "mns_exclusive_fraction"
        )
        winning_median = median(row[winning_key] for row in rows)
        losing_median = median(row[losing_key] for row in rows)
        ratio_pass = winning_median >= float(
            manifest["gates"]["minimum_exclusive_ratio"]
        ) * losing_median
        regime_binding = (
            all(row["prediction_correct"] and row["phase_stable"] for row in rows)
            and winning_median
            >= float(manifest["gates"]["minimum_exclusive_fraction"])
            and ratio_pass
        )
        binding_pass &= regime_binding
        ambiguity_median = median(
            float(run["binding"]["both_fraction"])
            + float(run["binding"]["waiting_unresolved_fraction"])
            for run in source_runs
        )
        score_gap_median = median(
            abs(
                float(run["binding"]["mns_exclusive_fraction"])
                - float(run["binding"]["mbbt_exclusive_fraction"])
            )
            for run in source_runs
        )
        kv_max_median = median(
            float(run["binding"]["kv_usage_max"]) for run in source_runs
        )
        any_preemption = any(
            int(run["binding"]["preemptions"]) > 0 for run in source_runs
        )
        regime_material = (
            ambiguity_median >= score_gap_median
            or kv_max_median >= float(manifest["gates"]["material_kv_usage"])
            or any_preemption
        )
        material_ambiguity |= regime_material
        regime_results[regime_name] = {
            "source": regime["source"],
            "actions": regime["actions"],
            "expected_winner": expected_winner,
            "crossed_response_pass": regime_crossed,
            "binding_pass": regime_binding,
            "winning_exclusive_median": winning_median,
            "losing_exclusive_median": losing_median,
            "exclusive_ratio_pass": ratio_pass,
            "ambiguity_median": ambiguity_median,
            "exclusive_gap_median": score_gap_median,
            "kv_usage_max_median": kv_max_median,
            "any_preemption": any_preemption,
            "material_ambiguity": regime_material,
            "repetitions": rows,
        }

    if not crossed_pass:
        decision = "STOP_WORKLOAD_NOT_CROSSED"
    elif not binding_pass:
        decision = "STOP_BINDING_NOT_PREDICTIVE"
    elif material_ambiguity:
        decision = "OPEN_EXACT_ATTRIBUTION_ABLATION"
    else:
        decision = "STOP_NO_NEW_INSTRUMENTATION_NEEDED"
    correct = sum(int(row["prediction_correct"]) for row in all_predictions)
    return {
        "decision": decision,
        "crossed_response_pass": crossed_pass,
        "binding_pass": binding_pass,
        "material_ambiguity": material_ambiguity,
        "regimes": regime_results,
        "baselines": {
            "always_mns_correct": sum(
                int(row["observed_winner"] == "mns") for row in all_predictions
            ),
            "always_mbbt_correct": sum(
                int(row["observed_winner"] == "mbbt") for row in all_predictions
            ),
            "binding_correct": correct,
            "decision_count": len(all_predictions),
        },
    }


def analyze(run_root: Path, manifest_path: Path) -> dict[str, Any]:
    manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
    if manifest.get("schema") not in {
        "action-aware-constraint-pilot-manifest-v0",
        "action-aware-constraint-pilot-manifest-v1",
    }:
        raise ValueError("unexpected manifest schema")
    duration_s = float(manifest["engine"]["duration_s"])
    phase_fractions = [float(value) for value in manifest["gates"]["phase_fractions"]]
    runs = []
    stream_audits = []
    for config in manifest["configs"]:
        session_root = run_root / "sessions" / str(config["id"])
        stream_records, stream_audit = load_stream(session_root)
        stream_audit["config_id"] = config["id"]
        stream_audits.append(stream_audit)
        for repetition in sorted(int(key) for key in manifest["repetitions"]):
            runs.append(
                analyze_run(
                    run_root=run_root,
                    config=config,
                    repetition=repetition,
                    expected=manifest["repetitions"][str(repetition)]["selection"],
                    stream_records=stream_records,
                    duration_s=duration_s,
                    phase_fractions=phase_fractions,
                )
            )
    invariants = {
        "fifteen_runs": len(runs) == 15,
        "five_streams": len(stream_audits) == 5,
        "all_run_invariants": all(
            all(bool(value) for value in run["invariants"].values()) for run in runs
        ),
        "all_stream_invariants": all(
            all(bool(value) for value in stream["invariants"].values())
            for stream in stream_audits
        ),
        "nonnegative_counters": all(
            all(
                float(run["binding"][key]) >= 0
                for key in (
                    "mns_exclusive_count",
                    "mbbt_exclusive_count",
                    "both_count",
                    "waiting_unresolved_count",
                    "preemptions",
                )
            )
            for run in runs
        ),
        "ratios_bounded": all(
            all(
                0.0 <= float(run["binding"][key]) <= 1.0
                for key in (
                    "mns_exclusive_fraction",
                    "mbbt_exclusive_fraction",
                    "both_fraction",
                    "waiting_unresolved_fraction",
                    "kv_usage_mean",
                    "kv_usage_max",
                )
            )
            for run in runs
        ),
        "per_config_results_not_all_identical": len(
            {float(run["outcome"]["pass_rate"]) for run in runs}
        )
        > 1,
    }
    red_flags = [name for name, passed in invariants.items() if not passed]
    decisions = (
        evaluate_decisions(runs, manifest)
        if not red_flags
        else {
            "decision": "STOP_DATA_INVALID",
            "crossed_response_pass": False,
            "binding_pass": False,
            "material_ambiguity": False,
            "regimes": {},
            "baselines": {},
        }
    )
    payload = {
        "schema": SCHEMA,
        "decision": decisions["decision"],
        "manifest": str(manifest_path),
        "manifest_sha256": sha256_file(manifest_path),
        "run_root": str(run_root),
        "runs": runs,
        "streams": stream_audits,
        "decision_audit": decisions,
        "sanity": {
            "runs": len(runs),
            "pass_rate": numeric(run["outcome"]["pass_rate"] for run in runs),
            "slo_goodput_req_s": numeric(
                run["outcome"]["slo_goodput_req_s"] for run in runs
            ),
            "telemetry_records_per_run": numeric(
                run["binding"]["records"] for run in runs
            ),
            "mns_values": numeric(run["mns"] for run in runs),
            "mbbt_values": numeric(run["mbbt"] for run in runs),
            "invariants": invariants,
            "red_flags": red_flags,
        },
    }
    return payload


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--run-root", type=Path, required=True)
    parser.add_argument("--manifest", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    payload = analyze(args.run_root, args.manifest)
    atomic_json(args.output, payload)
    print(
        json.dumps(
            {
                "decision": payload["decision"],
                "sanity": payload["sanity"],
                "decision_audit": payload["decision_audit"],
            },
            indent=2,
            sort_keys=True,
        )
    )


if __name__ == "__main__":
    main()
