"""
S3 upload service implementation.

Provides S3 upload operations including dataset chunk uploads, manifest management,
and file organization following CloudZero folder structure.
"""

from typing import List, Dict, Any
from datetime import datetime
import polars as pl

from ..interfaces import UploadServiceInterface
from ..models import DatasetChunk, UploadResult, UploadConfig
from ..s3.uploader import S3Uploader, S3UploadError
from ..s3.formatter import S3Formatter
from ..core.dataset_loader import S3DataService


class S3UploadService(UploadServiceInterface):
    """
    Implementation of S3 upload operations.
    
    Handles uploading dataset chunks to S3, managing manifests and drop IDs,
    and organizing files according to CloudZero folder structure.
    """
    
    def __init__(self, s3_uploader: S3Uploader, s3_formatter: S3Formatter, data_processing_service):
        """
        Initialize S3 upload service.
        
        Args:
            s3_uploader: S3 uploader for file operations
            s3_formatter: S3 formatter for file naming and manifest creation
            data_processing_service: Service for data processing operations
        """
        self.s3_uploader = s3_uploader
        self.s3_formatter = s3_formatter
        self.data_processing_service = data_processing_service
        self.s3_data_service = S3DataService()
    
    def upload_dataset_chunk(
        self,
        chunk: DatasetChunk,
        drop_id: str,
        billing_period: str
    ) -> UploadResult:
        """
        Upload a single dataset chunk to S3.
        
        Args:
            chunk: Dataset chunk with data and metadata
            drop_id: Unique drop identifier for this upload batch
            billing_period: Billing period in YYYYMMDD-YYYYMMDD format
            
        Returns:
            UploadResult with upload outcome and metadata
        """
        if chunk.is_empty():
            return UploadResult(
                s3_key="",
                filename="",
                size_mb=0.0,
                rows=0,
                provider=chunk.provider,
                month=chunk.month,
                billing_period=billing_period,
                drop_id=drop_id,
                error_message="Dataset chunk is empty"
            )
        
        try:
            # Process the data with provider context
            processing_result = self.data_processing_service.process_generated_dataset(
                chunk.data,
                provider=chunk.provider
            )

            if not processing_result.is_successful():
                return UploadResult(
                    s3_key="",
                    filename="",
                    size_mb=0.0,
                    rows=0,
                    provider=chunk.provider,
                    month=chunk.month,
                    billing_period=billing_period,
                    drop_id=drop_id,
                    error_message=f"Data processing failed: {processing_result.error_message}"
                )
            
            # Build S3 prefix for this upload
            s3_prefix = self.s3_uploader.build_s3_prefix(
                org_folder=chunk.organization,
                provider=chunk.provider,
                billing_period=billing_period,
                drop_id=drop_id
            )
            
            # Get next file number and generate filename
            file_number = self.s3_uploader.get_next_file_number(s3_prefix)
            filename = self.s3_formatter.generate_filename(file_number)
            
            # Upload the processed data to S3 (may result in multiple files)
            upload_results = self.s3_uploader.upload_dataframe_gzip(
                df=processing_result.processed_data,
                s3_prefix=s3_prefix,
                filename=filename
            )
            
            # For single chunk upload, return the first result with metadata
            if upload_results:
                first_result = upload_results[0]
                total_rows = sum(result["rows"] for result in upload_results)
                
                print(f"         ✅ Uploaded {total_rows:,} rows to S3 ({len(upload_results)} files)")
                
                return UploadResult(
                    s3_key=first_result["s3_key"],
                    filename=first_result["filename"],
                    size_mb=first_result["size_mb"],
                    rows=total_rows,
                    provider=chunk.provider,
                    month=chunk.month,
                    billing_period=billing_period,
                    drop_id=drop_id
                )
            else:
                return UploadResult(
                    s3_key="",
                    filename="",
                    size_mb=0.0,
                    rows=0,
                    provider=chunk.provider,
                    month=chunk.month,
                    billing_period=billing_period,
                    drop_id=drop_id,
                    error_message="No upload results returned"
                )
            
        except Exception as e:
            print(f"         ❌ Failed to upload {chunk.provider}/{chunk.month}: {str(e)}")
            return UploadResult(
                s3_key="",
                filename="",
                size_mb=0.0,
                rows=0,
                provider=chunk.provider,
                month=chunk.month,
                billing_period=billing_period,
                drop_id=drop_id,
                error_message=str(e)
            )
    
    def stream_upload_files(
        self,
        dataset_name: str,
        provider: str,
        months: List[str],
        stage: str,
        organization: str,
        drop_id: str
    ) -> List[UploadResult]:
        """
        Stream upload multiple files for a provider.
        
        Loads and uploads files one at a time to 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}")
        
        for month in months:
            try:
                # Load single file from S3
                df = self.data_processing_service.load_dataset_from_s3(dataset_name, provider, month, stage)
                print(f"      📊 {month}: {len(df):,} rows")
                
                # Create dataset chunk
                chunk = DatasetChunk(
                    data=df,
                    provider=provider,
                    month=month,
                    stage=stage,
                    organization=organization
                )
                
                # Calculate billing period for this month
                billing_period = self.calculate_billing_period(month)
                
                # Upload the chunk
                upload_result = self.upload_dataset_chunk(chunk, drop_id, billing_period)
                upload_results.append(upload_result)
                
                # DataFrame will be garbage collected when loop iteration ends
                
            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 update_manifest_for_drop(
        self,
        organization: str,
        provider: str,
        billing_period: str,
        drop_id: str
    ) -> None:
        """
        Update manifest file for a drop.
        
        Args:
            organization: Organization name
            provider: Provider name  
            billing_period: Billing period
            drop_id: Drop identifier to add to manifest
        """
        try:
            # Build manifest prefix
            manifest_prefix = self.s3_uploader.build_manifest_prefix(
                org_folder=organization,
                provider=provider,
                billing_period=billing_period
            )
            
            # Create manifest data
            manifest_data = self.s3_formatter.create_manifest_data(drop_id)
            
            # Upload manifest
            self.s3_uploader.upload_manifest(manifest_data, manifest_prefix)
            
        except S3UploadError as e:
            print(f"      ⚠️  Manifest update failed for {provider}/{billing_period}: {str(e)}")
            raise
    
    def prune_old_drops_for_period(
        self,
        organization: str,
        provider: str,
        billing_period: str,
        max_drops: int = 3
    ) -> List[str]:
        """
        Prune old drops for a billing period, keeping latest N drops.
        
        Args:
            organization: Organization name
            provider: Provider name
            billing_period: Billing period
            max_drops: Maximum number of drops to keep
            
        Returns:
            List of error messages from pruning operations
        """
        try:
            return self.s3_uploader.prune_old_drops(
                org_folder=organization,
                provider=provider,
                billing_period=billing_period,
                max_drops=max_drops
            )
        except S3UploadError as e:
            error_msg = f"Pruning failed for {provider}/{billing_period}: {str(e)}"
            print(f"      ❌ {error_msg}")
            return [error_msg]
    
    def calculate_billing_period(self, month: str) -> str:
        """
        Calculate billing period from month string.
        
        Args:
            month: Month in YYYY-MM format
            
        Returns:
            Billing period in YYYYMMDD-YYYYMMDD format
        """
        # Parse month
        year, month_num = month.split('-')
        year = int(year)
        month_num = int(month_num)
        
        # Calculate billing period (first day of month to first day of next month)
        start_date = f"{year:04d}{month_num:02d}01"
        
        # Calculate next month
        if month_num == 12:
            next_year = year + 1
            next_month = 1
        else:
            next_year = year
            next_month = month_num + 1
        
        end_date = f"{next_year:04d}{next_month:02d}01"
        
        return f"{start_date}-{end_date}"
    
    def validate_bucket_access(self, bucket: str, region: str) -> bool:
        """
        Validate S3 bucket access and permissions.
        
        Args:
            bucket: S3 bucket name
            region: S3 region
            
        Returns:
            True if bucket is accessible, False otherwise
            
        Raises:
            S3UploadError: If validation fails with details
        """
        try:
            self.s3_uploader.validate_bucket_access()
            return True
        except S3UploadError:
            raise
    
    def generate_drop_id(self) -> str:
        """
        Generate unique drop ID for upload batch.
        
        Returns:
            Unique drop identifier (timestamp-based)
        """
        return self.s3_formatter.generate_drop_id()
    
    def update_manifests_and_prune_for_uploads(
        self,
        organization: str,
        upload_results: List[UploadResult],
        drop_id: str,
        max_drops: int = 3
    ) -> List[str]:
        """
        Update manifests and prune old drops for a list of upload results.
        
        Args:
            organization: Organization name
            upload_results: List of successful upload results
            drop_id: Drop identifier
            max_drops: Maximum drops to keep per period
            
        Returns:
            List of error messages from operations
        """
        errors = []
        
        # Group upload results by provider and billing period for manifest updates
        manifest_updates = {}
        
        for upload_result in upload_results:
            if upload_result.is_successful():
                provider = upload_result.provider
                billing_period = upload_result.billing_period
                
                key = (provider, billing_period)
                if key not in manifest_updates:
                    manifest_updates[key] = []
                manifest_updates[key].append(upload_result)
        
        # Update manifest and prune for each provider/period combination
        for (provider, billing_period), results in manifest_updates.items():
            try:
                # Update manifest
                self.update_manifest_for_drop(organization, provider, billing_period, drop_id)
                
                # Prune old drops
                prune_errors = self.prune_old_drops_for_period(
                    organization, provider, billing_period, max_drops
                )
                errors.extend(prune_errors)
                
            except Exception as e:
                error_msg = f"Manifest/pruning failed for {provider}/{billing_period}: {str(e)}"
                errors.append(error_msg)
                print(f"      ⚠️  {error_msg}")
        
        return errors