#!/usr/bin/env python3
"""Stream-restore 1/3 of Waymo v1.4.3 training: each worker downloads a tfrecord from
GCS (ADC token), converts to JPEG@640 unified clips, deletes the tfrecord. Disk-frugal."""
import argparse, os, sys, time, traceback
sys.path.insert(0, "/tmp/wproto")

from mapgs.config import load_config
from mapgs.data.convert.waymo_to_unified import convert_waymo
import scripts.gcs_waymo as gcs


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--workers", type=int, default=4)
    ap.add_argument("--worker-id", type=int, default=0)
    ap.add_argument("--n-seg", type=int, default=266)      # 1/3 of 798
    ap.add_argument("--longside", type=int, default=640)
    ap.add_argument("--out", default="/mnt/william/data/unified/waymo")
    args = ap.parse_args()

    cfg = load_config()
    tok = gcs.access_token(); t_tok = time.time()
    segs = gcs.list_training_tfrecords(tok)[: args.n_seg]
    mine = segs[args.worker_id :: args.workers]
    tmp = f"/mnt/william/_waymo_stream/w{args.worker_id}"
    os.makedirs(tmp, exist_ok=True)
    print(f"[w{args.worker_id}] {len(mine)} segments", flush=True)

    done = 0
    for name in mine:
        base = os.path.basename(name)
        sid = base.split("segment-")[-1][:12]
        # skip if already converted (idempotent restart)
        if os.path.isdir(os.path.join(args.out, "train")) and \
           any(d.startswith(f"train_{sid}") for d in os.listdir(os.path.join(args.out, "train"))):
            done += 1; continue
        if time.time() - t_tok > 2400:                      # refresh token every 40 min
            tok = gcs.access_token(); t_tok = time.time()
        dst = os.path.join(tmp, base)
        try:
            gcs.download(tok, name, dst)
            convert_waymo(tmp, args.out, cfg, only_segments=[base],
                          store_longside=args.longside, image_ext="jpg")
        except Exception:
            print(f"[w{args.worker_id}] FAIL {base}", flush=True); traceback.print_exc()
        finally:
            if os.path.exists(dst):
                os.remove(dst)
        done += 1
        if done % 5 == 0:
            print(f"[w{args.worker_id}] {done}/{len(mine)}", flush=True)
    print(f"[w{args.worker_id}] DONE {done}/{len(mine)}", flush=True)


if __name__ == "__main__":
    main()