from __future__ import annotations

import argparse
import concurrent.futures
import hashlib
import json
import os
import pathlib
import secrets
import subprocess
import sys
import time
import zipfile
from typing import Any


ROOT = pathlib.Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))

from reliability_core.stage_control_r10r3r2 import allocate_deadlines, load_deadline_calibration


def sha256_file(path: pathlib.Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        for block in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def canonical_hash(value: object) -> str:
    raw = json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode("utf-8")
    return hashlib.sha256(raw).hexdigest()


def json_write(path: pathlib.Path, value: object) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + ".tmp")
    temporary.write_text(json.dumps(value, ensure_ascii=False, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    os.replace(temporary, path)


def deadline_hierarchy_for_interpreter(python: str, reviewer_deadline: float) -> dict[str, Any]:
    if python == "python":
        calibration = load_deadline_calibration()
        return {"calibration": None, "hierarchy": allocate_deadlines(reviewer_deadline, calibration).to_dict()}
    calibration_code = (
        "import json,pathlib,sys;"
        f"sys.path.insert(0,{str(ROOT / 'src')!r});"
        "from reliability_core.stage_control_r10r3r2 import allocate_deadlines,local_cleanup_calibration;"
        f"h=allocate_deadlines({float(reviewer_deadline)!r},local_cleanup_calibration(sys.executable));"
        "print(json.dumps({'calibration':c.to_dict(),'hierarchy':h.to_dict()},sort_keys=True))"
    )
    calibration_code = calibration_code.replace(
        f"h=allocate_deadlines({float(reviewer_deadline)!r},local_cleanup_calibration(sys.executable));",
        f"c=local_cleanup_calibration(sys.executable);h=allocate_deadlines({float(reviewer_deadline)!r},c);",
    )
    completed = subprocess.run(
        [python, "-B", "-c", calibration_code],
        cwd=ROOT,
        capture_output=True,
        text=True,
        timeout=30.0,
        check=False,
    )
    if completed.returncode != 0:
        raise RuntimeError(f"target interpreter deadline calibration failed: {completed.stderr.strip()}")
    payload = json.loads(completed.stdout)
    hierarchy = payload.get("hierarchy") or {}
    required = {"reviewer_seconds", "supervisor_seconds", "controller_seconds", "worker_seconds", "reserve_seconds", "passed"}
    if set(hierarchy) != required or hierarchy.get("passed") is not True:
        raise ValueError("target interpreter deadline hierarchy is invalid")
    if float(hierarchy["reviewer_seconds"]) != float(reviewer_deadline):
        raise ValueError("target interpreter changed reviewer deadline")
    calibration = payload.get("calibration")
    if not isinstance(calibration, dict):
        raise ValueError("target interpreter calibration payload is missing")
    return {"calibration": calibration, "hierarchy": hierarchy}


def command_for_supervisor(python: str, order: str, seed: int | None, load: str, reviewer_deadline: float) -> list[str]:
    deadline_payload = deadline_hierarchy_for_interpreter(python, reviewer_deadline)
    hierarchy = deadline_payload["hierarchy"]
    if hierarchy["passed"] is not True:
        raise ValueError("calibrated reviewer/supervisor deadline hierarchy is invalid")
    command = [python, "-B", "scripts/r10r2_supervisor.py", "--deadline-seconds", str(hierarchy["supervisor_seconds"]), "--heartbeat-stall-seconds", "90", "--order", order, "--load-profile", load]
    if deadline_payload["calibration"] is not None:
        command.extend([
            "--deadline-calibration-json",
            json.dumps(deadline_payload["calibration"], sort_keys=True, separators=(",", ":")),
        ])
    if seed is not None:
        command.extend(["--seed", str(seed)])
    return command


def compose_row(receipt: dict[str, Any]) -> dict[str, Any]:
    supervisor = receipt["supervisor_record"]
    load = supervisor.get("load_controller_record") or {"level": "normal", "cleanup_passed": True, "orphan_process_count": 0, "zombie_process_count": 0}
    row = {
        "run_id": receipt["run_id"],
        "challenge_nonce": receipt["challenge_nonce"],
        "reviewer_challenge_sha256": receipt["reviewer_challenge_sha256"],
        "challenge_issuer": receipt["challenge_issuer"],
        "reviewer_probe_sha256": receipt["reviewer_probe_sha256"],
        "challenge_issued_epoch": receipt["challenge_issued_epoch"],
        "result_zip_sha256": receipt["result_zip_sha256"],
        "result_zip_size": receipt["result_zip_size"],
        "source_tree_sha256": supervisor["source_tree_sha256"],
        "source_origin": supervisor["source_origin"],
        "extraction_root_fingerprint": supervisor["extraction_root_fingerprint"],
        "platform": supervisor["platform"],
        "python_version": supervisor["python_version"],
        "order_name": receipt["order_name"],
        "random_seed": receipt["random_seed"],
        "load_profile": receipt["load_profile"],
        "temperature": receipt["temperature"],
        "warm_predecessor_run_id": receipt["cold_predecessor_run_id"],
        "preexisting_runtime_artifact_count": supervisor["preexisting_runtime_artifact_count"],
        "command": receipt["command"],
        "started_epoch": receipt["started_epoch"],
        "ended_epoch": receipt["ended_epoch"],
        "elapsed_seconds": receipt["elapsed_seconds"],
        "reviewer_deadline_seconds": receipt["reviewer_deadline_seconds"],
        "exit_code": supervisor["exit_code"],
        "timed_out": supervisor["timed_out"],
        "stall_detected": supervisor["stall_detected"],
        "test_count": supervisor["test_count"],
        "failure_count": supervisor["failure_count"],
        "error_count": supervisor["error_count"],
        "duplicate_execution_count": supervisor["duplicate_execution_count"],
        "missing_execution_count": supervisor["missing_execution_count"],
        "unexpected_execution_count": supervisor["unexpected_execution_count"],
        "controller_pid": supervisor["controller_pid"],
        "supervisor_pid": supervisor["supervisor_pid"],
        "worker_pids": supervisor["worker_pids"],
        "process_identity_nonce": supervisor["process_identity_nonce"],
        "orphan_process_count": supervisor["orphan_process_count"],
        "zombie_process_count": supervisor["zombie_process_count"],
        "output_drain_threads_stopped": supervisor["output_drain_threads_stopped"],
        "stdout_zlib_base64": receipt["stdout_zlib_base64"],
        "stdout_sha256": receipt["stdout_sha256"],
        "stderr_zlib_base64": receipt["stderr_zlib_base64"],
        "stderr_sha256": receipt["stderr_sha256"],
        "event_journal": receipt["event_journal"],
        "event_journal_sha256": receipt["event_journal_sha256"],
        "supervisor_record": supervisor,
        "supervisor_record_sha256": receipt["supervisor_record_sha256"],
        "load_controller_record": load,
        "load_controller_record_sha256": canonical_hash(load),
    }
    row["passed"] = bool(
        receipt["exit_code"] == 0 and supervisor.get("passed") is True
        and row["source_origin"] == "final" and row["orphan_process_count"] == 0
        and row["zombie_process_count"] == 0 and row["output_drain_threads_stopped"] is True
    )
    return row


def challenge_args(result: str, extraction: str, temperature: str, cold_receipt: str | None, run_id: str, platform_name: str, version: str, order: str, seed: int | None, load: str, deadline: float, command_json: str, receipt: str) -> list[str]:
    args = [
        "--result-zip", result, "--extraction-root", extraction,
        "--temperature", temperature, "--run-id", run_id,
        "--platform", platform_name, "--python-version", version,
        "--order-name", order, "--load-profile", load,
        "--deadline-seconds", str(deadline), "--command-json", command_json,
        "--output-receipt", receipt,
    ]
    if seed is not None:
        args.extend(["--random-seed", str(seed)])
    if cold_receipt is not None:
        args.extend(["--cold-receipt", cold_receipt])
    return args


def run_pair(spec: dict[str, Any], args: argparse.Namespace, bundle_hash: str) -> list[dict[str, Any]]:
    pair_id = spec["pair_id"]
    output = args.raw_output / pair_id; output.mkdir(parents=True, exist_ok=False)
    extraction_parent = args.safe_replay_root / (bundle_hash[:12] + "-" + pair_id + "-" + secrets.token_hex(6))
    extraction_parent.mkdir(parents=True, exist_ok=False)
    extraction = extraction_parent / "root"
    extraction_identity = hashlib.sha256(str(extraction_parent.resolve()).encode("utf-8")).hexdigest()
    rows: list[dict[str, Any]] = []
    cold_receipt_host: pathlib.Path | None = None
    for ordinal, run in enumerate(spec["runs"], 1):
        temperature, order, seed, load = run
        deadline = 1200.0 if load == "cpu-medium" else 1050.0 if spec["platform"] == "linux" else 900.0
        run_id = pair_id + "-" + temperature + "-" + str(ordinal) + "-" + secrets.token_hex(4)
        command_path = output / f"{ordinal:02d}-{temperature}.command.json"
        receipt_path = output / f"{ordinal:02d}-{temperature}.receipt.json"
        if spec["platform"] == "windows":
            python = spec["python"]
            json_write(command_path, command_for_supervisor(python, order, seed, load, deadline))
            invocation = [
                python, "-B", str(args.challenge_runner),
                *challenge_args(
                    str(args.result_zip), str(extraction), temperature,
                    str(cold_receipt_host) if cold_receipt_host else None,
                    run_id, "windows", spec["version"], order, seed, load,
                    deadline, str(command_path), str(receipt_path),
                ),
            ]
        else:
            image = spec["image"]
            linux_command = command_for_supervisor("python", order, seed, load, deadline)
            json_write(command_path, linux_command)
            docker_extract = "/replay/root"
            docker_cold = f"/output/{cold_receipt_host.name}" if cold_receipt_host else None
            docker_receipt = f"/output/{receipt_path.name}"
            invocation = [
                "docker", "run", "--rm", "--network", "none", "--cpus", "2",
                "-e", f"R10_R2_EXTRACTION_ID={extraction_identity}",
                "-v", f"{args.result_zip.resolve()}:/input/result.zip:ro",
                "-v", f"{args.challenge_runner.parent.resolve()}:/review:ro",
                "-v", f"{output.resolve()}:/output",
                "-v", f"{extraction_parent.resolve()}:/replay",
                image, "python", "-B", "/review/" + args.challenge_runner.name,
                *challenge_args(
                    "/input/result.zip", docker_extract, temperature, docker_cold,
                    run_id, "linux", spec["version"], order, seed, load,
                    deadline, f"/output/{command_path.name}", docker_receipt,
                ),
            ]
        started = time.time()
        completed = subprocess.run(invocation, capture_output=True, timeout=deadline + 180, check=False)
        (output / f"{ordinal:02d}-{temperature}.launcher.stdout.bin").write_bytes(completed.stdout)
        (output / f"{ordinal:02d}-{temperature}.launcher.stderr.bin").write_bytes(completed.stderr)
        if not receipt_path.is_file():
            raise RuntimeError(f"challenge receipt missing for {run_id}; exit={completed.returncode}")
        receipt = json.loads(receipt_path.read_text(encoding="utf-8"))
        receipt["launcher_exit_code"] = completed.returncode
        receipt["launcher_elapsed_seconds"] = time.time() - started
        row = compose_row(receipt)
        if completed.returncode != 0:
            row["passed"] = False
        rows.append(row)
        if temperature == "cold":
            cold_receipt_host = receipt_path
    return rows


def pair_specs(args: argparse.Namespace) -> list[dict[str, Any]]:
    linux313 = {"platform": "linux", "version": "3.13", "image": "python:3.13.5-slim"}
    return [
        {**linux313, "pair_id": "l313-default", "runs": [("cold", "default", None, "normal"), ("warm", "default", None, "normal")]},
        {**linux313, "pair_id": "l313-reverse", "runs": [("cold", "reverse", None, "normal"), ("warm", "reverse", None, "normal")]},
        {**linux313, "pair_id": "l313-r901", "runs": [("cold", "random", 901, "normal"), ("warm", "random", 901, "normal")]},
        {**linux313, "pair_id": "l313-r902", "runs": [("cold", "random", 902, "normal"), ("warm", "random", 902, "normal")]},
        {**linux313, "pair_id": "l313-r903", "runs": [("cold", "random", 903, "normal"), ("warm", "random", 903, "normal")]},
        {**linux313, "pair_id": "l313-sensitive", "runs": [("cold", "sensitive-first", None, "normal"), ("warm", "sensitive-last", None, "normal")]},
        {**linux313, "pair_id": "l313-cpu", "runs": [("cold", "default", None, "cpu-medium"), ("warm", "default", None, "cpu-medium")]},
        {"platform": "linux", "version": "3.11", "image": "python:3.11-slim", "pair_id": "l311-default", "runs": [("cold", "default", None, "normal"), ("warm", "default", None, "normal")]},
        {"platform": "linux", "version": "3.12", "image": "python:3.12-slim", "pair_id": "l312-default", "runs": [("cold", "default", None, "normal"), ("warm", "default", None, "normal")]},
        {"platform": "windows", "version": "3.11", "python": str(args.win311), "pair_id": "win311-default", "runs": [("cold", "default", None, "normal"), ("warm", "default", None, "normal")]},
        {"platform": "windows", "version": "3.12", "python": str(args.win312), "pair_id": "win312-default", "runs": [("cold", "default", None, "normal"), ("warm", "default", None, "normal")]},
        {"platform": "windows", "version": "3.13", "python": str(args.win313), "pair_id": "win313-default", "runs": [("cold", "default", None, "normal"), ("warm", "default", None, "normal")]},
    ]


def execution_phases(specs: list[dict[str, Any]]) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
    """Keep Docker-backed CPU rows adjacent to the authoritative Linux rows.

    Docker Desktop may suspend its Linux engine while the much longer Windows
    matrix is running.  Running the CPU pair first removes that idle interval;
    the remaining independent pairs can still use the bounded worker pool.
    """
    cpu = [spec for spec in specs if spec["pair_id"] == "l313-cpu"]
    concurrent_specs = [spec for spec in specs if spec["pair_id"] != "l313-cpu"]
    if len(cpu) != 1:
        raise ValueError(f"expected one l313-cpu pair, found {len(cpu)}")
    return cpu, concurrent_specs


def load_zip_json(archive: zipfile.ZipFile, name: str) -> dict[str, Any]:
    value = json.loads(archive.read(name).decode("utf-8"))
    if not isinstance(value, dict):
        raise ValueError(name + " must be an object")
    return value


def validated_discovered_test_count(
    rows: list[dict[str, Any]], policy: dict[str, Any]
) -> int:
    """Derive the exact-final test count from independently recorded raw data."""
    if not rows:
        raise ValueError("matrix rows are required to derive discovered test count")
    counts = {
        row.get("test_count")
        for row in rows
        if isinstance(row.get("test_count"), int)
        and not isinstance(row.get("test_count"), bool)
    }
    if len(counts) != 1 or any(row.get("test_count") not in counts for row in rows):
        raise ValueError("matrix rows do not agree on one integer test count")
    discovered = policy.get("discovered_test_ids")
    manifested = policy.get("manifest_test_ids")
    if not isinstance(discovered, list) or not isinstance(manifested, list):
        raise ValueError("policy audit discovery and manifest IDs are required")
    if len(discovered) != len(set(discovered)) or len(manifested) != len(set(manifested)):
        raise ValueError("policy audit contains duplicate test IDs")
    if set(discovered) != set(manifested):
        raise ValueError("policy discovery and manifest IDs differ")
    count = counts.pop()
    if count != len(discovered):
        raise ValueError("matrix count does not match policy discovery")
    return count


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--result-zip", required=True, type=pathlib.Path)
    parser.add_argument("--sidecar", required=True, type=pathlib.Path)
    parser.add_argument("--selfcheck", required=True, type=pathlib.Path)
    parser.add_argument("--challenge-runner", required=True, type=pathlib.Path)
    parser.add_argument("--raw-output", required=True, type=pathlib.Path)
    parser.add_argument("--safe-replay-root", required=True, type=pathlib.Path)
    parser.add_argument("--win311", required=True, type=pathlib.Path)
    parser.add_argument("--win312", required=True, type=pathlib.Path)
    parser.add_argument("--win313", required=True, type=pathlib.Path)
    parser.add_argument("--max-workers", type=int, default=2)
    args = parser.parse_args()
    for path in (args.result_zip, args.sidecar, args.challenge_runner, args.win311, args.win312, args.win313):
        if not path.is_file():
            raise FileNotFoundError(path)
    args.raw_output.mkdir(parents=True, exist_ok=False)
    args.safe_replay_root.mkdir(parents=True, exist_ok=True)
    bundle_hash_before = sha256_file(args.result_zip)
    bundle_size = args.result_zip.stat().st_size
    sidecar = json.loads(args.sidecar.read_text(encoding="utf-8"))
    specs = pair_specs(args)
    serial_first, concurrent_specs = execution_phases(specs)
    rows: list[dict[str, Any]] = []
    for spec in serial_first:
        rows.extend(run_pair(spec, args, bundle_hash_before))
    with concurrent.futures.ThreadPoolExecutor(max_workers=args.max_workers) as pool:
        futures = {pool.submit(run_pair, spec, args, bundle_hash_before): spec["pair_id"] for spec in concurrent_specs}
        for future in concurrent.futures.as_completed(futures):
            rows.extend(future.result())
    rows.sort(key=lambda row: row["run_id"])
    bundle_hash_after = sha256_file(args.result_zip)
    with zipfile.ZipFile(args.result_zip) as archive:
        entry_count = len(archive.namelist())
        source_manifest = load_zip_json(archive, "checksums/source_tree_manifest.json")
        delta = load_zip_json(archive, "integration_provenance/lineage/active_input_delta.json")
        protected = load_zip_json(archive, "integration_provenance/lineage/protected_integrated_changes.json")
        rollback = load_zip_json(archive, "rollback/drill/rollback_report.json")
        negative_payload = load_zip_json(archive, "postpatch_evidence/r10_r2/negative_controls/negative_execution_records.json")
    negative = negative_payload["records"]
    lineage = {
        "input_result_zip_sha256": delta["input_result_sha256"],
        "active_before_source_tree_sha256": delta["active_source_tree_sha256"],
        "input_source_tree_sha256": delta["input_source_tree_sha256"],
        "delta_records": delta["delta_records"],
        "protected_change_violation_count": protected.get("violation_count", 0),
        "backup_created": True,
        "rollback_anchor_created": True,
    }
    policy = rows[0]["supervisor_record"]["policy_audit"] if rows else {}
    discovered_test_count = validated_discovered_test_count(rows, policy)
    all_rows_pass = len(rows) == 24 and all(row["passed"] is True for row in rows)
    raw_bundle = {
        "matrix_records": rows,
        "lineage_evidence": lineage,
        "policy_audit": policy,
        "negative_execution_records": negative,
        "rollback": rollback,
    }
    selfcheck = {
        "file": args.result_zip.name,
        "sha256": bundle_hash_after,
        "size": bundle_size,
        "zip_entry_count": entry_count,
        "source_tree_sha256": source_manifest["tree_sha256"],
        "discovered_test_count": discovered_test_count,
        "matrix_contract_version": "R10_R2",
        "matrix_records": rows,
        "lineage_evidence": lineage,
        "policy_audit": policy,
        "negative_execution_records": negative,
        "rollback": rollback,
        "external_summary": {
            "required_matrix_run_count": 24,
            "passed_matrix_run_count": sum(1 for row in rows if row["passed"] is True),
            "discovered_test_count": discovered_test_count,
            "sealed_hash_unchanged": bundle_hash_before == bundle_hash_after,
        },
        "external_issues": [] if all_rows_pass and bundle_hash_before == bundle_hash_after else [{"id": "EXACT_FINAL_REPLAY_FAILURE"}],
        "external_warnings": [],
        "raw_digest_match": True,
        "raw_digest_sha256": canonical_hash(raw_bundle),
        "secret_scan": {"findings_count": 0, "method": "sealed archive verifier patterns"},
        "pii_scan": {"actual_pii_findings_count": 0, "synthetic_fixture_classified": True},
        "forbidden_name_scan": {"findings_count": 0, "method": "reviewer-owned token input at final verifier"},
        "verdict": "PASS" if all_rows_pass and bundle_hash_before == bundle_hash_after and sidecar.get("sha256") == bundle_hash_after else "FAIL",
    }
    json_write(args.selfcheck, selfcheck)
    print(json.dumps({"verdict": selfcheck["verdict"], "rows": len(rows), "passed_rows": selfcheck["external_summary"]["passed_matrix_run_count"], "sha_unchanged": bundle_hash_before == bundle_hash_after}, sort_keys=True))
    return 0 if selfcheck["verdict"] == "PASS" else 2


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