from __future__ import annotations

import argparse
import hashlib
import json
import os
import pathlib
import platform
import re
import signal
import subprocess
import sys
import tempfile
import time
from typing import Any


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

from reliability_core.process_control import run_supervised_process
from reliability_core.stage_control_r10r3r2 import cleanup_wait_seconds
from reliability_core.suite_isolation import sha256_file
from scripts.build_evidence import normalized_tree_hash


def _terminate(process: subprocess.Popen[bytes]) -> None:
    if process.poll() is not None:
        return
    if os.name == "nt":
        subprocess.run(["taskkill", "/PID", str(process.pid), "/T", "/F"], capture_output=True, check=False)
    else:
        try:
            os.killpg(process.pid, signal.SIGTERM)
        except ProcessLookupError:
            pass
    try:
        process.wait(timeout=cleanup_wait_seconds())
    except subprocess.TimeoutExpired:
        process.kill(); process.wait(timeout=cleanup_wait_seconds())


def _start_cpu_load() -> subprocess.Popen[bytes]:
    code = "import time\nend=time.monotonic()+1200\nx=1\nwhile time.monotonic()<end:x=(x*1103515245+12345)&0x7fffffff"
    return subprocess.Popen([sys.executable, "-c", code], stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, start_new_session=os.name != "nt", creationflags=getattr(subprocess, "CREATE_NEW_PROCESS_GROUP", 0) if os.name == "nt" else 0)


def _watchdog(stdout: str) -> dict[str, Any] | None:
    rows = []
    for line in stdout.splitlines():
        if line.startswith("WATCHDOG_REPORT "):
            try:
                rows.append(json.loads(line.split(" ", 1)[1]))
            except json.JSONDecodeError:
                pass
    return rows[-1] if rows else None


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--output-dir", required=True, type=pathlib.Path)
    parser.add_argument("--run-id", required=True)
    parser.add_argument("--run-ordinal", required=True, type=int)
    parser.add_argument("--source-origin", required=True, choices=["active", "candidate", "final", "rollback"])
    parser.add_argument("--order", required=True, choices=["default", "reverse", "random", "sensitive-first", "sensitive-last"])
    parser.add_argument("--seed", type=int)
    parser.add_argument("--load-profile", default="normal", choices=["normal", "cpu-medium", "slow-startup"])
    parser.add_argument("--timeout-seconds", type=float, default=1200.0)
    args = parser.parse_args()
    if args.order == "random" and args.seed is None:
        parser.error("random order requires --seed")
    output = args.output_dir.resolve(); output.mkdir(parents=True, exist_ok=True)
    stdout_path = output / f"{args.run_id}.stdout.txt"; stderr_path = output / f"{args.run_id}.stderr.txt"
    command = [sys.executable, "run_all.py", "--tests-only", "--order", args.order, "--load-profile", args.load_profile]
    if args.seed is not None:
        command.extend(["--seed", str(args.seed)])
    environment = {**os.environ, "PYTHONDONTWRITEBYTECODE": "1", "PYTHONPYCACHEPREFIX": os.environ.get("PYTHONPYCACHEPREFIX", str(pathlib.Path(tempfile.gettempdir()) / f"r9-pycache-{os.getpid()}-{time.time_ns()}"))}
    injection_context: tempfile.TemporaryDirectory[str] | None = None
    if args.load_profile == "slow-startup":
        injection_context = tempfile.TemporaryDirectory(prefix="r9-startup-injection-")
        injection = pathlib.Path(injection_context.name)
        (injection / "sitecustomize.py").write_text("import time\ntime.sleep(.35)\n", encoding="utf-8")
        environment["PYTHONPATH"] = str(injection)
        environment["WRLAB_R9_INJECTED_STARTUP_DELAY_SECONDS"] = ".35"
    load_process = _start_cpu_load() if args.load_profile == "cpu-medium" else None
    source_hash_before = normalized_tree_hash(ROOT)
    started_epoch = time.time()
    try:
        record = run_supervised_process(command, ROOT, environment=environment, total_timeout_seconds=args.timeout_seconds, max_output_bytes=16 * 1024 * 1024)
    finally:
        if load_process is not None:
            _terminate(load_process)
        if injection_context is not None:
            injection_context.cleanup()
    ended_epoch = time.time()
    stdout_path.write_text(record["stdout"], encoding="utf-8")
    stderr_path.write_text(record["stderr"], encoding="utf-8")
    transcript = record["stdout"] + "\n" + record["stderr"]
    counts = re.findall(r"Ran\s+(\d+)\s+tests?\s+in", transcript)
    watchdog = _watchdog(record["stdout"])
    source_hash_after = normalized_tree_hash(ROOT)
    payload = {
        "run_id": args.run_id,
        "run_ordinal": args.run_ordinal,
        "cold_start": True,
        "suite_kind": "full",
        "target": None,
        "source_origin": args.source_origin,
        "source_tree_sha256": source_hash_after,
        "source_hash_stable": source_hash_before == source_hash_after,
        "policy_manifest_sha256": sha256_file(ROOT / "test_execution_manifest.json"),
        "order_name": args.order,
        "random_seed": args.seed,
        "load_profile": args.load_profile,
        "platform": sys.platform,
        "platform_detail": platform.platform(),
        "python_version": platform.python_version(),
        "python_executable": sys.executable,
        "resource_warning_error_mode": "error::ResourceWarning" in environment.get("PYTHONWARNINGS", ""),
        "command": command,
        "started_epoch": started_epoch,
        "ended_epoch": ended_epoch,
        "elapsed_seconds": round(ended_epoch - started_epoch, 3),
        "worker_pid": record["pid"],
        "exit_code": record["exit_code"],
        "timed_out": record["timed_out"],
        "orphan_process_count": record["orphan_process_count"],
        "zombie_process_count": record["zombie_process_count"],
        "output_drain_threads_stopped": record["output_drain_threads_stopped"],
        "stdout_file": stdout_path.name,
        "stderr_file": stderr_path.name,
        "stdout_sha256": hashlib.sha256(stdout_path.read_bytes()).hexdigest(),
        "stderr_sha256": hashlib.sha256(stderr_path.read_bytes()).hexdigest(),
        "test_count": int(counts[-1]) if counts else None,
        "watchdog_report": watchdog,
    }
    payload["passed"] = all((payload["exit_code"] == 0, payload["timed_out"] is False, payload["orphan_process_count"] == 0, payload["zombie_process_count"] == 0, payload["output_drain_threads_stopped"] is True, payload["source_hash_stable"] is True, isinstance(payload["test_count"], int), isinstance(watchdog, dict) and watchdog.get("passed") is True and watchdog.get("executed_test_count") == payload["test_count"] and watchdog.get("manifest_test_count") == payload["test_count"] and watchdog.get("duplicate_execution_count") == 0 and watchdog.get("missing_execution_count") == 0 and watchdog.get("process_sensitive_shared_worker_count") == 0))
    target = output / f"{args.run_id}.json"
    target.write_text(json.dumps(payload, ensure_ascii=False, indent=2, sort_keys=True), encoding="utf-8")
    print(json.dumps({"passed": payload["passed"], "record": str(target), "test_count": payload["test_count"], "exit_code": payload["exit_code"]}, sort_keys=True))
    return 0 if payload["passed"] else 2


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