import argparse
import json
import os
import shutil
import subprocess
import time
from pathlib import Path
from glob import glob

REPO_ROOT = Path(__file__).resolve().parents[2]


def _bool(v):
    if isinstance(v, bool):
        return v
    return str(v).lower() in ("1", "true", "yes", "y", "on")


def _normalize_local_root(path: str):
    path = path.rstrip("/")
    for suffix in ("cv_unified_debug", "cv_unified", "cv_webbed_debug", "cv_webbed"):
        if path.endswith(f"/{suffix}"):
            return path[: -(len(suffix) + 1)]
    return path


def _normalize_mode(value, *, choices, name):
    value = str(value).strip().lower()
    if value not in choices:
        raise ValueError(f"{name} must be one of {sorted(choices)}, got {value!r}")
    return value


def _float_list(value):
    if value is None:
        return None
    if isinstance(value, (list, tuple)) and len(value) == 1 and isinstance(value[0], str):
        value = value[0]
    if isinstance(value, (list, tuple)):
        values = value
    else:
        raw = str(value).strip()
        if raw.startswith("[") and raw.endswith("]"):
            raw = raw[1:-1]
        raw = raw.replace(",", " ")
        values = raw.split()
    return [float(v) for v in values]


def _str_list(value):
    if value is None:
        return None
    if isinstance(value, str):
        raw = value.strip()
        if raw.startswith("[") and raw.endswith("]"):
            raw = raw[1:-1]
        raw = raw.replace(",", " ")
        return [v for v in raw.split() if v]
    if isinstance(value, (list, tuple)):
        values = []
        for item in value:
            if isinstance(item, str):
                values.extend(_str_list(item))
            else:
                values.append(str(item))
        return values
    return [str(value)]


def parse_args():
    parser = argparse.ArgumentParser()

    # Existing SM config style.
    parser.add_argument("--datasets", type=str, nargs="+", default=None)
    parser.add_argument("--dataset", type=str, default=None)
    parser.add_argument("--split_json", type=str, default=None)

    # web_local.py expects positional "<dataset> <store>", where store looks like "17b4".
    parser.add_argument("--store", type=str, default=None)
    parser.add_argument("--subset", type=str, default=None)
    parser.add_argument("--num_procs", type=int, default=16)
    parser.add_argument("--max_threads", type=int, default=None)
    parser.add_argument("--mode", type=str, default="frames", choices=["frames", "videos", "latents"])
    parser.add_argument("--from_mode", type=str, default=None, choices=["frames", "videos"])
    parser.add_argument("--resize", type=str, nargs="+", default=None)
    parser.add_argument("--labels", type=str, nargs="+", default=None)
    parser.add_argument("--web_only", type=_bool, default=False)
    parser.add_argument("--depth", type=_bool, default=False)
    parser.add_argument("--depth_bits", type=int, default=None)
    parser.add_argument("--camera_mode", type=str, default=None, choices=["strict", "flexible"])
    parser.add_argument("--snippet_formula", type=float, default=None)
    parser.add_argument("--metadata_shared", type=str, default=None)
    parser.add_argument("--name_suffix", type=str, default=None)
    parser.add_argument("--quiet", type=_bool, default=False)
    parser.add_argument("--official", type=_bool, default=False)
    parser.add_argument("--delete", type=_bool, default=False)
    parser.add_argument("--restart", type=_bool, default=False)
    parser.add_argument("--per_camera", type=_bool, default=False)

    # Output paths for web_simplified.
    parser.add_argument("--local_path", type=str, default="/tmp")
    parser.add_argument("--s3_path", type=str, default="s3://tri-ml-sandbox-16011-us-west-2-datasets")
    parser.add_argument("--folder", type=str, default="cv_webbed")
    parser.add_argument("--upload", type=_bool, default=True)

    # Optional pre-download on SM instances.
    parser.add_argument("--s3_download_data", type=str, default=None)
    parser.add_argument("--s3_download_data_folder", type=str, default=None)
    parser.add_argument("--unified_base", type=str, default=None)
    # Backward compatibility keys.
    parser.add_argument("--download_data", type=str, default=None)
    parser.add_argument("--download_local_bucket", type=str, default="/tmp")
    parser.add_argument("--download_s3_bucket", type=str, default="s3://tri-ml-sandbox-16011-us-west-2-datasets")
    parser.add_argument("--download_num_procs", type=int, default=16)
    parser.add_argument("--input_split", type=str, default="split_all.json")
    parser.add_argument("--split_output_index", type=int, default=None)

    # Accepted for old code compatibility; not used by web_local now.
    parser.add_argument("--parallelize", type=str, default="sequence", choices=["task", "sequence"])
    parser.add_argument("--merge", type=_bool, default=True)
    parser.add_argument("--s5cmd", type=_bool, default=False)

    args = parser.parse_args()

    args.mode = _normalize_mode(args.mode, choices={"frames", "videos", "latents"}, name="mode")
    if args.from_mode is None:
        args.from_mode = args.mode if args.mode in {"frames", "videos"} else "frames"
    args.from_mode = _normalize_mode(args.from_mode, choices={"frames", "videos"}, name="from_mode")
    args.resize = _float_list(args.resize)
    args.labels = _str_list(args.labels)

    # Resolve roots/folders from a single official/debug decision.
    args.local_path = _normalize_local_root(args.local_path)
    # Use folder exactly as provided in config; web_local.py applies official/debug suffixing.
    args.folder = args.folder.rstrip("/")

    # Match web_local.py, which reads unified input from cv_unified/<from_mode>.
    if args.s3_download_data_folder is None:
        args.s3_download_data_folder = f"cv_unified/{args.from_mode}"
    else:
        args.s3_download_data_folder = args.s3_download_data_folder.rstrip("/")

    return args


def resolve_dataset(args):
    dataset = args.dataset
    if dataset is None and args.datasets:
        dataset = args.datasets[0]
    if dataset is None:
        raise ValueError("Missing dataset. Provide --dataset or --datasets.")

    # Keep command compatibility: optional split_json overrides split name.
    if args.split_json:
        split = Path(args.split_json).name.replace(".json", "")
        if "/" in dataset:
            name = dataset.split("/")[0]
            dataset = f"{name}/{split}"
        else:
            dataset = f"{dataset}/{split}"
    return dataset

import yaml
from contextlib import contextmanager

@contextmanager
def remap_cfg_paths_to_tmp(dataset):
    ### renaming dataset paths from /data/... to /tmp/... for SM compatibility to avoid uploading to s3 again after job finishing.
    dataset_name = dataset.split("/")[0]
    cfg_path = REPO_ROOT / f"anydata/webdataset/configs/{dataset_name}.yaml"
    if not cfg_path.exists():
        print(f"Config YAML not found, skipping path remap: {cfg_path}")
        yield
        return
    original = cfg_path.read_text()

    cfg = yaml.safe_load(original)
    changed = False
    for k, v in cfg.items():
        if isinstance(v, dict) and "path" in v:
            new_paths = []
            for p in v["path"]:
                if isinstance(p, str):
                    p2 = p.replace("/data/cv_unified_debug", "/tmp/cv_unified_debug")
                    p2 = p2.replace("/data/cv_unified", "/tmp/cv_unified")
                    new_paths.append(p2)
                    changed = changed or (p2 != p)
                else:
                    new_paths.append(p)
            v["path"] = new_paths

    if changed:
        cfg_path.write_text(yaml.safe_dump(cfg, sort_keys=False))
    try:
        yield
    finally:
        if changed:
            cfg_path.write_text(original)

def ensure_unified_debug_alias(local_bucket, data_folder="cv_unified"):
    """Ensure <local_bucket>/<data_folder>_debug points to <local_bucket>/<data_folder>."""
    root = local_bucket.rstrip("/")
    src = os.path.join(root, data_folder)
    alias = os.path.join(root, f"{data_folder}_debug")
    if os.path.exists(alias) or os.path.islink(alias):
        return
    os.makedirs(os.path.dirname(alias), exist_ok=True)
    os.symlink(src, alias)
    print(f"Created alias: {alias} -> {src}")


def data_download(args):
    download_data = args.s3_download_data if args.s3_download_data is not None else args.download_data
    download_data_folder = args.s3_download_data_folder
    if not download_data:
        return
    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
    local_world_size = int(os.environ.get("LOCAL_WORLD_SIZE", "1"))

    split_rel = download_data
    if not split_rel.endswith(".json"):
        split_rel = f"{split_rel.rstrip('/')}/{args.input_split}"
    split_basename = os.path.basename(split_rel)
    done_marker = os.path.join(
        args.download_local_bucket,
        download_data_folder,
        split_rel + ".download_done",
    )
    local_split = os.path.join(
        args.download_local_bucket,
        download_data_folder,
        split_rel,
    )
    expected_split = os.path.join(
        os.path.dirname(local_split),
        args.input_split,
    )
    dataset_spec = args.dataset if args.dataset is not None else (args.datasets[0] if args.datasets else None)
    expected_dataset_split = None
    if dataset_spec and "/" in dataset_spec and dataset_spec.endswith(".json"):
        expected_dataset_split = os.path.join(
            os.path.dirname(local_split),
            os.path.basename(dataset_spec),
        )

    download_base_folder = download_data_folder
    if os.path.basename(download_base_folder.rstrip("/")) == args.from_mode:
        download_base_folder = os.path.dirname(download_base_folder.rstrip("/"))

    cmd = [
        "python3",
        "anydata/sync/download_unified.py",
        split_rel,
        f"--{args.from_mode}",
        "--num_procs",
        str(args.download_num_procs),
        "--s3_bucket",
        args.download_s3_bucket,
        "--data_folder",
        download_base_folder,
        "--local_bucket",
        args.download_local_bucket,
    ]
    env = os.environ.copy()
    env["PYTHONPATH"] = f"{REPO_ROOT}:{env.get('PYTHONPATH', '')}"

    # One downloader per node (LOCAL_RANK=0), other local ranks wait.
    if local_world_size == 1 or local_rank == 0:
        print("Running download:", " ".join(cmd))
        subprocess.run(cmd, check=True, cwd=str(REPO_ROOT), env=env)
        # If user passed a custom split filename, alias it to input_split
        # so dataset configs that expect split_all.json keep working.
        if split_basename != args.input_split and os.path.exists(local_split):
            os.makedirs(os.path.dirname(expected_split), exist_ok=True)
            shutil.copy2(local_split, expected_split)
            print(f"Created local split alias: {expected_split} <- {local_split}")
        if expected_dataset_split and expected_dataset_split != local_split and os.path.exists(local_split):
            os.makedirs(os.path.dirname(expected_dataset_split), exist_ok=True)
            shutil.copy2(local_split, expected_dataset_split)
            print(f"Created dataset split alias: {expected_dataset_split} <- {local_split}")
        ensure_unified_debug_alias(args.download_local_bucket, download_data_folder)
        Path(done_marker).parent.mkdir(parents=True, exist_ok=True)
        Path(done_marker).write_text("done\n")
    else:
        waited = 0
        print(f"LOCAL_RANK={local_rank} waiting for download marker: {done_marker}")
        while not os.path.exists(done_marker):
            time.sleep(5)
            waited += 5
            if waited >= 240 * 60 * 60:
                raise TimeoutError(f"Timed out waiting for download marker: {done_marker}")
        if split_basename != args.input_split and os.path.exists(local_split) and not os.path.exists(expected_split):
            os.makedirs(os.path.dirname(expected_split), exist_ok=True)
            shutil.copy2(local_split, expected_split)
        if expected_dataset_split and os.path.exists(local_split) and not os.path.exists(expected_dataset_split):
            os.makedirs(os.path.dirname(expected_dataset_split), exist_ok=True)
            shutil.copy2(local_split, expected_dataset_split)
        ensure_unified_debug_alias(args.download_local_bucket, download_data_folder)


def resolve_unified_base(args):
    if args.unified_base:
        return args.unified_base.rstrip("/")
    download_data = args.s3_download_data if args.s3_download_data is not None else args.download_data
    if not download_data:
        return None
    return os.path.join(args.download_local_bucket.rstrip("/"), args.s3_download_data_folder.rstrip("/"))


def rank_adjusted_subset(subset):
    world_size = int(os.environ.get("WORLD_SIZE", "1"))
    rank = int(os.environ.get("RANK", "0"))
    if world_size <= 1:
        return subset
    base = subset if subset else "1/0"
    if "/" in str(base):
        a, b = [int(v) for v in str(base).split("/")]
    else:
        a, b = 1, 0
    return f"{a * world_size}/{b + a * rank}"


def run_webbing(args, dataset):
    world_size = int(os.environ.get("WORLD_SIZE", "1"))
    subset = rank_adjusted_subset(args.subset)
    cmd = [
        "python3",
        "anydata/webdataset/web_local.py",
        dataset,
        args.store,
        "--num_procs",
        str(args.num_procs),
        "--local_path",
        args.local_path,
        "--s3_path",
        args.s3_path,
        "--folder",
        args.folder,
        f"--{args.mode}",
        f"--from_{args.from_mode}",
    ]
    unified_base = resolve_unified_base(args)
    if unified_base is not None:
        cmd.extend(["--unified_base", unified_base])
    if args.max_threads is not None:
        cmd.extend(["--max_threads", str(args.max_threads)])
    if args.resize is not None:
        cmd.append("--resize")
        cmd.extend(str(v) for v in args.resize)
    if args.labels is not None:
        cmd.append("--labels")
        cmd.extend(args.labels)
    if args.web_only:
        cmd.append("--web_only")
    if args.depth:
        cmd.append("--depth")
    if args.depth_bits is not None:
        cmd.extend(["--depth_bits", str(args.depth_bits)])
    if args.camera_mode is not None:
        cmd.extend(["--camera_mode", args.camera_mode])
    if args.snippet_formula is not None:
        cmd.extend(["--snippet_formula", str(args.snippet_formula)])
    if args.metadata_shared is not None:
        cmd.extend(["--metadata_shared", args.metadata_shared])
    if args.name_suffix is not None:
        cmd.extend(["--name_suffix", args.name_suffix])
    if args.official:
        cmd.append("--official")
    if args.delete:
        cmd.append("--delete")
    if args.restart:
        cmd.append("--restart")
    if args.per_camera:
        cmd.append("--per_camera")
    if args.quiet:
        cmd.append("--quiet")
    if subset is not None and str(subset) != "":
        cmd.extend(["--subset", subset])
    if args.upload:
        cmd.append("--upload")

    print("Running webbing:", " ".join(cmd))
    env = os.environ.copy()
    env.pop("LD_PRELOAD", None)
    env["PYTHONPATH"] = f"{REPO_ROOT}:{env.get('PYTHONPATH', '')}"
    subprocess.run(cmd, check=True, cwd=str(REPO_ROOT), env=env)


def split_output_name(args):
    if args.split_output_index is None:
        return "split_all.json"
    return f"split_all_{args.split_output_index}.json"


def stats_output_name(args):
    if args.split_output_index is None:
        return "stats_all.txt"
    return f"stats_all_{args.split_output_index}.txt"


def dataset_split_name(dataset):
    """Return the expected combined manifest name for the dataset spec."""
    if dataset.endswith(".json"):
        return os.path.basename(dataset)
    return "split_all.json"


def dataset_split_prefix(dataset):
    return os.path.splitext(dataset_split_name(dataset))[0]


def output_base_paths(args, dataset):
    dataset_name = dataset.split("/")[0]
    folder = args.folder if args.official else f"{args.folder}_debug"
    local_base = os.path.join(args.local_path, folder, args.mode)
    s3_base = f"{args.s3_path.rstrip('/')}/{folder}/{args.mode}"
    dataset_base = os.path.join(local_base, dataset_name)
    return local_base, s3_base, dataset_base


def promote_single_rank_split(args, dataset):
    """Rename single-rank fanout split/stat outputs to index-qualified shard names."""
    world_size = int(os.environ.get("WORLD_SIZE", "1"))
    if world_size > 1 or args.split_output_index is None:
        return

    local_base, s3_base, dataset_base = output_base_paths(args, dataset)
    expected_name = dataset_split_name(dataset)
    split_paths = sorted(
        p for p in glob(f"{dataset_base}/**/{expected_name}", recursive=True)
        if f"/tmp_{expected_name}" not in p
    )
    if not split_paths and expected_name != "split_all.json":
        split_paths = sorted(
            p for p in glob(f"{dataset_base}/**/split_all.json", recursive=True)
            if "/tmp_split_all" not in p
        )
    if not split_paths:
        print(f"No local {expected_name} found under {dataset_base}; skipping split promotion.")
        return

    for src in split_paths:
        dst = os.path.join(os.path.dirname(src), split_output_name(args))
        shutil.copy2(src, dst)
        os.remove(src)
        print(f"Promoted split output: {src} -> {dst}")

        if args.upload:
            src_s3 = src.replace(local_base, s3_base)
            dst_s3 = dst.replace(local_base, s3_base)
            subprocess.run(["aws", "s3", "cp", dst, dst_s3, "--quiet"], check=True)
            subprocess.run(["aws", "s3", "rm", src_s3, "--quiet"], check=False)
            print(f"Uploaded split output: {dst_s3}")

        stats_src = os.path.join(os.path.dirname(src), "stats_all.txt")
        if not os.path.exists(stats_src):
            print(f"No local stats_all.txt found beside {dst}; skipping stats promotion.")
            continue

        stats_dst = os.path.join(os.path.dirname(stats_src), stats_output_name(args))
        shutil.copy2(stats_src, stats_dst)
        os.remove(stats_src)
        print(f"Promoted stats output: {stats_src} -> {stats_dst}")

        if args.upload:
            stats_src_s3 = stats_src.replace(local_base, s3_base)
            stats_dst_s3 = stats_dst.replace(local_base, s3_base)
            subprocess.run(["aws", "s3", "cp", stats_dst, stats_dst_s3, "--quiet"], check=True)
            subprocess.run(["aws", "s3", "rm", stats_src_s3, "--quiet"], check=False)
            print(f"Uploaded stats output: {stats_dst_s3}")


def merge_rank_splits(args, dataset, timeout_s=2 * 60 * 60):
    """Merge per-rank split json files into a single consolidated split on rank 0."""
    world_size = int(os.environ.get("WORLD_SIZE", "1"))
    rank = int(os.environ.get("RANK", "0"))
    if world_size <= 1 or rank != 0:
        return

    local_base, s3_base, base_dir = output_base_paths(args, dataset)
    split_prefix = dataset_split_prefix(dataset)

    start = time.time()
    shard_paths = []
    while True:
        shard_paths = sorted(
            p for p in glob(f"{base_dir}/**/{split_prefix}_{world_size}_*.json", recursive=True)
            if f"/tmp_{split_prefix}_" not in p
        )
        if not shard_paths and split_prefix != "split_all":
            shard_paths = sorted(
                p for p in glob(f"{base_dir}/**/split_all_{world_size}_*.json", recursive=True)
                if "/tmp_split_all_" not in p
            )
        if len(shard_paths) >= world_size:
            break
        if time.time() - start > timeout_s:
            raise TimeoutError(
                f"Timed out waiting for shard splits under {base_dir}. "
                f"Found {len(shard_paths)}/{world_size}."
            )
        time.sleep(5)

    out_name = split_output_name(args)
    print(f"Merging {len(shard_paths)} shard splits into consolidated {out_name}")
    merged = None
    all_sequences = {}

    for sp in shard_paths:
        with open(sp, "r") as f:
            data = json.load(f)
        seqs = data.get("sequences", {})
        if merged is None:
            merged = data
            merged["sequences"] = {}
        for key, vals in seqs.items():
            if key not in all_sequences:
                all_sequences[key] = []
            all_sequences[key].extend(vals)

    # Deduplicate tar names per sequence while preserving order.
    dedup_sequences = {}
    for key, vals in all_sequences.items():
        seen = set()
        uniq = []
        for v in vals:
            if v in seen:
                continue
            seen.add(v)
            uniq.append(v)
        dedup_sequences[key] = uniq

    merged["sequences"] = dedup_sequences
    if "info" not in merged:
        merged["info"] = {}
    merged["info"]["num_sequences"] = len(dedup_sequences)
    merged["info"]["num_files"] = sum(len(v) for v in dedup_sequences.values())

    out_path = os.path.join(os.path.dirname(shard_paths[0]), out_name)
    with open(out_path, "w") as f:
        json.dump(merged, f, indent=4)
    print(f"Wrote consolidated split: {out_path}")

    if args.upload:
        out_s3 = out_path.replace(local_base, s3_base)
        subprocess.run(["aws", "s3", "cp", out_path, out_s3, "--quiet"], check=True)
        print(f"Uploaded consolidated split: {out_s3}")

        for sp in shard_paths:
            shard_s3 = sp.replace(local_base, s3_base)
            if shard_s3 != out_s3:
                subprocess.run(["aws", "s3", "rm", shard_s3, "--quiet"], check=False)

    for sp in shard_paths:
        if sp != out_path and os.path.exists(sp):
            os.remove(sp)


def main():
    args = parse_args()

    dataset = resolve_dataset(args)

    print(f"SAGEMAKER={os.environ.get('SAGEMAKER', '')}")
    print(f"Resolved dataset: {dataset}")
    print(f"Upload enabled: {args.upload}")

    data_download(args)
    with remap_cfg_paths_to_tmp(dataset):
        run_webbing(args, dataset)
    promote_single_rank_split(args, dataset)
    merge_rank_splits(args, dataset)


if __name__ == "__main__":
    main()