#!/usr/bin/env python3
"""
Parallel S3 upload script with progress tracking.
Efficiently uploads large numbers of files to S3.
"""

import os
import boto3
from botocore.config import Config
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from tqdm import tqdm
import logging
from botocore.exceptions import ClientError


# Setup logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('s3_upload.log'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)


def upload_file(args):
    """
    Upload a single file to S3.
    
    Args:
        args: tuple of (local_path, s3_client, bucket, s3_key, check_exists)
    
    Returns:
        tuple: (success: bool, local_path: str, error_message: str or None)
    """
    local_path, s3_client, bucket, s3_key, check_exists = args
    
    try:
        # Check if file already exists in S3 (optional optimization)
        if check_exists:
            try:
                s3_client.head_object(Bucket=bucket, Key=s3_key)
                # File exists and is accessible, skip upload
                return True, local_path, None
            except ClientError as e:
                if e.response['Error']['Code'] != '404':
                    # Some other error occurred
                    return False, local_path, f"Error checking file: {str(e)}"
                # File doesn't exist, proceed with upload
        
        # Upload the file
        s3_client.upload_file(
            local_path,
            bucket,
            s3_key,
            ExtraArgs={'ContentType': 'audio/opus'}  # Set content type for opus files
        )
        return True, local_path, None
        
    except Exception as e:
        return False, local_path, str(e)


def get_all_files(directory):
    """
    Recursively get all files in a directory using os.walk (faster for large dirs).
    
    Args:
        directory: Path to scan
        
    Yields:
        Path objects for each file
    """
    for root, dirs, files in os.walk(directory):
        for filename in files:
            yield Path(root) / filename


def main():
    # Configuration
    local_dir = "/app2/suno/data/sara/sfx_get_beats/opus_audio_truncated"
    s3_bucket = "suno-data"
    s3_prefix = "datasets/bundles/v5/sfx_all_processed/opus_audio_truncated"
    
    # Performance settings
    max_workers = 32  # Number of parallel uploads (adjust based on your network)
    check_exists = True  # Set to False to skip existence checks and always upload
    
    logger.info(f"Starting parallel S3 upload")
    logger.info(f"Local directory: {local_dir}")
    logger.info(f"S3 destination: s3://{s3_bucket}/{s3_prefix}")
    logger.info(f"Max workers: {max_workers}")
    logger.info(f"Check if files exist: {check_exists}")
    
    # Initialize S3 client with larger connection pool
    # Connection pool must be >= max_workers to avoid "Connection pool is full" errors
    config = Config(
        max_pool_connections=max_workers * 2,  # 2x workers for safety
        retries={'max_attempts': 3, 'mode': 'adaptive'}
    )
    s3_client = boto3.client('s3', config=config)
    
    # Start scanning and uploading immediately (streaming approach)
    logger.info("Starting file discovery and upload...")
    local_dir_path = Path(local_dir)
    
    successful = 0
    failed = 0
    skipped = 0
    errors = []
    files_processed = 0
    
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = {}
        
        # Progress bar without total (will update as we go)
        with tqdm(desc="Uploading", unit="files") as pbar:
            # Start discovering and submitting files
            for file_path in get_all_files(local_dir):
                # Calculate relative path and S3 key
                relative_path = file_path.relative_to(local_dir_path)
                s3_key = f"{s3_prefix}/{relative_path}".replace('\\', '/')  # Handle Windows paths
                
                upload_task = (
                    str(file_path),
                    s3_client,
                    s3_bucket,
                    s3_key,
                    check_exists
                )
                
                # Submit task
                future = executor.submit(upload_file, upload_task)
                futures[future] = upload_task
                
                # Process completed uploads while discovering more files
                # This prevents memory buildup from too many pending futures
                while len(futures) > max_workers * 10:  # Keep queue manageable
                    done_futures = [f for f in futures if f.done()]
                    for future in done_futures:
                        success, local_path, error = future.result()
                        
                        if success:
                            if error is None:
                                successful += 1
                            else:
                                skipped += 1
                        else:
                            failed += 1
                            errors.append((local_path, error))
                            logger.error(f"Failed to upload {local_path}: {error}")
                        
                        files_processed += 1
                        pbar.update(1)
                        pbar.set_postfix({
                            'success': successful,
                            'skipped': skipped,
                            'failed': failed,
                            'queued': len(futures)
                        })
                        del futures[future]
            
            # Process remaining futures
            logger.info("Finishing remaining uploads...")
            for future in as_completed(futures):
                success, local_path, error = future.result()
                
                if success:
                    if error is None:
                        successful += 1
                    else:
                        skipped += 1
                else:
                    failed += 1
                    errors.append((local_path, error))
                    logger.error(f"Failed to upload {local_path}: {error}")
                
                files_processed += 1
                pbar.update(1)
                pbar.set_postfix({
                    'success': successful,
                    'skipped': skipped,
                    'failed': failed
                })
    
    # Summary
    logger.info("\n" + "="*80)
    logger.info("Upload Summary:")
    logger.info(f"Total files processed: {files_processed:,}")
    logger.info(f"Successfully uploaded: {successful:,}")
    logger.info(f"Skipped (already exist): {skipped:,}")
    logger.info(f"Failed: {failed:,}")
    logger.info("="*80)
    
    # Save errors to file
    if errors:
        error_file = "s3_upload_errors.txt"
        logger.info(f"\nWriting errors to {error_file}...")
        with open(error_file, 'w') as f:
            for local_path, error in errors:
                f.write(f"{local_path}\t{error}\n")
        logger.info(f"Errors saved to {error_file}")
    
    logger.info("\nUpload complete!")


if __name__ == "__main__":
    main()