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

from anydata.sync.sync_utils import aws_s3_cp


def _load_json(path: str):
    if path.startswith("s3://"):
        with tempfile.NamedTemporaryFile(suffix=".json", delete=False) as fp:
            tmp = fp.name
        aws_s3_cp(path, tmp)
        with open(tmp, "r") as f:
            data = json.load(f)
        os.remove(tmp)
        return data
    with open(path, "r") as f:
        return json.load(f)


def _seq_weight(value):
    # split_all.json usually stores int frame counts per sequence.
    if isinstance(value, int):
        return value
    if isinstance(value, list):
        return len(value)
    if isinstance(value, dict):
        # webbed-style dicts can be nested; fall back to entry count.
        return len(value)
    return 1


def _partition_sequences_contiguous_by_count(sequences: dict, num_splits: int):
    if num_splits <= 0:
        raise ValueError("num_splits must be > 0")

    items = list(sequences.items())
    total = len(items)
    base, remainder = divmod(total, num_splits)

    buckets = []
    loads = []
    offset = 0
    for i in range(num_splits):
        take = base + (1 if i < remainder else 0)
        chunk_items = items[offset : offset + take]
        offset += take
        chunk = dict(chunk_items)
        buckets.append(chunk)
        loads.append(sum(_seq_weight(v) for _, v in chunk_items))

    return buckets, loads


def _updated_size(size_dict, seqs_dict):
    out = dict(size_dict) if isinstance(size_dict, dict) else {}
    src_samples = out.get("samples", 0)
    src_frames = out.get("frames", None)

    new_seqs = len(seqs_dict)
    new_samples = sum(_seq_weight(v) for v in seqs_dict.values())
    out["seqs"] = new_seqs
    out["samples"] = new_samples

    # Keep frames consistent with source metadata when frames/sample ratio exists.
    if isinstance(src_frames, (int, float)) and isinstance(src_samples, (int, float)) and src_samples > 0:
        ratio = float(src_frames) / float(src_samples)
        out["frames"] = int(round(new_samples * ratio))
    return out


def _default_dataset_name(src: str):
    if src.startswith("s3://"):
        # .../cv_unified/<dataset>/split_all.json
        parts = src.rstrip("/").split("/")
        if len(parts) >= 2:
            return parts[-2]
    return Path(src).parent.name


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "src_split_json",
        type=str,
        help="Local path or s3://.../split_all.json",
    )
    parser.add_argument("--dataset", type=str, default=None, help="Dataset folder name (e.g., LBM)")
    parser.add_argument("--num_splits", type=int, default=8)
    parser.add_argument("--out_root", type=str, default="anydata/sagemaker/splits")
    parser.add_argument(
        "--prefix",
        type=str,
        default="sm_split",
        help="Output file prefix. Files become <prefix>_0.json ... <prefix>_<N-1>.json",
    )
    parser.add_argument(
        "--upload_s3_prefix",
        type=str,
        default=None,
        help="Optional S3 prefix to upload generated JSONs, e.g. s3://.../cv_unified/LBM",
    )
    return parser.parse_args()


def main():
    args = parse_args()
    data = _load_json(args.src_split_json)
    if "sequences" not in data or not isinstance(data["sequences"], dict):
        raise ValueError("Invalid split json: missing dict field 'sequences'")

    dataset = args.dataset if args.dataset else _default_dataset_name(args.src_split_json)
    out_dir = Path(args.out_root) / dataset
    out_dir.mkdir(parents=True, exist_ok=True)

    buckets, loads = _partition_sequences_contiguous_by_count(data["sequences"], args.num_splits)

    print(f"Input sequences: {len(data['sequences'])}")
    print(f"Target splits: {args.num_splits}")
    print("Split strategy: contiguous_by_count")
    print(f"Output dir: {out_dir}")

    for i, seqs in enumerate(buckets):
        split_data = dict(data)
        split_data["sequences"] = seqs
        split_data["size"] = _updated_size(data.get("size", {}), seqs)

        out_file = out_dir / f"{args.prefix}_{i}.json"
        with open(out_file, "w") as f:
            json.dump(split_data, f, indent=2)

        print(f"[{i:02d}] {out_file}  seqs={len(seqs)}  approx_load={loads[i]}")

        if args.upload_s3_prefix:
            s3_dst = args.upload_s3_prefix.rstrip("/") + f"/{args.prefix}_{i}.json"
            aws_s3_cp(str(out_file), s3_dst)
            print(f"     uploaded -> {s3_dst}")
            # Uploaded successfully; do not keep local split json.
            os.remove(out_file)
            print(f"     removed local -> {out_file}")

    # Remove dataset split dir if it is empty after upload cleanup.
    if args.upload_s3_prefix:
        try:
            out_dir.rmdir()
        except OSError:
            pass


if __name__ == "__main__":
    main()