#!/usr/bin/env python3
"""Audit held-out action/measurement choices against the exact 2x2 surface."""

from __future__ import annotations

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


SCHEMA = "active-intervention-prospective-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: list[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_surface(
    manifest: Mapping[str, Any], run_root: Path
) -> tuple[dict[str, Any], list[dict[str, Any]]]:
    rows = []
    aggregate = {}
    duration_s = float(manifest["engine"]["duration_s"])
    tp = int(manifest["engine"]["tp"])
    for config in manifest["configs"]:
        config_id = str(config["id"])
        values = []
        for repetition in sorted(int(key) for key in manifest["repetitions"]):
            expected = manifest["repetitions"][str(repetition)]["selection"]
            result_path = (
                run_root / "sessions" / config_id / f"rep{repetition}" / "result.json"
            )
            result = json.loads(result_path.read_text(encoding="utf-8"))
            if result["selection"]["request_id_order_sha256"] != expected[
                "request_id_order_sha256"
            ]:
                raise ValueError(f"request hash mismatch: {config_id} rep{repetition}")
            offered_total = float(expected["offered_req_s_per_gpu"]) * tp
            normalized = float(result["slo_pass_count"]) / duration_s / offered_total
            values.append(normalized)
            rows.append(
                {
                    "config_id": config_id,
                    "mns": int(config["mns"]),
                    "mbbt": int(config["mbbt"]),
                    "repetition": repetition,
                    "normalized_slo_goodput": normalized,
                    "slo_goodput_req_s": float(result["slo_pass_count"]) / duration_s,
                    "pass_rate": float(result["pass_rate"]),
                    "elapsed_s": float(result["interval"]["elapsed_s"]),
                    "result": str(result_path),
                    "result_sha256": sha256_file(result_path),
                }
            )
        aggregate[config_id] = {
            "normalized_slo_goodput_values": values,
            "median_normalized_slo_goodput": float(statistics.median(values)),
            "sanity": numeric(values),
        }
    return aggregate, rows


def source_cost_estimate(
    *,
    source_session: Mapping[str, Any],
    source_rows: list[Mapping[str, Any]],
    cutoff_s: float,
    tp: int,
) -> dict[str, float]:
    actual_h20_hours = float(source_session["gpu_hours"])
    measured_replay_h20_hours = (
        tp * sum(float(row["elapsed_s"]) for row in source_rows) / 3600.0
    )
    fixed_h20_hours = max(0.0, actual_h20_hours - measured_replay_h20_hours)
    prefix_replay_h20_hours = tp * len(source_rows) * cutoff_s / 3600.0
    return {
        "actual_full_session_h20_hours": actual_h20_hours,
        "fixed_startup_warmup_burnin_cleanup_h20_hours": fixed_h20_hours,
        "prefix_replay_h20_hours_lower_bound": prefix_replay_h20_hours,
        "counterfactual_all_in_h20_hours_lower_bound": fixed_h20_hours
        + prefix_replay_h20_hours,
    }


def replay_policy(
    *,
    mode: str,
    manifest: Mapping[str, Any],
    decision: Mapping[str, Any],
    surface: Mapping[str, Any],
    session_costs: Mapping[str, float],
    source_cost: Mapping[str, float],
) -> dict[str, Any]:
    acceptable_regret = float(manifest["gates"]["acceptable_regret"])
    source_id = str(manifest["source_config_id"])
    oracle = max(
        float(item["median_normalized_slo_goodput"]) for item in surface.values()
    )
    cumulative = float(source_cost["counterfactual_all_in_h20_hours_lower_bound"])
    source_score = float(surface[source_id]["median_normalized_slo_goodput"])
    source_regret = 1.0 - source_score / oracle if oracle > 0 else 0.0
    points = [
        {
            "action_id": "noop",
            "config_id": source_id,
            "score": source_score,
            "regret": source_regret,
            "cumulative_h20_hours_lower_bound": cumulative,
        }
    ]
    hit = points[0] if source_regret <= acceptable_regret + 1e-12 else None
    seen = {source_id}
    for action_id in decision["decisions"][mode]["intervention_order"]:
        config_id = str(manifest["actions"][action_id])
        if config_id in seen:
            continue
        seen.add(config_id)
        cumulative += float(session_costs[config_id])
        score = float(surface[config_id]["median_normalized_slo_goodput"])
        regret = 1.0 - score / oracle if oracle > 0 else 0.0
        point = {
            "action_id": action_id,
            "config_id": config_id,
            "score": score,
            "regret": regret,
            "cumulative_h20_hours_lower_bound": cumulative,
        }
        points.append(point)
        if hit is None and regret <= acceptable_regret + 1e-12:
            hit = point
    return {
        "mode": mode,
        "measurement_cutoff_s": float(
            decision["decisions"][mode]["selected_cutoff_s"]
        ),
        "selected_action": decision["decisions"][mode]["selected_action"],
        "decision_kind": decision["decisions"][mode]["decision_kind"],
        "intervention_order": decision["decisions"][mode]["intervention_order"],
        "source_cost": dict(source_cost),
        "oracle_normalized_slo_goodput": oracle,
        "cost_to_acceptable": hit,
        "reached_acceptable": hit is not None,
        "points": points,
    }


def build_audit(
    *, manifest_path: Path, decision_path: Path, run_root: Path
) -> dict[str, Any]:
    manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
    decision = json.loads(decision_path.read_text(encoding="utf-8"))
    state_path = run_root / "controller-state.json"
    state = json.loads(state_path.read_text(encoding="utf-8"))
    if manifest.get("schema") != "active-intervention-prospective-manifest-v0":
        raise ValueError("unexpected prospective manifest schema")
    if decision.get("schema") != "active-intervention-prospective-decision-v0":
        raise ValueError("unexpected prospective decision schema")
    if decision["manifest_sha256"] != sha256_file(manifest_path):
        raise ValueError("decision does not match prospective manifest")
    surface, rows = load_surface(manifest, run_root)
    source_id = str(manifest["source_config_id"])
    sessions = state["sessions"]
    session_costs = {
        config_id: float(sessions[config_id]["gpu_hours"])
        for config_id in surface
    }
    source_rows = [row for row in rows if row["config_id"] == source_id]
    policies = {}
    for mode in ("outcome_only", "telemetry"):
        cost = source_cost_estimate(
            source_session=sessions[source_id],
            source_rows=source_rows,
            cutoff_s=float(decision["decisions"][mode]["selected_cutoff_s"]),
            tp=int(manifest["engine"]["tp"]),
        )
        policies[mode] = replay_policy(
            mode=mode,
            manifest=manifest,
            decision=decision,
            surface=surface,
            session_costs=session_costs,
            source_cost=cost,
        )
    outcome_hit = policies["outcome_only"]["cost_to_acceptable"]
    telemetry_hit = policies["telemetry"]["cost_to_acceptable"]
    if outcome_hit is None or telemetry_hit is None:
        reduction = None
    else:
        outcome_cost = float(outcome_hit["cumulative_h20_hours_lower_bound"])
        telemetry_cost = float(telemetry_hit["cumulative_h20_hours_lower_bound"])
        reduction = 1.0 - telemetry_cost / outcome_cost if outcome_cost > 0 else 0.0
    confirmation_trigger = bool(
        reduction is not None
        and reduction
        >= float(manifest["gates"]["confirmation_trigger_gpu_cost_reduction"])
        and policies["telemetry"]["reached_acceptable"]
    )
    contribution_gate = bool(
        reduction is not None
        and reduction >= float(manifest["gates"]["contribution_gpu_cost_reduction"])
        and policies["telemetry"]["reached_acceptable"]
    )
    status = (
        "TRIGGER_ACTUAL_EARLY_STOP_CONFIRMATION"
        if confirmation_trigger
        else "STOP_NO_PROSPECTIVE_GPU_COST_SIGNAL"
    )
    normalized_values = [float(row["normalized_slo_goodput"]) for row in rows]
    costs = list(session_costs.values())
    invariants = {
        "controller_complete": state.get("status") == "complete",
        "four_sessions_complete": len(sessions) == 4
        and all(item.get("status") == "complete" for item in sessions.values()),
        "twelve_surface_outcomes": len(rows) == 12,
        "nonnegative_goodput": all(value >= 0.0 for value in normalized_values),
        "normalized_goodput_bounded": all(value <= 1.0 + 1e-12 for value in normalized_values),
        "surface_not_all_identical": len(set(normalized_values)) > 1,
        "nonnegative_session_costs": all(value >= 0.0 for value in costs),
        "policy_replay_reaches_oracle_surface": all(
            policy["reached_acceptable"] for policy in policies.values()
        ),
    }
    red_flags = [name for name, passed in invariants.items() if not passed]
    if red_flags:
        status = "STOP_SANITY"
    return {
        "schema": SCHEMA,
        "status": status,
        "claim_boundary": (
            "Prospective exact-surface replay. Prefix source costs reconstruct the "
            "measured fixed overhead plus selected replay seconds; actual early-stop "
            "confirmation is required before claiming GPU-cost reduction."
        ),
        "manifest": str(manifest_path),
        "manifest_sha256": sha256_file(manifest_path),
        "decision": str(decision_path),
        "decision_sha256": sha256_file(decision_path),
        "controller_state": str(state_path),
        "controller_state_sha256": sha256_file(state_path),
        "surface": surface,
        "rows": rows,
        "session_costs_h20_hours": session_costs,
        "annotation_campaign_h20_hours": float(state["gpu_hours_total"]),
        "policies": policies,
        "comparison": {
            "telemetry_gpu_cost_reduction_fraction": reduction,
            "confirmation_trigger": confirmation_trigger,
            "contribution_gate": contribution_gate,
            "confirmation_trigger_threshold": manifest["gates"][
                "confirmation_trigger_gpu_cost_reduction"
            ],
            "contribution_threshold": manifest["gates"][
                "contribution_gpu_cost_reduction"
            ],
            "action_changed": policies["outcome_only"]["selected_action"]
            != policies["telemetry"]["selected_action"],
            "measurement_changed": policies["outcome_only"]["measurement_cutoff_s"]
            != policies["telemetry"]["measurement_cutoff_s"],
        },
        "sanity": {
            "invariants": invariants,
            "red_flags": red_flags,
            "normalized_slo_goodput": numeric(normalized_values),
            "session_h20_hours": numeric(costs),
        },
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--manifest", type=Path, required=True)
    parser.add_argument("--decision", 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()
    audit = build_audit(
        manifest_path=args.manifest,
        decision_path=args.decision,
        run_root=args.run_root,
    )
    atomic_json(args.output, audit)
    print(
        json.dumps(
            {
                "status": audit["status"],
                "comparison": audit["comparison"],
                "sanity": audit["sanity"],
            },
            sort_keys=True,
        )
    )


if __name__ == "__main__":
    main()