"""
Streaming-optimized S3 upload service.

Memory-efficient upload service with streaming support for large datasets,
dynamic memory management, and optimized chunking strategies.
"""

import gc
import psutil
from typing import List, Dict, Any, Iterator, Optional
import polars as pl

from .s3_upload_service import S3UploadService
from ..models import DatasetChunk, UploadResult


# Memory thresholds for streaming decisions
STREAMING_THRESHOLD_MB = 500  # 500MB threshold for streaming mode
MIN_CHUNK_SIZE_ROWS = 100_000  # Minimum chunk size
MAX_CHUNK_SIZE_ROWS = 2_000_000  # Maximum chunk size
MEMORY_SAFETY_FACTOR = 0.8  # Use 80% of available memory


class StreamingS3UploadService(S3UploadService):
    """
    Memory-optimized S3 upload service with streaming support.
    
    Extends base S3UploadService with dynamic memory management,
    streaming for large datasets, and optimized chunking.
    """
    
    def stream_upload_files(
        self,
        dataset_name: str,
        provider: str,
        months: List[str],
        stage: str,
        organization: str,
        drop_id: str,
        connection_suffix: Optional[str] = None
    ) -> List[UploadResult]:
        """
        Memory-optimized streaming upload for multiple files.
        
        Uses adaptive memory management and streaming for large datasets
        to prevent OOM errors and minimize memory usage.
        
        Args:
            dataset_name: Name of the dataset
            provider: Provider name
            months: List of months to process
            stage: Processing stage
            organization: Organization name
            drop_id: Unique drop identifier
            
        Returns:
            List of UploadResult for each uploaded file
        """
        upload_results = []
        
        print(f"   🌤️  Processing provider: {provider} (from {stage} stage)")
        print(f"      📅 Processing {len(months)} months: {months}")
        print(f"      🧠 Available memory: {self._get_available_memory_mb():.0f}MB")
        
        for month in months:
            try:
                # Estimate dataset size before loading
                dataset_size_mb = self._estimate_dataset_size_mb(dataset_name, provider, month, stage)
                print(f"      📊 {month}: ~{dataset_size_mb:.0f}MB estimated")
                
                if dataset_size_mb > STREAMING_THRESHOLD_MB:
                    # Stream process large datasets
                    print(f"         🌊 Using streaming mode (dataset > {STREAMING_THRESHOLD_MB}MB)")
                    streaming_results = list(self._stream_process_large_dataset(
                        dataset_name, provider, month, stage, organization, drop_id, connection_suffix
                    ))
                    upload_results.extend(streaming_results)
                else:
                    # Normal processing for small datasets with optimized loading
                    print(f"         💾 Using normal mode (dataset < {STREAMING_THRESHOLD_MB}MB)")
                    result = self._process_normal_dataset_optimized(
                        dataset_name, provider, month, stage, organization, drop_id, connection_suffix
                    )
                    upload_results.append(result)
                
                # Explicit memory cleanup between months
                self._cleanup_memory()
                
            except Exception as e:
                print(f"         ❌ Failed to process {provider}/{month}: {str(e)}")
                # Create failed upload result
                billing_period = self.calculate_billing_period(month)
                failed_result = UploadResult(
                    s3_key="",
                    filename="",
                    size_mb=0.0,
                    rows=0,
                    provider=provider,
                    month=month,
                    billing_period=billing_period,
                    drop_id=drop_id,
                    error_message=str(e)
                )
                upload_results.append(failed_result)
        
        return upload_results
    
    def _stream_process_large_dataset(
        self,
        dataset_name: str,
        provider: str,
        month: str,
        stage: str,
        organization: str,
        drop_id: str,
        connection_suffix: Optional[str] = None
    ) -> Iterator[UploadResult]:
        """
        Stream process large dataset in memory-efficient chunks.
        
        Args:
            dataset_name: Name of the dataset
            provider: Provider name
            month: Month to process
            stage: Processing stage
            organization: Organization name
            drop_id: Drop identifier
            
        Yields:
            UploadResult for each processed chunk
        """
        # Load dataset metadata to determine chunking strategy
        total_rows = self._get_dataset_row_count(dataset_name, provider, month, stage)
        chunk_size = self._calculate_optimal_chunk_size(total_rows)
        
        print(f"         📦 Chunking {total_rows:,} rows into chunks of {chunk_size:,}")
        
        # Process dataset in chunks
        chunk_num = 0
        for chunk_start in range(0, total_rows, chunk_size):
            chunk_num += 1
            chunk_end = min(chunk_start + chunk_size, total_rows)
            
            try:
                print(f"           🔄 Processing chunk {chunk_num}: rows {chunk_start:,}-{chunk_end:,}")
                
                # Load chunk with memory monitoring
                memory_before = self._get_memory_usage_mb()
                chunk_df = self._load_dataset_chunk(dataset_name, provider, month, stage, chunk_start, chunk_size)
                memory_after = self._get_memory_usage_mb()
                
                print(f"           📈 Memory usage: {memory_before:.0f}MB → {memory_after:.0f}MB (+{memory_after-memory_before:.0f}MB)")
                
                # Create dataset chunk
                chunk = DatasetChunk(
                    data=chunk_df,
                    provider=provider,
                    month=f"{month}_chunk_{chunk_num}",  # Unique identifier for chunk
                    stage=stage,
                    organization=organization
                )
                
                # Calculate billing period
                billing_period = self.calculate_billing_period(month)

                # Upload the chunk
                upload_result = self.upload_dataset_chunk(chunk, drop_id, billing_period, connection_suffix)
                
                # Update chunk metadata
                upload_result.month = month  # Restore original month (remove chunk suffix)
                
                # Explicit cleanup after processing
                del chunk_df, chunk
                self._cleanup_memory()
                
                yield upload_result
                
            except Exception as e:
                print(f"           ❌ Chunk {chunk_num} failed: {str(e)}")
                billing_period = self.calculate_billing_period(month)
                yield UploadResult(
                    s3_key="",
                    filename="",
                    size_mb=0.0,
                    rows=0,
                    provider=provider,
                    month=month,
                    billing_period=billing_period,
                    drop_id=drop_id,
                    error_message=f"Chunk {chunk_num} failed: {str(e)}"
                )
    
    def _process_normal_dataset_optimized(
        self,
        dataset_name: str,
        provider: str,
        month: str,
        stage: str,
        organization: str,
        drop_id: str,
        connection_suffix: Optional[str] = None
    ) -> UploadResult:
        """
        Process normal-sized dataset with memory optimization.
        
        Uses column selection and optimized loading for datasets under streaming threshold.
        """
        # Load dataset with required columns only
        df = self._load_dataset_with_column_selection(dataset_name, provider, month, stage)
        print(f"           📊 Loaded {len(df):,} rows with {len(df.columns)} columns")
        
        # Create dataset chunk
        chunk = DatasetChunk(
            data=df,
            provider=provider,
            month=month,
            stage=stage,
            organization=organization
        )
        
        # Calculate billing period
        billing_period = self.calculate_billing_period(month)

        # Upload the chunk
        return self.upload_dataset_chunk(chunk, drop_id, billing_period, connection_suffix)
    
    def _load_dataset_with_column_selection(
        self,
        dataset_name: str,
        provider: str,
        month: str,
        stage: str
    ) -> pl.DataFrame:
        """
        Load dataset with only required columns for CBF processing.
        
        Reduces memory usage by 30-50% by filtering unnecessary columns early.
        """
        # Get required columns from CBF mapping
        required_columns = self._get_required_columns_for_cbf()
        
        # Load with column selection
        s3_key = f"czbillgen/datasets/{dataset_name}/{stage}/{provider}/{month}.parquet"
        
        # Use lazy loading with column selection
        df = (
            self.s3_data_service.scan_parquet_from_s3(s3_key)  # Lazy scan
            .select([col for col in required_columns if col])  # Select only required columns
            .collect()  # Execute query
        )
        
        return df
    
    def _load_dataset_chunk(
        self,
        dataset_name: str,
        provider: str,
        month: str,
        stage: str,
        offset: int,
        limit: int
    ) -> pl.DataFrame:
        """Load a specific chunk of dataset using row offset and limit."""
        s3_key = f"czbillgen/datasets/{dataset_name}/{stage}/{provider}/{month}.parquet"
        
        # Use lazy loading with row slicing and column selection
        required_columns = self._get_required_columns_for_cbf()
        
        df = (
            self.s3_data_service.scan_parquet_from_s3(s3_key)
            .select([col for col in required_columns if col])
            .slice(offset, limit)
            .collect()
        )
        
        return df
    
    def _estimate_dataset_size_mb(self, dataset_name: str, provider: str, month: str, stage: str) -> float:
        """Estimate dataset size in MB without loading the full dataset."""
        try:
            # Get file size from S3 metadata
            s3_key = f"czbillgen/datasets/{dataset_name}/{stage}/{provider}/{month}.parquet"
            file_size_bytes = self.s3_data_service.get_file_size_from_s3(s3_key)
            return file_size_bytes / (1024 * 1024)  # Convert to MB
        except Exception:
            # Fallback: estimate based on row count and average row size
            row_count = self._get_dataset_row_count(dataset_name, provider, month, stage)
            estimated_mb = (row_count * 500) / (1024 * 1024)  # Assume ~500 bytes per row
            return estimated_mb
    
    def _get_dataset_row_count(self, dataset_name: str, provider: str, month: str, stage: str) -> int:
        """Get dataset row count without loading full dataset."""
        s3_key = f"czbillgen/datasets/{dataset_name}/{stage}/{provider}/{month}.parquet"
        return self.s3_data_service.get_parquet_row_count(s3_key)
    
    def _calculate_optimal_chunk_size(self, total_rows: int) -> int:
        """
        Calculate optimal chunk size based on available memory and dataset characteristics.
        
        Args:
            total_rows: Total rows in dataset
            
        Returns:
            Optimal chunk size in rows
        """
        # Get available memory in bytes
        available_memory_bytes = psutil.virtual_memory().available * MEMORY_SAFETY_FACTOR
        
        # Estimate memory per row (rough approximation)
        estimated_bytes_per_row = 1000  # Conservative estimate
        
        # Calculate chunk size based on memory
        memory_based_chunk_size = int(available_memory_bytes / (estimated_bytes_per_row * 2))  # Divide by 2 for safety
        
        # Apply constraints
        chunk_size = max(MIN_CHUNK_SIZE_ROWS, min(memory_based_chunk_size, MAX_CHUNK_SIZE_ROWS))
        
        # Don't make chunks larger than 25% of total rows
        max_chunk_for_dataset = max(MIN_CHUNK_SIZE_ROWS, total_rows // 4)
        chunk_size = min(chunk_size, max_chunk_for_dataset)
        
        return chunk_size
    
    def _get_required_columns_for_cbf(self) -> List[str]:
        """Get list of columns required for CBF processing."""
        # This would be implemented based on CBF mapping configuration
        # For now, return a placeholder
        return ["*"]  # Load all columns - to be optimized based on actual CBF mapping
    
    def _get_available_memory_mb(self) -> float:
        """Get available system memory in MB."""
        return psutil.virtual_memory().available / (1024 * 1024)
    
    def _get_memory_usage_mb(self) -> float:
        """Get current process memory usage in MB."""
        return psutil.Process().memory_info().rss / (1024 * 1024)
    
    def _cleanup_memory(self) -> None:
        """Force garbage collection to free memory."""
        gc.collect()
        # Additional cleanup can be added here if needed