#!/usr/bin/env python3
"""
Dataset Preparation Script for Voice Cloning Dataset
Developer: Molla Samser
Email: help@rskworld.in
Phone: +91 93305 39277
Website: https://rskworld.in
Year: 2026
"""

import os
import json
import argparse
from pathlib import Path
from collections import defaultdict


def load_config(config_path):
    """Load dataset configuration from JSON file."""
    with open(config_path, 'r') as f:
        return json.load(f)


def scan_audio_files(base_dir):
    """
    Scan directory for audio files and organize by speaker.
    
    Returns:
        Dictionary mapping speaker IDs to lists of audio files
    """
    base_path = Path(base_dir)
    audio_extensions = ['.wav', '.flac', '.mp3']
    
    speakers = defaultdict(list)
    
    # Look for speaker directories
    for speaker_dir in base_path.iterdir():
        if speaker_dir.is_dir():
            speaker_id = speaker_dir.name
            
            # Find audio files in speaker directory
            for ext in audio_extensions:
                audio_files = list(speaker_dir.rglob(f'*{ext}'))
                speakers[speaker_id].extend(audio_files)
    
    return dict(speakers)


def create_metadata(speakers_data, output_path):
    """Create metadata file for the dataset."""
    metadata = {
        'dataset_name': 'Voice Cloning Dataset',
        'version': '1.0',
        'total_speakers': len(speakers_data),
        'total_files': sum(len(files) for files in speakers_data.values()),
        'speakers': {}
    }
    
    for speaker_id, files in speakers_data.items():
        metadata['speakers'][speaker_id] = {
            'speaker_id': speaker_id,
            'num_recordings': len(files),
            'files': [str(f.relative_to(output_path.parent)) for f in files]
        }
    
    metadata_file = output_path / 'metadata.json'
    with open(metadata_file, 'w') as f:
        json.dump(metadata, f, indent=2)
    
    print(f"Metadata saved to {metadata_file}")
    return metadata


def create_train_test_split(speakers_data, train_ratio=0.8, output_path=None):
    """Create train/test split for the dataset."""
    train_files = []
    test_files = []
    
    for speaker_id, files in speakers_data.items():
        # Shuffle files
        import random
        random.shuffle(files)
        
        # Split files
        split_idx = int(len(files) * train_ratio)
        train_files.extend(files[:split_idx])
        test_files.extend(files[split_idx:])
    
    # Save split information
    if output_path:
        split_file = output_path / 'train_test_split.json'
        split_data = {
            'train_files': [str(f) for f in train_files],
            'test_files': [str(f) for f in test_files],
            'train_count': len(train_files),
            'test_count': len(test_files)
        }
        
        with open(split_file, 'w') as f:
            json.dump(split_data, f, indent=2)
        
        print(f"Train/test split saved to {split_file}")
        print(f"Train: {len(train_files)} files, Test: {len(test_files)} files")
    
    return train_files, test_files


def main():
    parser = argparse.ArgumentParser(description='Prepare dataset for training')
    parser.add_argument('--config', type=str, required=True, help='Path to dataset configuration file')
    parser.add_argument('--input', type=str, help='Input directory (overrides config)')
    parser.add_argument('--output', type=str, help='Output directory (overrides config)')
    
    args = parser.parse_args()
    
    print("=" * 60)
    print("Voice Cloning Dataset - Dataset Preparation")
    print("Developer: Molla Samser")
    print("Website: https://rskworld.in")
    print("=" * 60)
    
    # Load configuration
    config = load_config(args.config)
    
    # Determine input and output directories
    input_dir = args.input or config.get('input_directory', 'audio')
    output_dir = args.output or config.get('output_directory', 'prepared')
    
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)
    
    # Scan audio files
    print(f"\nScanning audio files in {input_dir}...")
    speakers_data = scan_audio_files(input_dir)
    
    print(f"Found {len(speakers_data)} speakers")
    for speaker_id, files in speakers_data.items():
        print(f"  {speaker_id}: {len(files)} files")
    
    # Create metadata
    print("\nCreating metadata...")
    metadata = create_metadata(speakers_data, output_path)
    
    # Create train/test split
    print("\nCreating train/test split...")
    train_ratio = config.get('train_ratio', 0.8)
    create_train_test_split(speakers_data, train_ratio, output_path)
    
    print("\nDataset preparation complete!")
    print(f"Output directory: {output_path}")


if __name__ == '__main__':
    main()