#!/usr/bin/env python
from __future__ import annotations

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

PROJECT_ROOT = Path(__file__).resolve().parents[1]


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description=(
            "Build a measured CTT outcome acceptance gate from rollout and "
            "train-clean selector artifacts."
        )
    )
    parser.add_argument(
        "--rollout-metrics",
        type=Path,
        default=Path("runs/ctt_base_context_obs_test_envclip_k16_rollout_comparison/metrics.json"),
    )
    parser.add_argument(
        "--selector-metrics",
        type=Path,
        default=Path(
            "runs/ctt_base_context_obs_learned_dominance_chartcompat_obs_utility_task_envclip_k16_train_to_test/metrics.json"
        ),
    )
    parser.add_argument(
        "--out-dir",
        type=Path,
        default=Path("runs/ctt_outcome_acceptance_gate"),
    )
    parser.add_argument("--selected-min", type=float, default=0.4745)
    parser.add_argument("--target-selected-min", type=float, default=0.50)
    parser.add_argument("--proposal-oracle-min", type=float, default=0.50)
    parser.add_argument("--success-support-gap-max", type=float, default=0.07)
    parser.add_argument("--success-selector-gap-max", type=float, default=0.03)
    parser.add_argument("--v0-outcome-ptr-floor", type=float, default=0.4435)
    parser.add_argument("--unsafe-slack", type=float, default=0.0)
    parser.add_argument("--min-train-seeds", type=int, default=3)
    args = parser.parse_args(argv)

    rollout = _load_json(args.rollout_metrics)
    selector = _load_json(args.selector_metrics)
    values = _observed_values(rollout, selector)
    gates = _gates(values, args)
    out_dir = args.out_dir
    out_dir.mkdir(parents=True, exist_ok=True)

    payload = {
        "report_type": "ctt_outcome_acceptance_gate",
        "schema_version": 1,
        "overall_pass": all(bool(gate["pass"]) for gate in gates),
        "status": "method_success" if all(bool(gate["pass"]) for gate in gates) else "diagnostic_not_method_success",
        "rollout_metrics": str(args.rollout_metrics),
        "selector_metrics": str(args.selector_metrics),
        "thresholds": {
            "selected_min": args.selected_min,
            "target_selected_min": args.target_selected_min,
            "proposal_oracle_min": args.proposal_oracle_min,
            "success_support_gap_max": args.success_support_gap_max,
            "success_selector_gap_max": args.success_selector_gap_max,
            "v0_outcome_ptr_floor": args.v0_outcome_ptr_floor,
            "unsafe_slack": args.unsafe_slack,
            "min_train_seeds": args.min_train_seeds,
        },
        "observed": values,
        "gates": gates,
        "data_hash": rollout.get("data_hash") or rollout.get("target_content_hash"),
        "split_hash": rollout.get("split_hash") or rollout.get("target_split_hash"),
        "selector_data_hash": selector.get("data_hash") or selector.get("eval_target_content_hash"),
        "selector_split_hash": selector.get("split_hash") or selector.get("eval_target_split_hash"),
    }
    (out_dir / "metrics.json").write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
    (out_dir / "metrics_by_task.json").write_text(
        json.dumps(
            _group_file_payload(args.selector_metrics, args.rollout_metrics, "metrics_by_task.json"),
            indent=2,
            sort_keys=True,
        )
        + "\n"
    )
    (out_dir / "metrics_by_seed.json").write_text(
        json.dumps(
            _group_file_payload(args.selector_metrics, args.rollout_metrics, "metrics_by_seed.json"),
            indent=2,
            sort_keys=True,
        )
        + "\n"
    )
    (out_dir / "table.tex").write_text(_table(gates) + "\n")
    (out_dir / "config.yaml").write_text(_config(args) + "\n")
    (out_dir / "command.txt").write_text(
        "python scripts/build_ctt_outcome_gate.py " + " ".join(sys.argv[1:]) + "\n"
    )
    (out_dir / "git_hash.txt").write_text(_run(["git", "rev-parse", "HEAD"]) + "\n")
    (out_dir / "data_hash.txt").write_text(str(payload.get("data_hash") or "") + "\n")
    (out_dir / "split_hash.txt").write_text(str(payload.get("split_hash") or "") + "\n")
    (out_dir / "train.log").write_text("acceptance gate artifact; no training\n")
    (out_dir / "eval.log").write_text(
        "\n".join(
            [
                f"rollout_metrics={args.rollout_metrics}",
                f"selector_metrics={args.selector_metrics}",
                f"overall_pass={payload['overall_pass']}",
                f"status={payload['status']}",
            ]
        )
        + "\n"
    )
    print(json.dumps({"out_dir": str(out_dir), "overall_pass": payload["overall_pass"]}, indent=2))
    return 0


def _observed_values(rollout: dict[str, Any], selector: dict[str, Any]) -> dict[str, Any]:
    rollout_success = rollout.get("success_summary", {})
    rollout_summary = rollout.get("summary", {})
    selector_summary = selector.get("eval_summary", {})
    selected_success = _number(selector_summary.get("selected_success"))
    proposal_oracle = _number(selector_summary.get("proposal_oracle_success"))
    success_support_gap = _number(selector_summary.get("success_support_gap"))
    success_selector_gap = _number(selector_summary.get("success_selector_gap"))
    outcome_ptr = _summary_mean(rollout_summary, "outcome_ptr_at_16")
    generated_unsafe = _summary_mean(rollout_summary, "generated_unsafe_rate_known_at_16")
    selected_unsafe = _summary_mean(rollout_summary, "selected_unsafe_known_at_16")
    base_unsafe = _summary_mean(rollout_summary, "base_unsafe_known")
    if generated_unsafe is None:
        generated_unsafe = _number(rollout_success.get("generated_unsafe_rate_known"))
    if selected_unsafe is None:
        selected_unsafe = _number(rollout_success.get("selected_unsafe_rate_known"))
    if base_unsafe is None:
        base_unsafe = _number(rollout_success.get("base_unsafe_rate_known"))
    return {
        "base_success": _number(selector_summary.get("base_success")),
        "selected_success": selected_success,
        "selected_success_ci": _summary_ci(selector.get("summary", {}), "selected_success"),
        "target_selected_success": selected_success,
        "proposal_oracle_success": proposal_oracle,
        "proposal_oracle_success_ci": _summary_ci(selector.get("summary", {}), "proposal_oracle_success"),
        "hidden_chart_oracle_success": _number(selector_summary.get("hidden_chart_oracle_success")),
        "success_support_gap": success_support_gap,
        "success_support_gap_ci": _summary_ci(rollout_summary, "success_support_gap_at_16"),
        "success_selector_gap": success_selector_gap,
        "success_selector_gap_ci": _summary_ci(selector.get("summary", {}), "success_selector_gap"),
        "outcome_ptr": outcome_ptr,
        "outcome_ptr_ci": _summary_ci(rollout_summary, "outcome_ptr_at_16"),
        "v0_outcome_ptr_reference": 0.4435,
        "generated_unsafe_rate_known": generated_unsafe,
        "selected_unsafe_rate_known": selected_unsafe,
        "base_unsafe_rate_known": base_unsafe,
        "pairwise_causal_calibration_ece": _number(selector_summary.get("pairwise_causal_calibration_ece")),
        "train_seeds": _train_seed_count(rollout, selector),
        "bootstrap_ci_reported": _has_ci(selector.get("summary", {}), "selected_success")
        and _has_ci(rollout_summary, "outcome_ptr_at_16")
        and _has_ci(rollout_summary, "proposal_oracle_success_at_16"),
    }


def _gates(values: dict[str, Any], args: argparse.Namespace) -> list[dict[str, Any]]:
    selected = _number(values.get("selected_success"))
    proposal = _number(values.get("proposal_oracle_success"))
    support_gap = _number(values.get("success_support_gap"))
    selector_gap = _number(values.get("success_selector_gap"))
    outcome_ptr = _number(values.get("outcome_ptr"))
    generated_unsafe = _number(values.get("generated_unsafe_rate_known"))
    selected_unsafe = _number(values.get("selected_unsafe_rate_known"))
    base_unsafe = _number(values.get("base_unsafe_rate_known"))
    unsafe_limit = (base_unsafe if base_unsafe is not None else 0.0) + float(args.unsafe_slack)
    return [
        _gate("selected clean success", selected, f">= {_fmt(args.selected_min)}", selected is not None and selected >= args.selected_min),
        _gate(
            "target selected clean success",
            selected,
            f">= {_fmt(args.target_selected_min)}",
            selected is not None and selected >= args.target_selected_min,
        ),
        _gate(
            "proposal oracle success",
            proposal,
            f">= {_fmt(args.proposal_oracle_min)}",
            proposal is not None and proposal >= args.proposal_oracle_min,
        ),
        _gate(
            "success support gap",
            support_gap,
            f"<= {_fmt(args.success_support_gap_max)}",
            support_gap is not None and support_gap <= args.success_support_gap_max,
        ),
        _gate(
            "success selector gap",
            selector_gap,
            f"<= {_fmt(args.success_selector_gap_max)}",
            selector_gap is not None and selector_gap <= args.success_selector_gap_max,
        ),
        _gate(
            "true OutcomePTR@16 vs V0",
            outcome_ptr,
            f"> {_fmt(args.v0_outcome_ptr_floor)}",
            outcome_ptr is not None and outcome_ptr > args.v0_outcome_ptr_floor,
        ),
        _gate(
            "generated unsafe rate",
            generated_unsafe,
            f"<= base+slack ({_fmt(unsafe_limit)})",
            generated_unsafe is not None and generated_unsafe <= unsafe_limit,
        ),
        _gate(
            "selected unsafe rate",
            selected_unsafe,
            f"<= base+slack ({_fmt(unsafe_limit)})",
            selected_unsafe is not None and selected_unsafe <= unsafe_limit,
        ),
        _gate(
            "three train seeds",
            values.get("train_seeds"),
            f">= {args.min_train_seeds}",
            int(values.get("train_seeds") or 0) >= int(args.min_train_seeds),
        ),
        _gate(
            "bootstrap confidence intervals",
            float(bool(values.get("bootstrap_ci_reported"))),
            "reported",
            bool(values.get("bootstrap_ci_reported")),
        ),
    ]


def _gate(name: str, observed: Any, required: str, passed: bool) -> dict[str, Any]:
    return {
        "name": name,
        "observed": observed,
        "required": required,
        "pass": bool(passed),
    }


def _group_file_payload(
    selector_metrics: Path,
    rollout_metrics: Path,
    filename: str,
) -> dict[str, Any]:
    return {
        "selector_summary": _load_optional_json(selector_metrics.parent / filename),
        "rollout_summary": _load_optional_json(rollout_metrics.parent / filename),
    }


def _summary_mean(summary: dict[str, Any], key: str) -> float | None:
    return _number(summary.get(key, {}).get("micro", {}).get("mean"))


def _summary_ci(summary: dict[str, Any], key: str) -> dict[str, float] | None:
    micro = summary.get(key, {}).get("micro", {})
    low = _number(micro.get("low"))
    high = _number(micro.get("high"))
    mean = _number(micro.get("mean"))
    if low is None or high is None or mean is None:
        return None
    return {"mean": mean, "low": low, "high": high}


def _has_ci(summary: dict[str, Any], key: str) -> bool:
    return _summary_ci(summary, key) is not None


def _train_seed_count(rollout: dict[str, Any], selector: dict[str, Any]) -> int:
    rollout_seeds = {str(seed) for seed in rollout.get("train_seeds", [])}
    selector_seeds = {
        str(row.get("train_seed"))
        for row in selector.get("rows", [])
        if row.get("train_seed") not in {None, "unknown"}
    }
    seeds = rollout_seeds | selector_seeds
    return len(seeds)


def _table(gates: list[dict[str, Any]]) -> str:
    lines = [
        "% Auto-generated by scripts/build_ctt_outcome_gate.py",
        "\\begin{tabular}{lllc}",
        "\\toprule",
        "Gate & Observed & Required & Status \\\\",
        "\\midrule",
    ]
    for gate in gates:
        status = "pass" if gate["pass"] else "fail"
        lines.append(
            f"{_latex_escape(str(gate['name']))} & {_latex_escape(_fmt(gate['observed']))} & "
            f"{_latex_required(str(gate['required']))} & {status} \\\\"
        )
    lines.extend(["\\bottomrule", "\\end{tabular}"])
    return "\n".join(lines)


def _config(args: argparse.Namespace) -> str:
    return "\n".join(f"{key}: {value}" for key, value in sorted(vars(args).items()))


def _load_json(path: Path) -> dict[str, Any]:
    return json.loads(path.read_text())


def _load_optional_json(path: Path) -> dict[str, Any]:
    if not path.exists():
        return {}
    return json.loads(path.read_text())


def _number(value: Any) -> float | None:
    if isinstance(value, (int, float)) and math.isfinite(float(value)):
        return float(value)
    return None


def _fmt(value: Any) -> str:
    number = _number(value)
    if number is None:
        return "n/a"
    return f"{number:.4f}"


def _latex_escape(value: str) -> str:
    return (
        value.replace("\\", "\\textbackslash{}")
        .replace("_", "\\_")
        .replace("%", "\\%")
        .replace("&", "\\&")
        .replace("<", "$<$")
        .replace(">", "$>$")
    )


def _latex_required(value: str) -> str:
    return (
        _latex_escape(value)
        .replace("$>$=", "$\\ge$")
        .replace("$<$=", "$\\le$")
    )


def _run(command: list[str]) -> str:
    try:
        return subprocess.check_output(command, cwd=PROJECT_ROOT, text=True).strip()
    except (subprocess.CalledProcessError, FileNotFoundError):
        return ""


if __name__ == "__main__":
    raise SystemExit(main())