#!/usr/bin/env python3
"""Stream Waymo v1.4.x (tfrecord) -> unified clips, conserving disk.

Per segment: gcloud-download the ~1 GB tfrecord, convert (TF-free) to unified,
delete the raw file. Sharded + resumable + disk-guarded, like the AV2 pipeline.
Stores JPEG by default to fit more clips on a small disk.

    python scripts/waymo_pipeline.py --split train --num-shards 4 --shard 0 \
        --n-logs 40 --out data/unified/waymo --image-ext jpg --min-free-gb 6
"""

import argparse, os, shutil, subprocess, time
from mapgs.config import load_config
from mapgs.data.convert import convert_waymo

GC = "/mnt/william/google-cloud-sdk/bin/gcloud"
BUCKET = "gs://waymo_open_dataset_v_1_4_3/individual_files"
SPLIT_DIR = {"train": "training", "val": "validation"}


def gs(*args):
    env = dict(os.environ, CLOUDSDK_PYTHON="python3")
    return subprocess.run([GC, "storage", *args], capture_output=True, text=True, env=env)


def list_segments(split):
    out = gs("ls", f"{BUCKET}/{SPLIT_DIR[split]}/*.tfrecord").stdout
    return [ln.strip().split("/")[-1] for ln in out.splitlines() if ln.strip().endswith(".tfrecord")]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--config", default="configs/base.yaml")
    ap.add_argument("--override", nargs="*", default=["data.num_frames=20"])
    ap.add_argument("--split", default="train")
    ap.add_argument("--n-logs", type=int, default=40)
    ap.add_argument("--out", default="data/unified/waymo")
    ap.add_argument("--raw-tmp", default="/mnt/william/_waymo_stream")
    ap.add_argument("--store-longside", type=int, default=1024)
    ap.add_argument("--image-ext", default="jpg")
    ap.add_argument("--max-clips-per-log", type=int, default=None)
    ap.add_argument("--min-free-gb", type=float, default=6.0)
    ap.add_argument("--shard", type=int, default=0)
    ap.add_argument("--num-shards", type=int, default=1)
    args = ap.parse_args()

    cfg = load_config(args.config, args.override)
    raw_dir = os.path.join(args.raw_tmp, f"shard{args.shard}")
    os.makedirs(raw_dir, exist_ok=True)
    out_split = os.path.join(args.out, args.split)
    os.makedirs(out_split, exist_ok=True)
    processed_path = os.path.join(args.out, f"processed_{args.split}_{args.shard}of{args.num_shards}.txt")
    processed = set(open(processed_path).read().split()) if os.path.exists(processed_path) else set()

    segs = list_segments(args.split)[args.shard::args.num_shards]
    segs = [s for s in segs if s not in processed]
    print(f"[shard {args.shard}/{args.num_shards}] {len(segs)} segments to do; target {args.n_logs}", flush=True)

    done = 0
    for seg in segs:
        if done >= args.n_logs:
            break
        if shutil.disk_usage(args.out).free / 1e9 < args.min_free_gb:
            print(f"stopping: low disk ({shutil.disk_usage(args.out).free/1e9:.1f} GB)", flush=True)
            break
        t = time.time()
        try:
            r = gs("cp", f"{BUCKET}/{SPLIT_DIR[args.split]}/{seg}", raw_dir + "/")
            if r.returncode != 0:
                raise RuntimeError(r.stderr.strip()[:200])
            convert_waymo(raw_dir, args.out, cfg, split=args.split, only_segments=[seg],
                          store_longside=args.store_longside, image_ext=args.image_ext,
                          max_clips_per_log=args.max_clips_per_log)
        except Exception as e:
            print(f"[skip {seg[:24]}] {type(e).__name__}: {e}", flush=True)
        finally:
            for f in os.listdir(raw_dir):
                os.remove(os.path.join(raw_dir, f))
        with open(processed_path, "a") as fh:
            fh.write(seg + "\n")
        done += 1
        n = len([d for d in os.listdir(out_split) if d.startswith(args.split)])
        print(f"[{done}/{args.n_logs}] {seg[:24]} {time.time()-t:.0f}s | {n} clips | "
              f"{shutil.disk_usage(args.out).free/1e9:.1f} GB free", flush=True)
    print(f"shard {args.shard} done: {done} segments", flush=True)


if __name__ == "__main__":
    main()