#!/usr/bin/env python3
"""
Audio Format Conversion 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
from pathlib import Path
from tqdm import tqdm


def convert_audio_file(input_path, output_path, target_format='wav', target_sr=22050):
    """
    Convert audio file to target format.
    
    Args:
        input_path: Path to input audio file
        output_path: Path to save converted audio
        target_format: Target format (wav, flac, mp3)
        target_sr: Target sample rate
    """
    try:
        # Load audio
        audio, sr = librosa.load(input_path, sr=target_sr, mono=True)
        
        # Ensure output directory exists
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        
        # Save in target format
        if target_format.lower() == 'wav':
            sf.write(output_path, audio, target_sr, format='WAV')
        elif target_format.lower() == 'flac':
            sf.write(output_path, audio, target_sr, format='FLAC')
        elif target_format.lower() == 'mp3':
            # Note: MP3 requires additional library (pydub or similar)
            sf.write(output_path, audio, target_sr, format='WAV')
            print(f"Warning: MP3 conversion may require additional libraries")
        else:
            sf.write(output_path, audio, target_sr)
        
        return True
    except Exception as e:
        print(f"Error converting {input_path}: {str(e)}")
        return False


def convert_directory(input_dir, output_dir, target_format='wav', target_sr=22050):
    """
    Convert all audio files in a directory.
    
    Args:
        input_dir: Input directory
        output_dir: Output directory
        target_format: Target format
        target_sr: Target sample rate
    """
    input_path = Path(input_dir)
    output_path = Path(output_dir)
    
    # Supported input formats
    audio_extensions = ['.wav', '.flac', '.mp3', '.m4a', '.ogg']
    
    # 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 convert")
    
    # Convert each file
    success_count = 0
    for audio_file in tqdm(audio_files, desc="Converting audio files"):
        # Create relative path structure
        relative_path = audio_file.relative_to(input_path)
        output_file = output_path / relative_path.with_suffix(f'.{target_format.lower()}')
        
        if convert_audio_file(str(audio_file), str(output_file), target_format, target_sr):
            success_count += 1
    
    print(f"Successfully converted {success_count}/{len(audio_files)} files")


def main():
    parser = argparse.ArgumentParser(description='Convert audio files to different format')
    parser.add_argument('--input', type=str, required=True, help='Input directory')
    parser.add_argument('--output', type=str, required=True, help='Output directory')
    parser.add_argument('--format', type=str, default='wav', choices=['wav', 'flac', 'mp3'],
                       help='Target format (default: wav)')
    parser.add_argument('--sample-rate', type=int, default=22050,
                       help='Target sample rate (default: 22050)')
    
    args = parser.parse_args()
    
    print("=" * 60)
    print("Voice Cloning Dataset - Format Conversion")
    print("Developer: Molla Samser")
    print("Website: https://rskworld.in")
    print("=" * 60)
    
    convert_directory(
        args.input,
        args.output,
        target_format=args.format,
        target_sr=args.sample_rate
    )
    
    print("\nConversion complete!")


if __name__ == '__main__':
    main()