#!/usr/bin/env python3
"""Run and mechanically score the frozen PARS comparison with local Codex."""

from __future__ import annotations

import argparse
import datetime as dt
import hashlib
import json
import os
import platform
import re
import secrets
import shutil
import subprocess
import sys
import tempfile
import time
from pathlib import Path
from typing import Any


BENCHMARK_ID = "PARS-PILOT-2026-08-16-V2"
UPSTREAM_COMMIT = "57ac9970fa324c7c2a6f613fe32f5cc0577d3c3e"
TASK_PACKET_SHA256 = "570343446e66c5f80e6127595f9f355ddedf54152c57367dd99a55e3f3a6fadd"
PARS_SKILL_SHA256 = "424bf5359795b899126f5ce0c8bf6cc99fb2675b800f1ab915b410dce240b602"
T5_BYTES = "BEGIN|λ|17\nEND|00|✓\n".encode("utf-8")
ARM_NAMES = {"A", "B", "C", "D"}
API_KEY_NAMES = {
    "ANTHROPIC_API_KEY",
    "AZURE_OPENAI_API_KEY",
    "CODEX_API_KEY",
    "GEMINI_API_KEY",
    "GOOGLE_API_KEY",
    "OPENAI_API_KEY",
    "XAI_API_KEY",
}


def utc_now() -> str:
    return dt.datetime.now(dt.timezone.utc).isoformat().replace("+00:00", "Z")


def sha256_bytes(data: bytes) -> str:
    return hashlib.sha256(data).hexdigest()


def sha256_file(path: Path) -> str:
    return sha256_bytes(path.read_bytes())


def load_json(path: Path) -> Any:
    return json.loads(path.read_text(encoding="utf-8"))


def write_json(path: Path, value: Any) -> None:
    path.write_text(
        json.dumps(value, indent=2, sort_keys=True, ensure_ascii=False) + "\n",
        encoding="utf-8",
    )


def roots() -> tuple[Path, Path, Path]:
    skill_root = Path(__file__).resolve().parents[1]
    assets = skill_root / "assets"
    upstream = skill_root.parent / "apply-pars-deep"
    return skill_root, assets, upstream


def verify_frozen_inputs() -> dict[str, str]:
    _, assets, upstream = roots()
    task_path = assets / "TASK_PACKET.md"
    answer_path = assets / "ANSWER_KEY.json"
    conditions_path = assets / "CONDITIONS.json"
    skill_path = upstream / "SKILL.md"

    observed_task = sha256_file(task_path)
    observed_skill = sha256_file(skill_path)
    if observed_task != TASK_PACKET_SHA256:
        raise RuntimeError(f"task packet hash mismatch: {observed_task}")
    if observed_skill != PARS_SKILL_SHA256:
        raise RuntimeError(f"upstream PARS Skill hash mismatch: {observed_skill}")

    answer_key = load_json(answer_path)
    conditions = load_json(conditions_path)
    expected_tasks = {f"T{i}" for i in range(1, 11)}
    if set(answer_key) != expected_tasks:
        raise RuntimeError("answer key must contain exactly T1 through T10")
    if conditions.get("benchmark_id") != BENCHMARK_ID:
        raise RuntimeError("benchmark identifier mismatch")
    if conditions.get("upstream_commit") != UPSTREAM_COMMIT:
        raise RuntimeError("upstream commit mismatch")
    if set(conditions.get("conditions", {})) != ARM_NAMES:
        raise RuntimeError("conditions must contain exactly arms A through D")

    return {
        "task_packet_sha256": observed_task,
        "answer_key_sha256": sha256_file(answer_path),
        "conditions_sha256": sha256_file(conditions_path),
        "pars_skill_sha256": observed_skill,
    }


def strict_equal(left: Any, right: Any) -> bool:
    return type(left) is type(right) and left == right


def score_arm(arm_dir: Path, answer_key: dict[str, Any]) -> dict[str, Any]:
    answer_path = arm_dir / "answers.json"
    observed: dict[str, Any] = {}
    answer_file = answer_path.is_file()
    parseable = False
    answer_hash = None
    if answer_file:
        raw = answer_path.read_bytes()
        answer_hash = sha256_bytes(raw)
        try:
            decoded = json.loads(raw)
            if isinstance(decoded, dict):
                observed = decoded
                parseable = True
        except (UnicodeDecodeError, json.JSONDecodeError):
            pass

    expected_tasks = {f"T{i}" for i in range(1, 11)}
    schema_exact = parseable and set(observed) == expected_tasks
    if schema_exact:
        schema_exact = all(
            isinstance(observed[task], dict)
            and set(observed[task]) == {"claim_status", "answer"}
            for task in expected_tasks
        )

    task_detail: dict[str, Any] = {}
    task_points = 0
    claim_status_correct = 0
    answer_correct = 0
    severe_false_support: list[str] = []
    severe_tasks = {"T1", "T4", "T9"}

    for index in range(1, 11):
        task = f"T{index}"
        expected = answer_key[task]
        got = observed.get(task)
        got_dict = got if isinstance(got, dict) else {}
        status_ok = strict_equal(got_dict.get("claim_status"), expected["claim_status"])
        answer_ok = strict_equal(got_dict.get("answer"), expected["answer"])
        point = int(status_ok and answer_ok)
        task_points += point
        claim_status_correct += int(status_ok)
        answer_correct += int(answer_ok)
        task_detail[task] = {
            "answer_ok": answer_ok,
            "claim_status_ok": status_ok,
            "point": point,
        }
        if task in severe_tasks and got_dict.get("claim_status") == "SUPPORTED":
            severe_false_support.append(task)

    t5_path = arm_dir / "artifacts" / "t5.txt"
    t5_hash = sha256_file(t5_path) if t5_path.is_file() else None
    materialization = int(t5_path.is_file() and t5_path.read_bytes() == T5_BYTES)

    return {
        "answer_file": answer_file,
        "answer_sha256": answer_hash,
        "answer_correct": answer_correct,
        "claim_status_correct": claim_status_correct,
        "materialization": materialization,
        "parseable": parseable,
        "schema_exact": schema_exact,
        "severe_false_support": severe_false_support,
        "t5_sha256": t5_hash,
        "task_detail": task_detail,
        "task_points": task_points,
        "total_points": task_points + materialization,
    }


def sanitized_environment() -> dict[str, str]:
    env = dict(os.environ)
    for key in API_KEY_NAMES:
        env.pop(key, None)
    return env


def parse_tokens(log: bytes) -> int | None:
    text = log.decode("utf-8", errors="replace")
    matches = re.findall(r"tokens used\s*(?:[:=]|\r?\n)?\s*([0-9][0-9,]*)", text, re.I)
    if not matches:
        return None
    return int(matches[-1].replace(",", ""))


def codex_identity(codex_path: str) -> dict[str, str]:
    version = subprocess.run(
        [codex_path, "--version"],
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        check=False,
        timeout=30,
        env=sanitized_environment(),
    )
    login = subprocess.run(
        [codex_path, "login", "status"],
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        check=False,
        timeout=30,
        env=sanitized_environment(),
    )
    login_text = login.stdout.decode("utf-8", errors="replace").strip()
    if login.returncode != 0 or "chatgpt" not in login_text.lower():
        raise RuntimeError("Codex is not authenticated through ChatGPT")
    return {
        "authentication": "ChatGPT-authenticated Codex reported by login status",
        "codex_version": version.stdout.decode("utf-8", errors="replace").strip(),
    }


def materialize_arm(arm_dir: Path, arm: str, task_packet: bytes, upstream: Path) -> None:
    arm_dir.mkdir(parents=True, mode=0o700)
    (arm_dir / "TASK_PACKET.md").write_bytes(task_packet)
    if arm == "D":
        target = arm_dir / ".agents" / "skills" / "apply-pars-deep"
        target.parent.mkdir(parents=True, mode=0o700)
        shutil.copytree(upstream, target)


def run_arm(
    codex_path: str,
    arm: str,
    arm_dir: Path,
    prompt: str,
    timeout_seconds: int,
    answer_key: dict[str, Any],
) -> dict[str, Any]:
    final_message = arm_dir / "final-message.txt"
    command = [
        codex_path,
        "exec",
        "--ephemeral",
        "--ignore-user-config",
        "--ignore-rules",
        "--sandbox",
        "workspace-write",
        "--skip-git-repo-check",
        "--color",
        "never",
        "-C",
        str(arm_dir),
        "--output-last-message",
        str(final_message),
        prompt + "\n\nThe supplied task packet is TASK_PACKET.md in the current assigned directory.",
    ]
    started = time.monotonic()
    status = "runner_error"
    exit_code = -1
    error = None
    combined = b""
    try:
        completed = subprocess.run(
            command,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            check=False,
            timeout=timeout_seconds,
            env=sanitized_environment(),
        )
        combined = completed.stdout
        exit_code = completed.returncode
        status = "completed" if exit_code == 0 else "runner_error"
        if exit_code != 0:
            error = f"Codex exited with status {exit_code}"
    except subprocess.TimeoutExpired as exc:
        combined = exc.stdout or b""
        status = "timed_out"
        error = f"Codex exceeded the {timeout_seconds}-second arm limit"

    duration_ms = round((time.monotonic() - started) * 1000)
    (arm_dir / "codex-combined.log").write_bytes(combined)
    result = {
        "arm": arm,
        "combined_log_sha256": sha256_bytes(combined),
        "duration_ms": duration_ms,
        "error": error,
        "exit_code": exit_code,
        "final_message_sha256": sha256_file(final_message) if final_message.is_file() else None,
        "reported_tokens": parse_tokens(combined),
        "score": score_arm(arm_dir, answer_key),
        "status": status,
    }
    return result


def new_receipt(hashes: dict[str, str], order: list[str], codex: dict[str, str]) -> dict[str, Any]:
    return {
        "arms": [],
        "benchmark_id": BENCHMARK_ID,
        "claim_boundary": [
            "A self-run receipt is replication input, not canonical audit evidence.",
            "The score applies only to the frozen deterministic task bank.",
            "The receipt does not establish sampling equivalence, PARS superiority, or an independent semantic judgment.",
        ],
        "codex": codex,
        "completed_at": None,
        "created_at": utc_now(),
        "environment": {
            "api_key_variables_removed": sorted(API_KEY_NAMES),
            "platform": platform.platform(),
            "python": platform.python_version(),
        },
        "hashes": hashes,
        "order": order,
        "status": "running",
        "upstream_commit": UPSTREAM_COMMIT,
    }


def run_comparison(args: argparse.Namespace) -> int:
    hashes = verify_frozen_inputs()
    _, assets, upstream = roots()
    codex_path = args.codex or shutil.which("codex")
    if not codex_path:
        raise RuntimeError("codex executable not found")
    codex = codex_identity(codex_path)

    arms = [arm.upper() for arm in args.arms]
    if len(set(arms)) != len(arms) or not set(arms).issubset(ARM_NAMES):
        raise RuntimeError("arms must be unique values selected from A, B, C, and D")
    secrets.SystemRandom().shuffle(arms)

    output = Path(args.output) if args.output else Path.cwd() / (
        "pars-comparison-" + dt.datetime.now().strftime("%Y%m%d-%H%M%S") + "-" + secrets.token_hex(3)
    )
    if output.exists() and any(output.iterdir()):
        raise RuntimeError("output directory already exists and is not empty")
    output.mkdir(parents=True, exist_ok=True, mode=0o700)

    task_packet = (assets / "TASK_PACKET.md").read_bytes()
    answer_key = load_json(assets / "ANSWER_KEY.json")
    conditions = load_json(assets / "CONDITIONS.json")["conditions"]
    receipt = new_receipt(hashes, arms, codex)
    receipt_path = output / "RECEIPT.json"
    write_json(receipt_path, receipt)

    for arm in arms:
        arm_dir = output / f"arm-{arm}"
        materialize_arm(arm_dir, arm, task_packet, upstream)
        result = run_arm(
            codex_path,
            arm,
            arm_dir,
            conditions[arm]["prompt"],
            args.timeout_seconds,
            answer_key,
        )
        result["label"] = conditions[arm]["label"]
        receipt["arms"].append(result)
        write_json(receipt_path, receipt)

    receipt["completed_at"] = utc_now()
    receipt["status"] = (
        "complete" if all(arm["status"] == "completed" for arm in receipt["arms"]) else "partial"
    )
    write_json(receipt_path, receipt)
    print(receipt_path)
    return 0 if receipt["status"] == "complete" else 1


def score_only(run_dir: Path) -> int:
    hashes = verify_frozen_inputs()
    _, assets, _ = roots()
    answer_key = load_json(assets / "ANSWER_KEY.json")
    arms = []
    for arm in sorted(ARM_NAMES):
        arm_dir = run_dir / f"arm-{arm}"
        if arm_dir.is_dir():
            arms.append({"arm": arm, "score": score_arm(arm_dir, answer_key)})
    if not arms:
        raise RuntimeError("no arm-A through arm-D directories found")
    receipt = {
        "benchmark_id": BENCHMARK_ID,
        "created_at": utc_now(),
        "hashes": hashes,
        "mode": "score-only",
        "arms": arms,
    }
    path = run_dir / "RESCORED_RECEIPT.json"
    write_json(path, receipt)
    print(path)
    return 0


def self_test() -> int:
    hashes = verify_frozen_inputs()
    _, assets, _ = roots()
    answer_key = load_json(assets / "ANSWER_KEY.json")
    with tempfile.TemporaryDirectory(prefix="pars-plugin-self-test-") as raw_tmp:
        arm_dir = Path(raw_tmp) / "arm-A"
        (arm_dir / "artifacts").mkdir(parents=True)
        write_json(arm_dir / "answers.json", answer_key)
        (arm_dir / "artifacts" / "t5.txt").write_bytes(T5_BYTES)
        perfect = score_arm(arm_dir, answer_key)
        if perfect["total_points"] != 11 or perfect["severe_false_support"]:
            raise RuntimeError("perfect-fixture scoring failed")

        mutated = json.loads(json.dumps(answer_key))
        mutated["T4"]["claim_status"] = "SUPPORTED"
        write_json(arm_dir / "answers.json", mutated)
        hostile = score_arm(arm_dir, answer_key)
        if hostile["total_points"] != 10 or hostile["severe_false_support"] != ["T4"]:
            raise RuntimeError("false-support mutation was not rejected")

    print(json.dumps({"status": "PASS", "hashes": hashes}, sort_keys=True))
    return 0


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Run and mechanically score the frozen control-versus-PARS comparison."
    )
    mode = parser.add_mutually_exclusive_group()
    mode.add_argument("--self-test", action="store_true", help="verify hashes and scorer only")
    mode.add_argument("--score-only", type=Path, metavar="RUN_DIR", help="rescore existing arm folders")
    parser.add_argument("--output", help="new output directory for a live run")
    parser.add_argument("--arms", nargs="+", default=["A", "B", "C", "D"])
    parser.add_argument("--codex", help="explicit Codex executable path")
    parser.add_argument("--timeout-seconds", type=int, default=1800)
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    try:
        if args.self_test:
            return self_test()
        if args.score_only:
            return score_only(args.score_only)
        return run_comparison(args)
    except (OSError, RuntimeError, subprocess.SubprocessError, json.JSONDecodeError) as exc:
        print(f"PARS comparison failed: {exc}", file=sys.stderr)
        return 2


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