#!/usr/bin/env python3
"""Replay the P1 simulator shortlist under full and prefix policies."""

from __future__ import annotations

import argparse
import json
import math
import subprocess
from pathlib import Path
from typing import Any

from analyze_prefixes import numeric, sha256_file


AITUNER_ROOT = Path(__file__).resolve().parents[2]
FROZEN_K = 2
CUTOFF_S = 5.0
THRESHOLD = 0.95


def git_capture(*arguments: str) -> str:
    return subprocess.run(
        ["git", "-C", str(AITUNER_ROOT), *arguments],
        check=True,
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    ).stdout


def setup_costs(state: dict[str, Any]) -> dict[str, float]:
    result = {}
    for cell, payload in state["cells"].items():
        tp = int(payload["tp"])
        annotation_intervals = sum(
            float(run["elapsed_s"]) * tp / 3600.0
            for run in payload["runs"]
            if run["role"] not in {"low1", "high1"}
        )
        primary_intervals = sum(
            float(run["elapsed_s"]) * tp / 3600.0
            for run in payload["runs"]
            if run["role"] in {"low1", "high1"}
        )
        setup = float(payload["gpu_hours"]) - annotation_intervals - primary_intervals
        if setup < -1e-12:
            raise ValueError(f"negative inferred setup cost: {cell}={setup}")
        result[cell] = max(0.0, setup)
    return result


def build_candidates(
    manifest: dict[str, Any],
    state: dict[str, Any],
    strong: dict[str, Any],
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
    baseline_probability = strong["headline"]["sim_plus_outcome"]["probability"]
    instrument_probability = strong["headline"][
        "sim_plus_outcome_plus_instrumentation"
    ]["probability"]
    setup = setup_costs(state)
    anchors = []
    for detail, baseline_p, instrument_p in zip(
        strong["pilot_examples"], baseline_probability, instrument_probability
    ):
        cell = str(detail["cell"])
        level = str(detail["level"])
        role = f"{level}1"
        selection = manifest["cells"][cell]["targets"][level]["selections"][role]
        run = next(
            item for item in state["cells"][cell]["runs"] if item["role"] == role
        )
        tp = int(state["cells"][cell]["tp"])
        full_cost = float(run["elapsed_s"]) * tp / 3600.0
        prefix_cost = min(CUTOFF_S, float(run["elapsed_s"])) * tp / 3600.0
        anchors.append(
            {
                "cell": cell,
                "level": level,
                "role": role,
                "tp": tp,
                "real_feasible": bool(detail["adjudicated_feasible"]),
                "real_goodput_req_s_per_gpu": float(
                    selection["offered_req_s_per_gpu"]
                ),
                "sim_feasible": bool(detail["sim_slo_feasible"]),
                "sim_pass_rate": float(detail["sim_slo_pass_rate"]),
                "sim_throughput_req_s_per_gpu": float(
                    detail["sim_completed_throughput_per_gpu"]
                ),
                "baseline_probability": float(baseline_p),
                "instrument_probability": float(instrument_p),
                "setup_h20_hours": setup[cell],
                "full_trial_h20_hours": full_cost,
                "prefix_h20_hours": prefix_cost,
            }
        )

    candidates = []
    for cell in sorted(manifest["cells"]):
        feasible = [
            anchor for anchor in anchors if anchor["cell"] == cell and anchor["sim_feasible"]
        ]
        if not feasible:
            continue
        candidates.append(
            max(feasible, key=lambda anchor: anchor["sim_throughput_req_s_per_gpu"])
        )
    candidates.sort(
        key=lambda anchor: (
            -anchor["sim_throughput_req_s_per_gpu"],
            anchor["cell"],
        )
    )
    return anchors, candidates


def expanded_top_k(candidates: list[dict[str, Any]], k: int) -> list[dict[str, Any]]:
    if not candidates or k <= 0:
        return []
    boundary = candidates[min(k, len(candidates)) - 1][
        "sim_throughput_req_s_per_gpu"
    ]
    return [
        candidate
        for candidate in candidates
        if candidate["sim_throughput_req_s_per_gpu"] >= boundary - 1e-12
    ]


def selected_result(
    evaluated: list[dict[str, Any]], feasible_key: str
) -> tuple[str | None, float | None]:
    feasible = [candidate for candidate in evaluated if candidate[feasible_key]]
    if not feasible:
        return None, None
    best = max(feasible, key=lambda candidate: candidate["real_goodput_req_s_per_gpu"])
    return str(best["cell"]), float(best["real_goodput_req_s_per_gpu"])


def replay(
    shortlist: list[dict[str, Any]],
    *,
    probability_key: str | None,
    oracle_goodput: float,
    common_failure_h20_hours: float,
) -> dict[str, Any]:
    evaluated = []
    online_cost = 0.0
    early_accept = 0
    early_reject = 0
    false_accept = 0
    false_reject = 0
    for candidate in shortlist:
        current = dict(candidate)
        online_cost += current["setup_h20_hours"]
        if probability_key is None:
            predicted_feasible = current["real_feasible"]
            online_cost += current["full_trial_h20_hours"]
            action = "full"
        else:
            probability = float(current[probability_key])
            if probability >= THRESHOLD:
                predicted_feasible = True
                early_accept += 1
                online_cost += current["prefix_h20_hours"]
                action = "early_accept"
                false_accept += int(not current["real_feasible"])
            elif probability <= 1.0 - THRESHOLD:
                predicted_feasible = False
                early_reject += 1
                online_cost += current["prefix_h20_hours"]
                action = "early_reject"
                false_reject += int(current["real_feasible"])
            else:
                predicted_feasible = current["real_feasible"]
                online_cost += current["full_trial_h20_hours"]
                action = "continue_full"
        current["policy_feasible"] = predicted_feasible
        current["action"] = action
        evaluated.append(current)
    selected_cell, selected_goodput = selected_result(evaluated, "policy_feasible")
    regret = (
        1.0 - selected_goodput / oracle_goodput
        if selected_goodput is not None and oracle_goodput > 0
        else None
    )
    return {
        "selected_cell": selected_cell,
        "selected_real_goodput_req_s_per_gpu": selected_goodput,
        "real_regret": regret,
        "online_h20_hours": online_cost,
        "conservative_h20_hours_with_prior_failure": (
            online_cost + common_failure_h20_hours
        ),
        "early_accept": early_accept,
        "early_reject": early_reject,
        "false_accept": false_accept,
        "false_reject": false_reject,
        "evaluated": [
            {
                "cell": item["cell"],
                "level": item["level"],
                "action": item["action"],
                "real_feasible": item["real_feasible"],
                "policy_feasible": item["policy_feasible"],
            }
            for item in evaluated
        ],
    }


def analyze(
    manifest_path: Path,
    state_path: Path,
    prior_state_path: Path,
    strong_path: Path,
) -> dict[str, Any]:
    manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
    state = json.loads(state_path.read_text(encoding="utf-8"))
    prior = json.loads(prior_state_path.read_text(encoding="utf-8"))
    strong = json.loads(strong_path.read_text(encoding="utf-8"))
    anchors, candidates = build_candidates(manifest, state, strong)
    oracle_anchor = max(
        (anchor for anchor in anchors if anchor["real_feasible"]),
        key=lambda anchor: anchor["real_goodput_req_s_per_gpu"],
    )
    oracle_goodput = float(oracle_anchor["real_goodput_req_s_per_gpu"])
    common_failure = float(prior["gpu_hours_total"])
    by_k = {}
    for k in (1, 2, 3, 6):
        shortlist = expanded_top_k(candidates, k)
        full = replay(
            shortlist,
            probability_key=None,
            oracle_goodput=oracle_goodput,
            common_failure_h20_hours=common_failure,
        )
        baseline = replay(
            shortlist,
            probability_key="baseline_probability",
            oracle_goodput=oracle_goodput,
            common_failure_h20_hours=common_failure,
        )
        instrument = replay(
            shortlist,
            probability_key="instrument_probability",
            oracle_goodput=oracle_goodput,
            common_failure_h20_hours=common_failure,
        )
        for result in (baseline, instrument):
            result["online_cost_reduction_vs_full"] = (
                1.0 - result["online_h20_hours"] / full["online_h20_hours"]
            )
            result["conservative_cost_reduction_vs_full"] = 1.0 - (
                result["conservative_h20_hours_with_prior_failure"]
                / full["conservative_h20_hours_with_prior_failure"]
            )
        by_k[str(k)] = {
            "actual_shortlist_size": len(shortlist),
            "shortlist": [candidate["cell"] for candidate in shortlist],
            "sim_top_k_plus_real_final": full,
            "sim_plus_outcome": baseline,
            "sim_plus_outcome_plus_instrumentation": instrument,
        }

    frozen = by_k[str(FROZEN_K)]
    full = frozen["sim_top_k_plus_real_final"]
    baseline = frozen["sim_plus_outcome"]
    instrument = frozen["sim_plus_outcome_plus_instrumentation"]
    baseline_safe = baseline["false_accept"] == 0 and baseline["false_reject"] == 0
    instrument_safe = (
        instrument["false_accept"] == 0 and instrument["false_reject"] == 0
    )
    incremental_reduction = (
        1.0 - instrument["online_h20_hours"] / baseline["online_h20_hours"]
        if baseline_safe and instrument_safe and baseline["online_h20_hours"] > 0
        else None
    )
    contribution_gate = {
        "frozen_k": FROZEN_K,
        "instrument_safe": instrument_safe,
        "outcome_baseline_safe": baseline_safe,
        "instrument_regret_at_most_5pct": (
            instrument["real_regret"] is not None
            and instrument["real_regret"] <= 0.05
        ),
        "instrument_cost_reduction_vs_full_at_least_30pct": (
            instrument["online_cost_reduction_vs_full"] >= 0.30
        ),
        "instrument_cost_reduction_vs_outcome_at_least_20pct": (
            incremental_reduction is not None and incremental_reduction >= 0.20
        ),
        "incremental_reduction_vs_outcome": incremental_reduction,
    }
    contribution_gate["passes"] = all(
        contribution_gate[key]
        for key in (
            "instrument_safe",
            "outcome_baseline_safe",
            "instrument_regret_at_most_5pct",
            "instrument_cost_reduction_vs_full_at_least_30pct",
            "instrument_cost_reduction_vs_outcome_at_least_20pct",
        )
    )

    red_flags = []
    if state["status"] != "complete" or int(state["completed_cells"]) != 6:
        red_flags.append("pilot_incomplete")
    if strong["status"] != "PASS" or strong["sanity"]["red_flags"]:
        red_flags.append("strong_input_invalid")
    if len(anchors) != 12 or len(candidates) != 6:
        red_flags.append("unexpected_surface_size")
    probabilities = [
        value
        for anchor in anchors
        for value in (anchor["baseline_probability"], anchor["instrument_probability"])
    ]
    costs = [
        value
        for anchor in anchors
        for value in (
            anchor["setup_h20_hours"],
            anchor["full_trial_h20_hours"],
            anchor["prefix_h20_hours"],
        )
    ]
    if not all(0.0 <= value <= 1.0 for value in probabilities):
        red_flags.append("probability_out_of_range")
    if not all(value >= 0.0 and math.isfinite(value) for value in costs):
        red_flags.append("invalid_cost")

    return {
        "schema": "fidelity-pilot-e2e-v1",
        "status": "PASS" if not red_flags else "STOP",
        "scope": "held-out P1 replay; gate diagnostic, not paper-facing evidence",
        "ranking": [
            {
                "rank": rank,
                "cell": candidate["cell"],
                "level": candidate["level"],
                "sim_throughput_req_s_per_gpu": candidate[
                    "sim_throughput_req_s_per_gpu"
                ],
                "real_feasible": candidate["real_feasible"],
                "real_goodput_req_s_per_gpu": candidate[
                    "real_goodput_req_s_per_gpu"
                ],
            }
            for rank, candidate in enumerate(candidates, start=1)
        ],
        "real_oracle": {
            "cell": oracle_anchor["cell"],
            "level": oracle_anchor["level"],
            "goodput_req_s_per_gpu": oracle_goodput,
        },
        "by_k": by_k,
        "contribution_gate": contribution_gate,
        "analysis": {
            "script": str(Path(__file__).resolve()),
            "script_sha256": sha256_file(Path(__file__).resolve()),
            "aituner_git_head": git_capture("rev-parse", "HEAD").strip(),
            "aituner_git_status_short": git_capture("status", "--short"),
        },
        "provenance": {
            "manifest": str(manifest_path.resolve()),
            "manifest_sha256": sha256_file(manifest_path),
            "controller_state": str(state_path.resolve()),
            "controller_state_sha256": sha256_file(state_path),
            "prior_state": str(prior_state_path.resolve()),
            "prior_state_sha256": sha256_file(prior_state_path),
            "strong_metrics": str(strong_path.resolve()),
            "strong_metrics_sha256": sha256_file(strong_path),
        },
        "sanity": {
            "red_flags": red_flags,
            "anchors": numeric([1 for _ in anchors]),
            "candidates": numeric([1 for _ in candidates]),
            "probabilities": numeric(probabilities),
            "costs_h20_hours": numeric(costs),
            "invariants": {
                "anchors_12": len(anchors) == 12,
                "candidates_6": len(candidates) == 6,
                "probabilities_bounded": all(
                    0.0 <= value <= 1.0 for value in probabilities
                ),
                "costs_nonnegative": all(value >= 0.0 for value in costs),
                "per_config_not_all_identical": len(
                    {candidate["sim_throughput_req_s_per_gpu"] for candidate in candidates}
                )
                > 1,
                "tie_expansion_applied": True,
            },
        },
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--manifest", type=Path, required=True)
    parser.add_argument("--controller-state", type=Path, required=True)
    parser.add_argument("--prior-state", type=Path, required=True)
    parser.add_argument("--strong-metrics", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    result = analyze(
        args.manifest,
        args.controller_state,
        args.prior_state,
        args.strong_metrics,
    )
    args.output.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n")
    print(
        json.dumps(
            {
                "status": result["status"],
                "red_flags": result["sanity"]["red_flags"],
                "contribution_gate": result["contribution_gate"],
            },
            sort_keys=True,
        )
    )
    if result["status"] != "PASS":
        raise RuntimeError(result["sanity"]["red_flags"])


if __name__ == "__main__":
    main()