#!/usr/bin/env python3
"""Merge two preference datasets by concatenating mmap, metadata, and info files.

This module provides functionality to merge two preference datasets that were
created by the make_dataset function. The merge operation concatenates:
- Memory-mapped binary files (data_*.bin)
- Metadata JSONL files (meta_*.jsonl)
- Info JSON files (info_*.json)

The merge is done by appending the second dataset to the first, updating all
indices appropriately.
"""
import json
import os
from typing import Dict, List, Any, Tuple
import numpy as np
import tqdm
from collections import defaultdict


# Constants from the original file
SEMANTIC_N_CODEBOOKS = 1
N_TOKENS_AUDIO = 25 * 8 * 60  # max 8 mins of audio


def read_jsonl(filepath: str) -> List[Dict[str, Any]]:
    """Read JSONL file and return list of dictionaries.
    
    Args:
        filepath: Path to the JSONL file
        
    Returns:
        List of metadata dictionaries
    """
    data = []
    with open(filepath, 'r') as f:
        for line in f:
            data.append(json.loads(line.strip()))
    return data


def write_jsonl(data: List[Dict[str, Any]], filepath: str, do_append: bool = False) -> None:
    """Write list of dictionaries to JSONL file.
    
    Args:
        data: List of metadata dictionaries
        filepath: Path to output JSONL file
        do_append: Whether to append to existing file
    """
    mode = 'a' if do_append else 'w'
    with open(filepath, mode) as f:
        for item in data:
            f.write(json.dumps(item) + '\n')


def read_json(filepath: str) -> Dict[str, Any]:
    """Read JSON file and return dictionary.
    
    Args:
        filepath: Path to the JSON file
        
    Returns:
        Dictionary containing info data
    """
    with open(filepath, 'r') as f:
        return json.load(f)


def write_json(data: Dict[str, Any], filepath: str) -> None:
    """Write dictionary to JSON file.
    
    Args:
        data: Dictionary to write
        filepath: Path to output JSON file
    """
    with open(filepath, 'w') as f:
        json.dump(data, f, indent=2)


def load_sample_from_mmap(
    mmap_path: str, 
    sample_idx: int, 
    t_data_memmap: int = N_TOKENS_AUDIO
) -> np.ndarray:
    """Load a single sample from the memory-mapped file.
    
    Args:
        mmap_path: Path to the mmap binary file
        sample_idx: Index of the sample to load
        t_data_memmap: Number of tokens per sample
        
    Returns:
        Array of shape (t_data_memmap, SEMANTIC_N_CODEBOOKS)
    """
    sample_size = t_data_memmap * SEMANTIC_N_CODEBOOKS
    offset = sample_idx * sample_size
    
    mmap = np.memmap(mmap_path, dtype=np.uint16, mode='r')
    sample_data = mmap[offset:offset + sample_size].copy()
    del mmap
    
    return sample_data.reshape(t_data_memmap, SEMANTIC_N_CODEBOOKS)


def merge_preference_datasets(
    input_dir1: str,
    input_dir2: str,
    output_dir: str,
    is_val: bool = False,
    t_data_memmap: int = N_TOKENS_AUDIO,
    validate: bool = True,
) -> None:
    """Merge two preference datasets by concatenating them.
    
    This function merges two preference datasets (created by make_dataset) by:
    1. Concatenating the mmap binary files
    2. Concatenating the metadata JSONL files
    3. Merging the info JSON files with updated indices
    4. Optionally validating the merged dataset
    
    Args:
        input_dir1: Path to the first dataset directory
        input_dir2: Path to the second dataset directory
        output_dir: Path to output directory for merged dataset
        is_val: Whether this is validation set (affects file naming)
        t_data_memmap: Number of tokens per sample in the mmap
        validate: Whether to validate the merge by checking samples
        
    Raises:
        ValueError: If input files don't exist or validation fails
    """
    dset_type = "val" if is_val else "tr"
    
    # Define file paths
    mmap1_path = os.path.join(input_dir1, f"data_{dset_type}.bin")
    meta1_path = os.path.join(input_dir1, f"meta_{dset_type}.jsonl")
    info1_path = os.path.join(input_dir1, f"info_{dset_type}.json")
    
    mmap2_path = os.path.join(input_dir2, f"data_{dset_type}.bin")
    meta2_path = os.path.join(input_dir2, f"meta_{dset_type}.jsonl")
    info2_path = os.path.join(input_dir2, f"info_{dset_type}.json")
    
    output_mmap_path = os.path.join(output_dir, f"data_{dset_type}.bin")
    output_meta_path = os.path.join(output_dir, f"meta_{dset_type}.jsonl")
    output_info_path = os.path.join(output_dir, f"info_{dset_type}.json")
    
    # Validate input files exist
    for path in [mmap1_path, meta1_path, info1_path, mmap2_path, meta2_path, info2_path]:
        if not os.path.exists(path):
            raise ValueError(f"Input file does not exist: {path}")
    
    # Create output directory if it doesn't exist
    os.makedirs(output_dir, exist_ok=True)
    
    print(f"Merging {dset_type} datasets...")
    print(f"  Input 1: {input_dir1}")
    print(f"  Input 2: {input_dir2}")
    print(f"  Output:  {output_dir}")
    
    # Step 1: Load metadata to determine sample counts
    print("\n[1/4] Reading metadata files...")
    meta1 = read_jsonl(meta1_path)
    meta2 = read_jsonl(meta2_path)
    n_samples1 = len(meta1)
    n_samples2 = len(meta2)
    
    print(f"  Dataset 1: {n_samples1:,} samples")
    print(f"  Dataset 2: {n_samples2:,} samples")
    print(f"  Total:     {n_samples1 + n_samples2:,} samples")
    
    # Step 2: Concatenate mmap files
    print("\n[2/4] Concatenating mmap files...")
    sample_size = t_data_memmap * SEMANTIC_N_CODEBOOKS
    
    # Load dataset 1
    mmap1 = np.memmap(mmap1_path, dtype=np.uint16, mode='r')
    size1 = n_samples1 * sample_size
    assert len(mmap1) == size1, f"Dataset 1 size mismatch: {len(mmap1)} != {size1}"
    
    # Load dataset 2
    mmap2 = np.memmap(mmap2_path, dtype=np.uint16, mode='r')
    size2 = n_samples2 * sample_size
    assert len(mmap2) == size2, f"Dataset 2 size mismatch: {len(mmap2)} != {size2}"
    
    # Create output mmap and copy data
    total_size = size1 + size2
    output_mmap = np.memmap(output_mmap_path, dtype=np.uint16, mode='w+', shape=(total_size,))
    
    # Copy dataset 1
    print("  Copying dataset 1...")
    output_mmap[:size1] = mmap1[:]
    
    # Copy dataset 2
    print("  Copying dataset 2...")
    output_mmap[size1:total_size] = mmap2[:]
    
    output_mmap.flush()
    del output_mmap, mmap1, mmap2
    
    print(f"  Total mmap size: {total_size:,} elements ({total_size * 2 / 1024**3:.2f} GB)")
    
    # Step 3: Concatenate metadata files
    print("\n[3/4] Merging metadata files...")
    write_jsonl(meta1, output_meta_path, do_append=False)
    write_jsonl(meta2, output_meta_path, do_append=True)
    print(f"  Wrote {n_samples1 + n_samples2:,} metadata entries")
    
    # Step 4: Merge info files with updated indices
    print("\n[4/4] Merging info files...")
    info1 = read_json(info1_path)
    info2 = read_json(info2_path)
    
    # Create merged info
    merged_info = defaultdict(dict)
    
    # Copy dataset 1 info as-is
    for dataset_name, dataset_info in info1.items():
        merged_info[dataset_name] = dataset_info.copy()
    
    # Add dataset 2 info with offset indices
    for dataset_name, dataset_info in info2.items():
        if dataset_name in merged_info:
            # Dataset exists in both - merge idx_list with offset
            if "idx_list" in dataset_info:
                offset_idx_list = [idx + n_samples1 for idx in dataset_info["idx_list"]]
                merged_info[dataset_name]["idx_list"].extend(offset_idx_list)
        else:
            # New dataset - add with offset
            merged_info[dataset_name] = dataset_info.copy()
            if "idx_list" in dataset_info:
                merged_info[dataset_name]["idx_list"] = [
                    idx + n_samples1 for idx in dataset_info["idx_list"]
                ]
    
    write_json(dict(merged_info), output_info_path)
    
    # Print summary
    print("\nMerge Summary:")
    for dataset_name, dataset_info in sorted(merged_info.items()):
        n_clips = len(dataset_info.get("idx_list", []))
        print(f"  {dataset_name}: {n_clips:,} clips")
    
    # Step 5: Validation
    if validate:
        print("\n[Validation] Checking merged dataset integrity...")
        validate_merged_dataset(
            output_mmap_path,
            output_meta_path,
            output_info_path,
            t_data_memmap,
            num_samples_to_check=min(10, n_samples1 + n_samples2),
        )
        print("✓ Validation passed!")
    
    print("\n✅ Merge complete!")


def validate_merged_dataset(
    mmap_path: str,
    meta_path: str,
    info_path: str,
    t_data_memmap: int = N_TOKENS_AUDIO,
    num_samples_to_check: int = 10,
) -> None:
    """Validate the merged dataset by checking sample integrity.
    
    This function validates that:
    1. The number of metadata entries matches the mmap size
    2. Info indices point to valid metadata entries
    3. Sample data can be loaded correctly
    
    Args:
        mmap_path: Path to the mmap binary file
        meta_path: Path to the metadata JSONL file
        info_path: Path to the info JSON file
        t_data_memmap: Number of tokens per sample
        num_samples_to_check: Number of samples to check in detail
        
    Raises:
        AssertionError: If validation fails
    """
    # Load metadata and info
    metadata = read_jsonl(meta_path)
    info = read_json(info_path)
    
    n_samples = len(metadata)
    sample_size = t_data_memmap * SEMANTIC_N_CODEBOOKS
    
    # Check mmap size
    mmap = np.memmap(mmap_path, dtype=np.uint16, mode='r')
    expected_size = n_samples * sample_size
    assert len(mmap) == expected_size, f"Size mismatch: {len(mmap)} != {expected_size}"
    del mmap
    
    print(f"  ✓ Mmap size matches metadata count: {n_samples:,} samples")
    
    # Check info indices
    all_indices = set()
    for dataset_name, dataset_info in info.items():
        if "idx_list" in dataset_info:
            idx_list = dataset_info["idx_list"]
            # Check indices are in valid range
            for idx in idx_list:
                assert 0 <= idx < n_samples, f"Invalid index {idx} for dataset {dataset_name}"
                all_indices.add(idx)
    
    print(f"  ✓ All {len(all_indices):,} indices in info are valid")
    
    # Sample some entries and verify they can be loaded
    print(f"  Checking {num_samples_to_check} random samples...")
    check_indices = np.random.choice(
        n_samples, 
        size=min(num_samples_to_check, n_samples), 
        replace=False
    )
    
    for idx in tqdm.tqdm(check_indices, desc="  Validating samples"):
        # Load from mmap
        sample = load_sample_from_mmap(mmap_path, idx, t_data_memmap)
        assert sample.shape == (t_data_memmap, SEMANTIC_N_CODEBOOKS)
        
        # Check metadata exists
        meta = metadata[idx]
        assert "id" in meta, f"Metadata at index {idx} missing 'id' field"
        assert "dataset" in meta, f"Metadata at index {idx} missing 'dataset' field"
    
    print(f"  ✓ All checked samples are valid")


if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(
        description="Merge two preference datasets by concatenating them"
    )
    parser.add_argument(
        "--input_dir1",
        type=str,
        required=True,
        help="Path to the first dataset directory"
    )
    parser.add_argument(
        "--input_dir2",
        type=str,
        required=True,
        help="Path to the second dataset directory"
    )
    parser.add_argument(
        "--output_dir",
        type=str,
        required=True,
        help="Path to output directory for merged dataset"
    )
    parser.add_argument(
        "--is_val",
        action="store_true",
        help="Whether this is validation set (default: False for train set)"
    )
    parser.add_argument(
        "--t_data_memmap",
        type=int,
        default=N_TOKENS_AUDIO,
        help=f"Number of tokens per sample (default: {N_TOKENS_AUDIO})"
    )
    parser.add_argument(
        "--no_validate",
        action="store_true",
        help="Skip validation step"
    )
    
    args = parser.parse_args()
    
    merge_preference_datasets(
        input_dir1=args.input_dir1,
        input_dir2=args.input_dir2,
        output_dir=args.output_dir,
        is_val=args.is_val,
        t_data_memmap=args.t_data_memmap,
        validate=not args.no_validate,
    )