#!/usr/bin/env python3
"""
S3 File Download Script with Parallel Processing from File List
Downloads specific files from S3 based on a list of dictionaries with file info.
"""

import os
from pathlib import Path
from joblib import Parallel, delayed
from typing import Callable, List, Dict, Optional, Any
import logging
from suno_utils.utils.text import read_jsonl
from suno_utils.audio import Audio
from tqdm import tqdm

# Configure logging (reduced level since we'll use progress bars)
logging.basicConfig(
    level=logging.WARNING,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


def default_naming_function(file_info: Dict[str, Any]) -> str:
    """
    Default file naming function - uses original filename from S3 filepath.
    
    Args:
        file_info: Dictionary containing file information with 's3_filepath' key
    
    Returns:
        Local filename
    """
    s3_filepath = file_info.get('s3_filepath', '')
    return os.path.basename(s3_filepath)


def download_single_file_from_info(
    file_info: Dict[str, Any], 
    local_dir: Path, 
    naming_function: Callable[[Dict[str, Any]], str],
    skip_existing: bool = True
) -> Dict[str, Any]:
    """
    Download a single audio file from S3 using file info dictionary.
    
    Args:
        file_info: Dictionary containing audio file information
        local_dir: Local directory to save file
        naming_function: Function to transform file_info to local filename
        skip_existing: Skip download if file already exists locally
    Returns:
        Dict with success status and metadata
    """
    try:
        file_id = file_info["id"]
        
        # Apply naming function to get local filename
        local_filename = naming_function(file_info)
        local_path = local_dir / local_filename
        
        # Skip if file already exists (optional optimization)
        if skip_existing and local_path.exists():
            return {"success": True, "skipped": True, "file_id": file_id}
        
        # Create directory if it doesn't exist
        local_path.parent.mkdir(parents=True, exist_ok=True)
        
        # Download file
        audio = Audio.from_s3(file_info["s3_filepath"])
        audio.write_opus(str(local_path))
        
        return {
            "success": True, 
            "skipped": False, 
            "file_id": file_id,
            "duration": file_info.get('duration_s', 0),
            "audio_type": file_info.get('audio_type', 'unknown')
        }
        
    except KeyError as e:
        logger.error(f"Missing required field in file_info: {e}")
        return {"success": False, "error": f"Missing field: {e}", "file_id": file_info.get('id', 'unknown')}
    except Exception as e:
        logger.error(f"Error downloading file ID {file_info.get('id', 'unknown')}: {type(e).__name__}: {e}")
        return {"success": False, "error": f"{type(e).__name__}: {e}", "file_id": file_info.get('id', 'unknown')}


def download_files_from_list(
    file_list: List[Dict[str, Any]],
    local_dir: str,
    naming_function: Optional[Callable[[Dict[str, Any]], str]] = None,
    n_jobs: int = -1,
    skip_existing: bool = True
) -> bool:
    """
    Download audio files from S3 based on a list of file information dictionaries.
    
    Args:
        file_list: List of dictionaries containing audio file information
        local_dir: Local directory to save files
        naming_function: Function to transform file_info to local filenames
        n_jobs: Number of parallel jobs (-1 uses all available cores)
        skip_existing: Skip download if files already exist locally
    Returns:
        True if all downloads successful, False otherwise
    """
    
    # Use default naming function if none provided
    if naming_function is None:
        naming_function = default_naming_function
    
    # Create local directory
    local_path = Path(local_dir)
    local_path.mkdir(parents=True, exist_ok=True)
    
    if not file_list:
        print("No files to download (empty file list)")
        return True
    
    try:
        # Log summary statistics
        total_files = len(file_list)
        audio_types = {}
        total_duration = 0
        
        for file_info in file_list:
            audio_type = file_info.get('audio_type', 'unknown')
            audio_types[audio_type] = audio_types.get(audio_type, 0) + 1
            total_duration += file_info.get('duration_s', 0)
        
        print(f"Starting download of {total_files} audio files:")
        print(f"  Total duration: {total_duration/60:.1f} minutes")
        print(f"  Audio types: {dict(audio_types)}")
        print(f"  Using {n_jobs} parallel jobs")
        print(f"  Skip existing files: {skip_existing}")
        
        # Use tqdm with joblib's backend
        with tqdm(total=total_files, desc="Downloading", unit="files") as pbar:
            
            def download_with_progress(file_info):
                result = download_single_file_from_info(file_info, local_path, naming_function, skip_existing)
                pbar.update(1)
                # Update description based on result
                return result
            
            # Download files in parallel
            results = Parallel(n_jobs=n_jobs, verbose=0, backend="threading")(
                delayed(download_with_progress)(file_info)
                for file_info in file_list
            )
        
        # Analyze results
        successful = sum(1 for r in results if r["success"])
        failed = len(results) - successful
        skipped = sum(1 for r in results if r.get("skipped", False))
        downloaded = successful - skipped
        
        # Print summary
        print(f"\n📊 Download Summary:")
        print(f"  ✅ Successfully downloaded: {downloaded}")
        print(f"  ⏭️  Skipped (already exist): {skipped}")
        print(f"  ❌ Failed: {failed}")
        print(f"  📁 Total processed: {len(results)}")
        
        if failed > 0:
            print(f"\n❌ Failed downloads:")
            failed_files = [r for r in results if not r["success"]]
            for fail in failed_files[:10]:  # Show first 10 failures
                print(f"  - {fail['file_id']}: {fail.get('error', 'Unknown error')}")
            if len(failed_files) > 10:
                print(f"  ... and {len(failed_files) - 10} more failures")
        
        return failed == 0
        
    except Exception as e:
        logger.error(f"Error during download: {type(e).__name__}: {e}")
        return False


def main():
    """
    Example usage of the S3 audio file download function.
    """
    
    # Example audio file list with your metadata structure
    audio_file_list = read_jsonl("/home/sara/sfx/metas_v1_gpt.jsonl")
    
    LOCAL_DIR = "/home/sara/sfx/raw_audio_sfx_v0/"
    N_JOBS = -1
    
    def custom_naming_function(file_info: Dict[str, Any]) -> str:
        file_id = file_info.get('id', 'unknown')
        return f"{file_id}.opus"
    
    print(f"Downloading {len(audio_file_list)} audio files from S3")
    print(f"To local directory: {LOCAL_DIR}")
    print(f"Using {N_JOBS} parallel jobs")
    
    success = download_files_from_list(
        file_list=audio_file_list,
        local_dir=LOCAL_DIR,
        naming_function=custom_naming_function,
        n_jobs=N_JOBS,
        skip_existing=True  # Skip files that already exist
    )
    
    if success:
        print("✅ Audio download completed successfully!")
    else:
        print("❌ Download failed. Check logs for details.")


if __name__ == "__main__":
    main()