#!/usr/bin/env python3
"""
Audio Processing 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 argparse
import librosa
import soundfile as sf
import numpy as np
from pathlib import Path
from tqdm import tqdm


def process_audio_file(input_path, output_path, target_sr=22050, normalize=True):
    """
    Process a single audio file: load, resample, normalize, and save.
    
    Args:
        input_path: Path to input audio file
        output_path: Path to save processed audio
        target_sr: Target sample rate (default: 22050)
        normalize: Whether to normalize audio (default: True)
    """
    try:
        # Load audio file
        audio, sr = librosa.load(input_path, sr=target_sr, mono=True)
        
        # Normalize audio if requested
        if normalize:
            audio = librosa.util.normalize(audio)
        
        # Ensure output directory exists
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        
        # Save processed audio
        sf.write(output_path, audio, target_sr)
        
        return True
    except Exception as e:
        print(f"Error processing {input_path}: {str(e)}")
        return False


def process_directory(input_dir, output_dir, target_sr=22050, normalize=True):
    """
    Process all audio files in a directory.
    
    Args:
        input_dir: Input directory containing audio files
        output_dir: Output directory for processed files
        target_sr: Target sample rate
        normalize: Whether to normalize audio
    """
    input_path = Path(input_dir)
    output_path = Path(output_dir)
    
    # Supported audio formats
    audio_extensions = ['.wav', '.flac', '.mp3', '.m4a']
    
    # Find all audio files
    audio_files = []
    for ext in audio_extensions:
        audio_files.extend(input_path.rglob(f'*{ext}'))
    
    print(f"Found {len(audio_files)} audio files to process")
    
    # Process each file
    success_count = 0
    for audio_file in tqdm(audio_files, desc="Processing audio files"):
        # Create relative path structure
        relative_path = audio_file.relative_to(input_path)
        output_file = output_path / relative_path.with_suffix('.wav')
        
        if process_audio_file(str(audio_file), str(output_file), target_sr, normalize):
            success_count += 1
    
    print(f"Successfully processed {success_count}/{len(audio_files)} files")


def main():
    parser = argparse.ArgumentParser(description='Process audio files for voice cloning dataset')
    parser.add_argument('--input', type=str, required=True, help='Input directory containing audio files')
    parser.add_argument('--output', type=str, required=True, help='Output directory for processed files')
    parser.add_argument('--sample-rate', type=int, default=22050, help='Target sample rate (default: 22050)')
    parser.add_argument('--no-normalize', action='store_true', help='Disable audio normalization')
    
    args = parser.parse_args()
    
    print("=" * 60)
    print("Voice Cloning Dataset - Audio Processing")
    print("Developer: Molla Samser")
    print("Website: https://rskworld.in")
    print("=" * 60)
    
    process_directory(
        args.input,
        args.output,
        target_sr=args.sample_rate,
        normalize=not args.no_normalize
    )
    
    print("\nProcessing complete!")


if __name__ == '__main__':
    main()