import argparse
import glob
import os
import subprocess
import sys
import time


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--dir", required=True, help="Folder containing *.yaml")
    ap.add_argument("--mode", choices=["tpe", "grid"], required=True)
    ap.add_argument("--pattern", default="*.yaml", help="Glob pattern (default: *.yaml)")
    ap.add_argument("--max_files", type=int, default=None, help="Run only the first N files")
    args = ap.parse_args()

    pattern = os.path.join(args.dir, args.pattern)
    files = sorted(glob.glob(pattern))
    if not files:
        print(f"[RUN-FOLDER] No YAML files found in: {pattern}")
        sys.exit(1)

    if args.max_files is not None:
        files = files[: args.max_files]

    print(f"[RUN-FOLDER] Mode={args.mode}  Files={len(files)}  Pattern={pattern}", flush=True)

    for i, fpath in enumerate(files, 1):
        cmd = [
            sys.executable,
            "-u",  # unbuffered output from child
            "-m",
        ]
        if args.mode == "tpe":
            cmd += ["lkb.scripts.run_tpe", "--tpe", fpath]
        else:
            cmd += ["lkb.scripts.run_sweep", "--sweep", fpath]

        print(f"\n[RUN-FOLDER] ({i}/{len(files)}) START: {fpath}", flush=True)
        t0 = time.time()

        proc = subprocess.Popen(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            universal_newlines=True,
            bufsize=1,
        )
        try:
            for line in proc.stdout:
                print(line, end="")
        except KeyboardInterrupt:
            proc.terminate()
            raise

        ret = proc.wait()
        dt = time.time() - t0
        if ret != 0:
            print(f"[RUN-FOLDER] ({i}/{len(files)}) FAIL ({dt:.1f}s): {fpath} (exit={ret})", flush=True)
        else:
            print(f"[RUN-FOLDER] ({i}/{len(files)}) DONE ({dt:.1f}s): {fpath}", flush=True)


if __name__ == "__main__":
    main()
