from __future__ import annotations
import argparse
import os
import subprocess
import sys


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--checkpoint", default="/root/chess/ckpt/model.pt")
    ap.add_argument("--out-dir", default="/root/chess/selfplay_data")
    ap.add_argument("--games", type=int, default=4000, help="total games across all workers")
    ap.add_argument("--workers", type=int, default=28)
    ap.add_argument("--movetime", type=int, default=200)
    ap.add_argument("--temperature-plies", type=int, default=10)
    args = ap.parse_args()
    per_worker = max(1, args.games // args.workers)
    procs = []
    logs = []
    for i in range(args.workers):
        worker_out = os.path.join(args.out_dir, f"worker{i}")
        log_path = os.path.join(args.out_dir, f"worker{i}.log")
        os.makedirs(args.out_dir, exist_ok=True)
        log_f = open(log_path, "w")
        cmd = [
            sys.executable,
            "selfplay_finetune.py",
            "--checkpoint",
            args.checkpoint,
            "--out-dir",
            worker_out,
            "--games",
            str(per_worker),
            "--movetime",
            str(args.movetime),
            "--temperature-plies",
            str(args.temperature_plies),
        ]
        p = subprocess.Popen(
            cmd,
            stdout=log_f,
            stderr=subprocess.STDOUT,
            cwd=os.path.dirname(os.path.abspath(__file__)) or ".",
        )
        procs.append(p)
        logs.append(log_f)
        print(f"[launcher] started worker {i} (pid {p.pid}), {per_worker} games -> {worker_out}")
    exit_codes = []
    for i, p in enumerate(procs):
        code = p.wait()
        exit_codes.append(code)
        print(f"[launcher] worker {i} exited with code {code}")
    for f in logs:
        f.close()
    failed = [i for i, c in enumerate(exit_codes) if c != 0]
    if failed:
        print(
            f"[launcher] WARNING: workers failed: {failed} -- check worker*.log in {args.out_dir}"
        )
    print(
        f"[launcher] DONE: {args.workers} workers, {per_worker * args.workers} games total requested"
    )


if __name__ == "__main__":
    main()