#!/usr/bin/env python3
"""
Merge multiple task-specific h=16 datasets into single training-ready dataset.
Creates unified manifest and symlinks to shards.
"""
import json
import shutil
from pathlib import Path
import argparse

def merge_datasets(task_dirs, output_dir):
    """Merge multiple CIL datasets into one."""
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    # Collect all manifests
    all_groups = []
    all_group_index = []
    all_record_index = []
    total_records = 0
    shard_counter = 0

    for task_dir in task_dirs:
        task_path = Path(task_dir)
        manifest_path = task_path / "manifest.json"

        if not manifest_path.exists():
            print(f"⚠️  Skipping {task_path.name} (no manifest)")
            continue

        print(f"Processing {task_path.name}...")
        manifest = json.loads(manifest_path.read_text())

        # Copy shard files (rename to .jsonl for compatibility)
        for shard_info in manifest.get("shards", []):
            shard_file = task_path / shard_info["path"]
            if shard_file.exists():
                dest = output_dir / f"shard_{shard_counter:04d}.jsonl"
                shutil.copy2(shard_file, dest)
                shard_counter += 1

        # Fallback: glob for .shard files if shards key missing
        if not manifest.get("shards"):
            for shard_file in task_path.glob("*.shard"):
                dest = output_dir / f"shard_{shard_counter:04d}.jsonl"
                shutil.copy2(shard_file, dest)
                shard_counter += 1

        # Read and merge group index
        group_index_path = task_path / manifest.get("group_index_path", "group_index.jsonl")
        if group_index_path.exists():
            with open(group_index_path) as f:
                for line in f:
                    if line.strip():
                        entry = json.loads(line)
                        all_groups.append(entry["group_id"])
                        all_group_index.append(entry)

        # Read and merge record index
        record_index_path = task_path / manifest.get("record_index_path", "record_index.jsonl")
        if record_index_path.exists():
            with open(record_index_path) as f:
                for line in f:
                    if line.strip():
                        all_record_index.append(json.loads(line))

        # Accumulate metadata
        total_records += manifest["record_count"]

        print(f"  ✅ {manifest['group_count']} groups, {manifest['record_count']} records")

    # Write merged indexes
    group_index_out = output_dir / "group_index.jsonl"
    with open(group_index_out, 'w') as f:
        for entry in all_group_index:
            f.write(json.dumps(entry) + '\n')

    record_index_out = output_dir / "record_index.jsonl"
    with open(record_index_out, 'w') as f:
        for entry in all_record_index:
            f.write(json.dumps(entry) + '\n')

    # Create unified manifest
    shard_list = [{"path": f"shard_{i:04d}.jsonl"} for i in range(shard_counter)]

    merged_manifest = {
        "backend": "maniskill",
        "created_at": Path(task_dirs[0]).joinpath("manifest.json").stat().st_mtime if task_dirs else 0,
        "dataset_name": "cil_h16_merged",
        "format": "dovla_cil",
        "group_count": len(all_groups),
        "group_index_path": "group_index.jsonl",
        "index_format": "jsonl",
        "k": 16,
        "num_groups": len(all_groups),
        "num_records": total_records,
        "record_count": total_records,
        "record_index_path": "record_index.jsonl",
        "schema_version": "0.1",
        "shard_count": shard_counter,
        "shard_format": "jsonl",
        "shards": shard_list,
        "horizon": 16,
        "tasks": [Path(d).name for d in task_dirs if (Path(d) / "manifest.json").exists()]
    }

    manifest_out = output_dir / "manifest.json"
    manifest_out.write_text(json.dumps(merged_manifest, indent=2))

    print("\n" + "="*60)
    print(f"✅ Merged dataset created: {output_dir}")
    print(f"   Total groups: {merged_manifest['group_count']}")
    print(f"   Total records: {merged_manifest['record_count']}")
    print(f"   Shards: {merged_manifest['shard_count']}")
    print(f"   Tasks: {', '.join(merged_manifest['tasks'])}")
    print("="*60)

    return merged_manifest

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", required=True, help="Parent dir containing task subdirs")
    parser.add_argument("--output", required=True, help="Output merged dataset dir")
    args = parser.parse_args()

    input_path = Path(args.input)
    task_dirs = [d for d in input_path.iterdir() if d.is_dir() and (d / "manifest.json").exists()]

    if not task_dirs:
        print(f"❌ No valid task directories found in {input_path}")
        return 1

    print(f"Found {len(task_dirs)} tasks to merge:")
    for d in task_dirs:
        print(f"  - {d.name}")
    print()

    merge_datasets(task_dirs, args.output)
    return 0

if __name__ == "__main__":
    import sys
    sys.exit(main())