"""Convert the adverse-weather (rain/snow) Unified data to AnyData *Webbed* shards
(videos mode), with two extra always-loaded streams:

  <scene>/<cam>/<st>_<fn>.tar        rgb mp4 window + lowdim(camera marker)   [one per RGB camera]
  <scene>/text/<st>_<fn>.tar         raw caption (lowdim.language)            [one per window]
  <scene>/text_emb/<st>_<fn>.tar     UMT5 embedding (lat_language, from store) [one per window]

Bypasses web_local.py's CLI/dataset_obj (hard-wired to cv_unified/<mode>/<DATASET> and mis-parses
cv_unified_adverse_weather/parquet_all_weather). Reads scene mp4s directly and reuses web_local's
proven writers for byte-compatible tar members. The text streams carry a `lowdim.camera` marker
('text' / 'text_emb') so they survive decode_sample's camera filter and are selected as always-on
aux streams independent of the RGB camera. The compact lat_language stream requires the
BaseWebbed._decode_value patch (decodes 'lat_language', stored once-per-window).

Usage (in container):
  python3 -m custom.dataset.web_adverse_weather --split split_rain_snow --context 32 \
      --out /workspace/data/cv_webbed_adverse_weather --num_procs 16
"""
import argparse
import json
import os
from functools import partial
from multiprocessing import Pool

import numpy as np

from anydata.webdataset.web_local import _video_sample_for_cam, write_tarfile
from anydata.utils.write import create_folder, write_json, write_video_seq  # noqa: F401
from custom.dataset.adverse_weather_captions import AdverseWeatherCaptions, DEFAULT_DATA_ROOT

CAMS7 = ["front_wide", "front_tele", "left_side_forward", "left_side_rearward",
         "rear_tele", "right_side_forward", "right_side_rearward"]
CAMS_PVM = ["front_pvm", "left_side_pvm", "rear_pvm", "right_side_pvm"]
# Canonical global camera ordering. The per-window tar member keys embed the camera's
# index from this list (e.g. front_pvm -> 7). The loader (BaseWebbed.decode_sample)
# uses the in-tar `lowdim.camera` marker as the ground-truth camera identity and remaps
# indices, but distinct indices are required so cameras opened together (multi_tarfiles>1
# or cameras_core_sample>1) don't collide. Keep PVMs at 7..10 to stay disjoint from CAMS7.
CAMS_ALL = CAMS7 + CAMS_PVM
TEXT_IDX = 90
TEXT_EMB_IDX = 91
SCENES_SUBDIR = "cv_unified_adverse_weather/videos/parquet_all_weather"

_CAPS = None  # per-process lazy caption->embedding store


def _caps():
    global _CAPS
    if _CAPS is None:
        _CAPS = AdverseWeatherCaptions()
    return _CAPS


def _read_mp4(path):
    import decord
    return decord.VideoReader(path)[:].asnumpy()  # (T, H, W, 3) uint8


def _suffix_and_key(sample, seq, i):
    length = len(sample.keys()) - 1
    out = {"__key__": f"{seq}_{i}"}
    for k, v in sample.items():
        if not k.startswith("__"):
            k = f"{k[:-4]}.{length}{k[-4:]}"
        out[k] = v
    return out


def _lowdim_marker(cam_name, n):
    return {"camera": np.array([cam_name] * n, dtype=object),
            "timestep": np.arange(n, dtype=np.int64)}


def _build_md(meta, cam_name, H, W, args):
    info = meta.get("info") if isinstance(meta.get("info"), dict) else {}
    return {
        "info": {"name": info.get("name", "adverse_weather"),
                 "tags": info.get("tags", ["driving", "real"]), "storage": "videos"},
        "labels": ["rgb", "language"],
        "cameras": [cam_name],
        "resolution": [int(H), int(W)],
        "framerate": meta.get("framerate", 10),
        "stride": args.stride,
        "length": args.context,
    }


def process_scene(scene_item, args, scenes_root):
    """Web one scene; returns (manifest_key, [tar basenames]) or None. Writes tars to disk."""
    seq_hash, (scene_key, _nframes) = scene_item  # seq_hash = stable int for __key__ uniqueness
    scene_dir = os.path.join(scenes_root, scene_key)
    try:
        meta = json.load(open(os.path.join(scene_dir, "metadata.json")))
    except Exception:
        return None
    caption = meta.get("language", {}).get("prompt")
    # Weather-transfer val: override the clear scene's own caption with a rain/snow transfer
    # prompt, picked deterministically by the scene's split index (seq_hash) so the bake is
    # reproducible and balanced across prompts. Lets the webbed val mirror the local transfer val.
    if getattr(args, "transfer_prompts", None):
        caption = args.transfer_prompts[seq_hash % len(args.transfer_prompts)]
    # In --rgb_only mode the text/text_emb streams already exist on S3; skip caption work.
    if args.rgb_only:
        emb = None
    else:
        emb = _caps().get(caption) if isinstance(caption, str) else None

    # Web only the requested cameras, keyed by their canonical global index (CAMS_ALL).
    frames = {}
    for cam in args.cameras:
        ci = CAMS_ALL.index(cam)
        p = os.path.join(scene_dir, "rgb", f"{cam}.mp4")
        if os.path.exists(p):
            try:
                frames[ci] = _read_mp4(p)
            except Exception:
                pass
    if not frames:
        return None
    T = min(len(v) for v in frames.values())
    out_base = f"{args.out}/{args.name}/{args.split}/{scene_key}"

    tar_names = []
    win = 0
    for st in range(0, T - 1, args.context):
        fn = min(st + args.context, T)
        if fn - st < 2:
            continue
        wb = "%04d_%04d.tar" % (st, fn)
        ci0 = sorted(frames)[0]
        H0, W0 = frames[ci0].shape[1], frames[ci0].shape[2]

        for ci in sorted(frames):
            H, W = frames[ci].shape[1], frames[ci].shape[2]
            consolidated = {"rgb": {ci: frames[ci]}, "lowdim": {ci: _lowdim_marker(CAMS_ALL[ci], T)}}
            sample = _video_sample_for_cam(args, consolidated, ci, st, fn, store_raw=True)
            sample["metadata.npz"] = dict(metadata=np.array(_build_md(meta, CAMS_ALL[ci], H, W, args), dtype=object))
            sample = _suffix_and_key(sample, seq_hash, win)
            fname = f"{out_base}/{CAMS_ALL[ci]}/{wb}"
            create_folder(fname); write_tarfile(args, fname, sample)

        if not args.rgb_only:
            tc = f"({st}:{fn},{TEXT_IDX})"
            tsample = {f"lowdim.{tc}.npz": {**_lowdim_marker("text", fn - st),
                                           "language": np.array([{"prompt": caption}] * (fn - st), dtype=object)}}
            tsample["metadata.npz"] = dict(metadata=np.array(_build_md(meta, "text", H0, W0, args), dtype=object))
            tsample = _suffix_and_key(tsample, seq_hash, win)
            fname = f"{out_base}/text/{wb}"
            create_folder(fname); write_tarfile(args, fname, tsample)

        if not args.rgb_only and emb is not None:
            ec = f"({st}:{fn},{TEXT_EMB_IDX})"
            lat = {"t5_text_embeddings": emb["t5_text_embeddings"][0].float().cpu().numpy(),
                   "t5_text_mask": emb["t5_text_mask"][0].cpu().numpy(),
                   "prompt": np.array(caption, dtype=object)}
            esample = {f"lat_language.{ec}.npz": lat, f"lowdim.{ec}.npz": _lowdim_marker("text_emb", fn - st)}
            esample["metadata.npz"] = dict(metadata=np.array(_build_md(meta, "text_emb", H0, W0, args), dtype=object))
            esample = _suffix_and_key(esample, seq_hash, win)
            fname = f"{out_base}/text_emb/{wb}"
            create_folder(fname); write_tarfile(args, fname, esample)
        tar_names.append(wb)
        win += 1

    if not tar_names:
        return None
    return scene_key, sorted(set(tar_names)), emb is not None


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data_root", default=DEFAULT_DATA_ROOT)
    ap.add_argument("--split", default="split_rain_snow")
    ap.add_argument("--out", required=True)
    ap.add_argument("--context", type=int, default=32)
    ap.add_argument("--stride", type=int, default=1)
    ap.add_argument("--num_scenes", type=int, default=None)
    ap.add_argument("--num_procs", type=int, default=16)
    ap.add_argument("--name", default="adverse_weather")
    ap.add_argument("--cameras", default=None,
                    help="Comma-separated camera names to web (default = the 7 16:9 cams). "
                         "Must be a subset of CAMS_ALL; indices come from CAMS_ALL ordering.")
    ap.add_argument("--rgb_only", action="store_true",
                    help="Skip the text/text_emb streams (they already exist on S3); web only the RGB cams.")
    ap.add_argument("--transfer-prompts", dest="transfer_prompts", default=None,
                    help="JSON list of weather-transfer prompts; if set, each scene's caption is "
                         "overridden with prompts[scene_idx %% len] (for the webbed transfer val).")
    args = ap.parse_args()
    if args.transfer_prompts:
        args.transfer_prompts = json.load(open(args.transfer_prompts))
    args.cameras = [c.strip() for c in args.cameras.split(",")] if args.cameras else list(CAMS7)
    unknown = [c for c in args.cameras if c not in CAMS_ALL]
    if unknown:
        raise ValueError(f"Unknown camera(s) {unknown}; valid: {CAMS_ALL}")
    args.global_cam_idx = None
    args.depth_bits = None
    args.length = args.context
    args.upload = False
    args.delete = False
    args.s3_path = ""
    args.local_path = args.out

    scenes_root = os.path.join(args.data_root, SCENES_SUBDIR)
    split = json.load(open(os.path.join(scenes_root, f"{args.split}.json")))
    seqs = list(split["sequences"].items())
    if args.num_scenes:
        seqs = seqs[: args.num_scenes]
    tasks = list(enumerate(seqs))  # (seq_hash, (scene_key, nframes))

    fn = partial(process_scene, args=args, scenes_root=scenes_root)
    manifest_sequences, n_emb, done = {}, 0, 0
    with Pool(args.num_procs) as pool:
        for res in pool.imap_unordered(fn, tasks, chunksize=1):
            done += 1
            if res is None:
                continue
            scene_key, basenames, has_emb = res
            manifest_sequences[scene_key] = basenames
            n_emb += int(has_emb)
            if done % 100 == 0:
                print(f"[{done}/{len(tasks)}] scenes={len(manifest_sequences)} with_emb={n_emb}", flush=True)

    cameras = list(args.cameras) + ([] if args.rgb_only else ["text", "text_emb"])
    info = {"name": args.name, "tags": ["driving", "real"], "storage": "videos",
            "stride": args.stride, "num_sequences": len(manifest_sequences),
            "num_files": sum(len(v) for v in manifest_sequences.values()),
            "cameras": cameras}
    manifest = {"info": info, "sequences": manifest_sequences}
    man_path = f"{args.out}/{args.name}/{args.split}/manifest.json"
    create_folder(man_path); write_json(man_path, manifest)
    print(f"\nDONE. Manifest: {man_path} ({len(manifest_sequences)} scenes, "
          f"{info['num_files']} tars, {n_emb} with embeddings)")


if __name__ == "__main__":
    main()