#!/usr/bin/env python3
"""Stream Argoverse 2 Sensor -> unified clips, conserving disk.

For each log: selectively download (3 front cameras + lidar timestamps + map +
calibration + annotations + poses, ~0.5 GB vs 1.1 GB full), convert to unified,
then delete the raw log. Resumable (tracks processed log ids) and disk-guarded.

    python scripts/av2_pipeline.py --split train --n-logs 100 \
        --out data/unified/av2 --store-longside 1024 --min-free-gb 15
"""

import argparse
import os
import shutil
import subprocess
import time

from mapgs.config import load_config
from mapgs.data.convert import convert_argoverse2

BUCKET = "s3://argoverse/datasets/av2/sensor"
NEEDED_DIRS = ["calibration", "map", "sensors/lidar"]
FRONT = ["ring_front_left", "ring_front_center", "ring_front_right"]
FILES = ["annotations.feather", "city_SE3_egovehicle.feather"]


def s5(s5cmd, *args):
    subprocess.run([s5cmd, "--no-sign-request", *args], check=False,
                   stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)


def list_logs(s5cmd, split):
    out = subprocess.run([s5cmd, "--no-sign-request", "ls", f"{BUCKET}/{split}/"],
                         capture_output=True, text=True).stdout
    return [ln.split()[-1].rstrip("/") for ln in out.splitlines() if "DIR" in ln]


def download_log(s5cmd, split, lid, raw_root):
    base = f"{BUCKET}/{split}/{lid}"
    dst = os.path.join(raw_root, split, lid)
    for d in NEEDED_DIRS:
        os.makedirs(os.path.join(dst, d), exist_ok=True)
        s5(s5cmd, "cp", f"{base}/{d}/*", os.path.join(dst, d) + "/")
    for c in FRONT:
        os.makedirs(os.path.join(dst, "sensors/cameras", c), exist_ok=True)
        s5(s5cmd, "cp", f"{base}/sensors/cameras/{c}/*", os.path.join(dst, "sensors/cameras", c) + "/")
    for f in FILES:
        s5(s5cmd, "cp", f"{base}/{f}", dst + "/")
    return dst


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=100)
    ap.add_argument("--out", default="data/unified/av2")
    ap.add_argument("--raw-tmp", default="/mnt/william/_av2_stream")
    ap.add_argument("--s5cmd", default="/tmp/s5cmd")
    ap.add_argument("--store-longside", type=int, default=1024)
    ap.add_argument("--max-clips-per-log", type=int, default=None)
    ap.add_argument("--min-free-gb", type=float, default=15.0)
    ap.add_argument("--keep-raw", action="store_true")
    ap.add_argument("--shard", type=int, default=0, help="this worker's shard index")
    ap.add_argument("--num-shards", type=int, default=1, help="total parallel workers")
    args = ap.parse_args()

    cfg = load_config(args.config, args.override)
    args.raw_tmp = os.path.join(args.raw_tmp, f"shard{args.shard}")  # isolate each worker's scratch
    out_split = os.path.join(args.out, args.split)
    os.makedirs(out_split, exist_ok=True)
    # per-shard processed file (avoids concurrent-append races between workers)
    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()

    all_logs = list_logs(args.s5cmd, args.split)
    shard_logs = all_logs[args.shard::args.num_shards]  # disjoint slice for this worker
    logs = [l for l in shard_logs if l not in processed]
    print(f"[shard {args.shard}/{args.num_shards}] {len(logs)} logs to process "
          f"(of {len(shard_logs)} in shard, {len(processed)} done); target {args.n_logs}")

    done = 0
    for lid in logs:
        if done >= args.n_logs:
            break
        free_gb = shutil.disk_usage(args.out).free / 1e9
        if free_gb < args.min_free_gb:
            print(f"stopping: only {free_gb:.1f} GB free (< {args.min_free_gb})")
            break
        t = time.time()
        try:
            download_log(args.s5cmd, args.split, lid, args.raw_tmp)
            convert_argoverse2(args.raw_tmp, args.out, cfg, split=args.split, only_logs=[lid],
                               store_longside=args.store_longside, max_clips_per_log=args.max_clips_per_log)
        except Exception as e:
            print(f"[skip {lid[:8]}] {type(e).__name__}: {e}")
        finally:
            if not args.keep_raw:
                shutil.rmtree(os.path.join(args.raw_tmp, args.split, lid), ignore_errors=True)
        with open(processed_path, "a") as fh:
            fh.write(lid + "\n")
        done += 1
        n_clips = len([d for d in os.listdir(out_split) if d.startswith(args.split)])
        print(f"[{done}/{args.n_logs}] {lid[:8]} done in {time.time()-t:.0f}s | "
              f"{n_clips} clips total | {shutil.disk_usage(args.out).free/1e9:.1f} GB free")
    print(f"pipeline finished: {done} logs processed")


if __name__ == "__main__":
    main()