#!/usr/bin/env python3
"""
Dataset Validation 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
import librosa


def validate_audio_file(audio_path):
    """
    Validate a single audio file.
    
    Returns:
        Dictionary with validation results
    """
    result = {
        'file': str(audio_path),
        'valid': False,
        'errors': [],
        'warnings': [],
        'info': {}
    }
    
    try:
        # Check if file exists
        if not Path(audio_path).exists():
            result['errors'].append('File does not exist')
            return result
        
        # Try to load audio
        try:
            audio, sr = librosa.load(audio_path, sr=None, mono=False)
            result['info']['sample_rate'] = sr
            result['info']['duration'] = librosa.get_duration(y=audio, sr=sr)
            result['info']['channels'] = 1 if len(audio.shape) == 1 else audio.shape[0]
            result['info']['shape'] = audio.shape
            
            # Check duration
            if result['info']['duration'] < 0.5:
                result['warnings'].append('Audio duration is very short (< 0.5s)')
            if result['info']['duration'] > 60:
                result['warnings'].append('Audio duration is very long (> 60s)')
            
            # Check sample rate
            if sr < 16000:
                result['warnings'].append(f'Sample rate is low ({sr} Hz)')
            if sr > 48000:
                result['warnings'].append(f'Sample rate is very high ({sr} Hz)')
            
            # Check for silence
            if abs(audio).max() < 0.01:
                result['warnings'].append('Audio appears to be silent or very quiet')
            
            result['valid'] = True
            
        except Exception as e:
            result['errors'].append(f'Error loading audio: {str(e)}')
            
    except Exception as e:
        result['errors'].append(f'Unexpected error: {str(e)}')
    
    return result


def validate_speaker_directory(speaker_dir):
    """Validate all files in a speaker directory."""
    speaker_path = Path(speaker_dir)
    results = {
        'speaker_id': speaker_path.name,
        'valid': True,
        'audio_files': [],
        'metadata_exists': False,
        'transcript_exists': False,
        'errors': [],
        'warnings': []
    }
    
    # Check for metadata.json
    metadata_file = speaker_path / 'metadata.json'
    if metadata_file.exists():
        results['metadata_exists'] = True
        try:
            with open(metadata_file, 'r') as f:
                metadata = json.load(f)
                results['metadata'] = metadata
        except Exception as e:
            results['errors'].append(f'Error reading metadata.json: {str(e)}')
    else:
        results['warnings'].append('metadata.json not found')
    
    # Check for transcripts.txt
    transcript_file = speaker_path / 'transcripts.txt'
    if transcript_file.exists():
        results['transcript_exists'] = True
    else:
        results['warnings'].append('transcripts.txt not found')
    
    # Validate audio files
    recordings_dir = speaker_path / 'recordings'
    if recordings_dir.exists():
        audio_extensions = ['.wav', '.flac', '.mp3']
        for ext in audio_extensions:
            for audio_file in recordings_dir.rglob(f'*{ext}'):
                validation = validate_audio_file(audio_file)
                results['audio_files'].append(validation)
                if not validation['valid']:
                    results['valid'] = False
    else:
        results['warnings'].append('recordings/ directory not found')
    
    return results


def validate_dataset(dataset_dir):
    """Validate entire dataset."""
    dataset_path = Path(dataset_dir)
    results = {
        'dataset_path': str(dataset_path),
        'total_speakers': 0,
        'valid_speakers': 0,
        'total_audio_files': 0,
        'valid_audio_files': 0,
        'speakers': [],
        'errors': [],
        'warnings': []
    }
    
    # Find all speaker directories
    for item in dataset_path.iterdir():
        if item.is_dir() and item.name.startswith('speaker_'):
            results['total_speakers'] += 1
            speaker_result = validate_speaker_directory(item)
            results['speakers'].append(speaker_result)
            
            if speaker_result['valid']:
                results['valid_speakers'] += 1
            
            results['total_audio_files'] += len(speaker_result['audio_files'])
            results['valid_audio_files'] += sum(
                1 for f in speaker_result['audio_files'] if f['valid']
            )
            
            if speaker_result['errors']:
                results['errors'].extend(speaker_result['errors'])
            if speaker_result['warnings']:
                results['warnings'].extend(speaker_result['warnings'])
    
    return results


def main():
    parser = argparse.ArgumentParser(description='Validate voice cloning dataset')
    parser.add_argument('--dataset', type=str, default='audio', 
                       help='Dataset directory (default: audio)')
    parser.add_argument('--output', type=str, 
                       help='Output JSON file for validation results')
    parser.add_argument('--verbose', action='store_true',
                       help='Show detailed validation information')
    
    args = parser.parse_args()
    
    print("=" * 60)
    print("Voice Cloning Dataset - Validation")
    print("Developer: Molla Samser")
    print("Website: https://rskworld.in")
    print("=" * 60)
    
    # Validate dataset
    results = validate_dataset(args.dataset)
    
    # Print summary
    print(f"\nValidation Summary:")
    print(f"  Total Speakers: {results['total_speakers']}")
    print(f"  Valid Speakers: {results['valid_speakers']}")
    print(f"  Total Audio Files: {results['total_audio_files']}")
    print(f"  Valid Audio Files: {results['valid_audio_files']}")
    
    if results['errors']:
        print(f"\n  Errors: {len(results['errors'])}")
        if args.verbose:
            for error in results['errors']:
                print(f"    - {error}")
    
    if results['warnings']:
        print(f"\n  Warnings: {len(results['warnings'])}")
        if args.verbose:
            for warning in results['warnings']:
                print(f"    - {warning}")
    
    # Save results if output specified
    if args.output:
        with open(args.output, 'w') as f:
            json.dump(results, f, indent=2)
        print(f"\nValidation results saved to {args.output}")
    
    # Exit with error code if validation failed
    if results['total_speakers'] > 0 and results['valid_speakers'] < results['total_speakers']:
        print("\n⚠ Validation completed with errors")
        return 1
    
    print("\n✓ Validation completed successfully")
    return 0


if __name__ == '__main__':
    exit(main())