from __future__ import annotations

import argparse
import json
import os
import pathlib
import subprocess
import sys
import time


def atomic_ready(path: pathlib.Path, payload: dict[str, object]) -> None:
    payload = {**payload, "event_epoch": time.time_ns() / 1_000_000_000}
    temporary = path.with_suffix(path.suffix + ".tmp")
    temporary.write_text(json.dumps(payload, sort_keys=True), encoding="utf-8")
    temporary.replace(path)


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("ready_path", type=pathlib.Path)
    parser.add_argument("--bootstrapped-path", type=pathlib.Path)
    parser.add_argument("--worker", action="store_true")
    parser.add_argument("--escape-session", action="store_true")
    parser.add_argument("--ready-delay", type=float, default=0.0)
    parser.add_argument("--never-ready", action="store_true")
    parser.add_argument("--single-process-ready", action="store_true")
    args = parser.parse_args()
    if args.single_process_ready:
        time.sleep(max(0.0, args.ready_delay))
        if not args.never_ready:
            atomic_ready(args.ready_path, {
                "status": "READY",
                "ready": True,
                "child_pid": os.getpid(),
                "grandchild_pid": None,
                "escaped_session": False,
                "phase": "workload_ready",
                "intended_ready_delay_seconds": args.ready_delay,
                "single_process_ready": True,
            })
        time.sleep(60)
        return 0
    if args.worker:
        if args.bootstrapped_path is not None:
            atomic_ready(args.bootstrapped_path, {"status": "BOOTSTRAPPED", "bootstrapped": True, "child_pid": os.getpid(), "phase": "python_started"})
        grandchild = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(60)"])
        time.sleep(max(0.0, args.ready_delay))
        if not args.never_ready:
            atomic_ready(args.ready_path, {"status": "READY", "ready": True, "child_pid": os.getpid(), "grandchild_pid": grandchild.pid, "escaped_session": args.escape_session, "phase": "workload_ready", "intended_ready_delay_seconds": args.ready_delay})
        time.sleep(60)
        return 0
    worker_command = [sys.executable, str(pathlib.Path(__file__).resolve()), str(args.ready_path), "--worker", "--ready-delay", str(args.ready_delay)]
    if args.bootstrapped_path is not None: worker_command.extend(["--bootstrapped-path", str(args.bootstrapped_path)])
    if args.escape_session: worker_command.append("--escape-session")
    if args.never_ready: worker_command.append("--never-ready")
    kwargs: dict[str, object] = {}
    if args.escape_session:
        if os.name == "nt": kwargs["creationflags"] = getattr(subprocess, "CREATE_NEW_PROCESS_GROUP", 0)
        else: kwargs["start_new_session"] = True
    subprocess.Popen(worker_command, **kwargs)
    time.sleep(60)
    return 0


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