#!/usr/bin/env python3
"""Merge SageMaker split-fanout webdatasetize outputs.

Expected shard naming:
  split_all_0.json ... split_all_<num_splits-1>.json
  stats_all_0.txt ... stats_all_<num_splits-1>.txt

The prefix can be a local directory or an S3 prefix.
"""

import argparse
import json
import os
import subprocess
import tempfile
from pathlib import Path


def _run(cmd):
    subprocess.run(cmd, check=True)


def _copy(src, dst):
    if src.startswith("s3://") or dst.startswith("s3://"):
        _run(["aws", "s3", "cp", src, dst, "--quiet"])
    else:
        Path(dst).parent.mkdir(parents=True, exist_ok=True)
        _run(["cp", src, dst])


def _remove(path):
    if path.startswith("s3://"):
        _run(["aws", "s3", "rm", path, "--quiet"])
    else:
        os.remove(path)


def _load_json(path):
    with open(path, "r") as f:
        return json.load(f)


def _merge(paths):
    merged = None
    all_sequences = {}

    for path in paths:
        data = _load_json(path)
        seqs = data.get("sequences", {})
        if merged is None:
            merged = data
            merged["sequences"] = {}

        for seq, tars in seqs.items():
            all_sequences.setdefault(seq, []).extend(tars)

    dedup_sequences = {}
    for seq, tars in sorted(all_sequences.items()):
        seen = set()
        dedup = []
        for tar in tars:
            if tar in seen:
                continue
            seen.add(tar)
            dedup.append(tar)
        dedup_sequences[seq] = dedup

    merged["sequences"] = dedup_sequences
    merged.setdefault("info", {})
    merged["info"]["num_sequences"] = len(dedup_sequences)
    merged["info"]["num_files"] = sum(len(v) for v in dedup_sequences.values())
    merged["info"]["source_manifests"] = [os.path.basename(p) for p in paths]
    return merged


def _merge_stats(paths):
    chunks = []
    for path in paths:
        with open(path, "r") as f:
            text = f.read().rstrip()
        if text:
            chunks.append(text)
    return "\n\n".join(chunks) + ("\n" if chunks else "")


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("prefix", help="Local directory or S3 prefix containing split_all_<i>.json")
    parser.add_argument("--num_splits", type=int, required=True)
    parser.add_argument("--out_name", default="split_all.json")
    parser.add_argument("--stats_out_name", default="stats_all.txt")
    parser.add_argument("--delete_shards", action="store_true")
    return parser.parse_args()


def main():
    args = parse_args()
    prefix = args.prefix.rstrip("/")

    with tempfile.TemporaryDirectory() as tmpdir:
        local_paths = []
        src_paths = []
        local_stats_paths = []
        src_stats_paths = []
        for i in range(args.num_splits):
            name = f"split_all_{i}.json"
            src = f"{prefix}/{name}"
            local = os.path.join(tmpdir, name)
            _copy(src, local)
            src_paths.append(src)
            local_paths.append(local)

            stats_name = f"stats_all_{i}.txt"
            stats_src = f"{prefix}/{stats_name}"
            stats_local = os.path.join(tmpdir, stats_name)
            _copy(stats_src, stats_local)
            src_stats_paths.append(stats_src)
            local_stats_paths.append(stats_local)

        merged = _merge(local_paths)
        local_out = os.path.join(tmpdir, args.out_name)
        with open(local_out, "w") as f:
            json.dump(merged, f, indent=4)

        dst = f"{prefix}/{args.out_name}"
        _copy(local_out, dst)
        print(f"Merged {len(local_paths)} split outputs -> {dst}")

        merged_stats = _merge_stats(local_stats_paths)
        local_stats_out = os.path.join(tmpdir, args.stats_out_name)
        with open(local_stats_out, "w") as f:
            f.write(merged_stats)

        stats_dst = f"{prefix}/{args.stats_out_name}"
        _copy(local_stats_out, stats_dst)
        print(f"Merged {len(local_stats_paths)} stats outputs -> {stats_dst}")

        if args.delete_shards:
            for src in src_paths:
                _remove(src)
            for src in src_stats_paths:
                _remove(src)
            print(f"Deleted {len(src_paths)} shard split files")
            print(f"Deleted {len(src_stats_paths)} shard stats files")


if __name__ == "__main__":
    main()