"""
Optimized S3 operations with parallel uploads and memory buffers.

Provides high-performance S3 operations to replace the sequential file-based
approach with parallel, memory-efficient operations.
"""

import asyncio
import concurrent.futures
from typing import List, Tuple, Dict, Any
import io
import threading
from pathlib import Path

import polars as pl
from botocore.exceptions import ClientError

from .s3_data_service import S3DataService
from ..log_context import get_logger


class OptimizedS3Service(S3DataService):
    """
    Optimized S3 service with parallel operations and memory buffers.
    
    Features:
    - Parallel uploads using ThreadPoolExecutor
    - Memory buffers instead of temporary files
    - Batch upload operations
    - Async-style operations for better throughput
    """
    
    def __init__(self, max_workers: int = 4):
        super().__init__()
        self.max_workers = max_workers
        self.upload_executor = concurrent.futures.ThreadPoolExecutor(max_workers=max_workers)
        self.logger = get_logger(__name__)
        self._upload_stats = {
            "total_uploads": 0,
            "total_bytes": 0,
            "failed_uploads": 0
        }
    
    def upload_dataset_optimized(
        self, 
        dataset_name: str, 
        provider: str, 
        month: str, 
        data: pl.DataFrame
    ) -> None:
        """
        Optimized single dataset upload with memory buffers.
        
        Args:
            dataset_name: Human-readable timestamp
            provider: Cloud provider
            month: Month in YYYY-MM format
            data: DataFrame to upload
            
        Raises:
            ValueError: If S3 configuration invalid or data is empty
        """
        self._validate_s3_config()
        
        if data.height == 0:
            raise ValueError("Cannot upload empty dataset")
        
        s3_key = self._get_dataset_s3_key(dataset_name, provider, month)
        
        # Write parquet to memory buffer (faster than temp file)
        buffer = io.BytesIO()
        data.write_parquet(buffer, compression="snappy")
        buffer_size = buffer.tell()
        buffer.seek(0)
        
        # Upload directly to S3 (no local caching needed for generated datasets)
        self._upload_from_buffer(buffer, s3_key, buffer_size)
        
        # Update statistics
        self._upload_stats["total_uploads"] += 1
        self._upload_stats["total_bytes"] += buffer_size
        
        self.logger.info(f"Successfully uploaded dataset: {dataset_name}/{provider}/{month} ({buffer_size:,} bytes)")
    
    def upload_multiple_datasets_batch(
        self, 
        uploads: List[Tuple[str, str, str, pl.DataFrame]]
    ) -> Dict[str, Any]:
        """
        Batch upload multiple datasets in parallel with progress tracking.
        
        Args:
            uploads: List of (dataset_name, provider, month, data) tuples
            
        Returns:
            Dictionary with upload results and statistics
        """
        if not uploads:
            return {"success": True, "uploaded": 0, "failed": 0, "errors": []}
        
        self.logger.info(f"Starting batch upload of {len(uploads)} datasets")
        
        # Submit all uploads concurrently
        upload_futures = []
        for dataset_name, provider, month, data in uploads:
            future = self.upload_executor.submit(
                self.upload_dataset_optimized, dataset_name, provider, month, data
            )
            upload_futures.append({
                "future": future,
                "dataset_name": dataset_name,
                "provider": provider,
                "month": month,
                "data_size": len(data) if not data.is_empty() else 0
            })
        
        # Collect results - fail-fast on any upload failure
        successful = 0
        
        for upload_info in upload_futures:
            upload_info["future"].result(timeout=600)  # 10 minute timeout per upload
            successful += 1
            self.logger.debug(f"Upload completed: {upload_info['dataset_name']}/{upload_info['provider']}/{upload_info['month']}")
        
        self.logger.info(f"Batch upload completed: {successful} successful")
        
        return {
            "success": True,
            "uploaded": successful,
            "failed": 0,
            "errors": [],
            "total_datasets": len(uploads)
        }
    
    def _upload_from_buffer(self, buffer: io.BytesIO, s3_key: str, buffer_size: int) -> None:
        """
        Upload from memory buffer instead of temp file.
        
        Args:
            buffer: Memory buffer containing parquet data
            s3_key: S3 key for upload
            buffer_size: Size of buffer for logging
        """
        buffer.seek(0)
        
        self.s3_client.upload_fileobj(
            Fileobj=buffer,
            Bucket=self.bucket,
            Key=s3_key,
            ExtraArgs={
                'ContentType': 'application/octet-stream',
                'Metadata': {
                    'content-type': 'parquet',
                    'compression': 'snappy'
                }
            }
        )
        self.logger.debug(f"S3 upload completed: s3://{self.bucket}/{s3_key} ({buffer_size:,} bytes)")
    
    
    def get_upload_statistics(self) -> Dict[str, Any]:
        """
        Get upload performance statistics.
        
        Returns:
            Dictionary with upload statistics
        """
        return {
            **self._upload_stats,
            "success_rate": (
                (self._upload_stats["total_uploads"] - self._upload_stats["failed_uploads"]) 
                / max(self._upload_stats["total_uploads"], 1) * 100
            ),
            "avg_bytes_per_upload": (
                self._upload_stats["total_bytes"] / max(self._upload_stats["total_uploads"], 1)
            )
        }
    
    def close(self) -> None:
        """Close the executor and clean up resources."""
        if hasattr(self, 'upload_executor'):
            self.upload_executor.shutdown(wait=True)
    
    def __enter__(self):
        """Context manager entry."""
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit."""
        self.close()


class StreamingUploadManager:
    """
    Manages streaming uploads with backpressure for very large datasets.
    
    Provides memory-bounded uploads that can handle datasets larger than memory
    by processing and uploading in streaming chunks.
    """
    
    def __init__(self, optimized_s3_service: OptimizedS3Service, max_concurrent_uploads: int = 3):
        self.s3_service = optimized_s3_service
        self.max_concurrent_uploads = max_concurrent_uploads
        self.upload_queue: List[Tuple[str, str, str, pl.DataFrame]] = []
        self.upload_semaphore = threading.Semaphore(max_concurrent_uploads)
        self.logger = get_logger(__name__)
    
    def queue_upload(self, dataset_name: str, provider: str, month: str, data: pl.DataFrame) -> None:
        """
        Queue an upload for batch processing.
        
        Args:
            dataset_name: Dataset name
            provider: Cloud provider
            month: Month string
            data: DataFrame to upload
        """
        self.upload_queue.append((dataset_name, provider, month, data))
    
    def process_queue_with_backpressure(self) -> Dict[str, Any]:
        """
        Process upload queue with memory backpressure control.
        
        Returns:
            Upload results dictionary
        """
        if not self.upload_queue:
            return {"success": True, "uploaded": 0, "failed": 0}
        
        self.logger.info(f"Processing upload queue with {len(self.upload_queue)} items")
        
        # Process in chunks to maintain memory bounds
        chunk_size = self.max_concurrent_uploads
        results = {"uploaded": 0, "failed": 0, "errors": []}
        
        for i in range(0, len(self.upload_queue), chunk_size):
            chunk = self.upload_queue[i:i + chunk_size]
            
            # Process chunk - will fail-fast on any upload failure
            chunk_result = self.s3_service.upload_multiple_datasets_batch(chunk)
            
            # Aggregate results
            results["uploaded"] += chunk_result["uploaded"]
            results["failed"] += chunk_result["failed"]
            results["errors"].extend(chunk_result["errors"])
            
            # Clear processed items to free memory
            del chunk
        
        # Clear the queue
        self.upload_queue.clear()
        
        results["success"] = results["failed"] == 0
        return results