"""
S3 uploader with retry logic and multi-org support.

Extracted from cz-demo-tools S3 service patterns and adapted
for cz-billing-generator upload operations.
"""

import io
import json
import boto3
from typing import List, Dict, Any, Optional
from botocore.exceptions import ClientError, NoCredentialsError
from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type
import polars as pl

from .config import S3Config

# CloudZero file upload limits and safety buffer configuration
CLOUDZERO_MAX_ROWS_PER_FILE = 1_000_000  # CloudZero's hard limit per file
CLOUDZERO_SAFETY_BUFFER_PERCENT = 0.25   # 25% safety buffer for reliable uploads
CLOUDZERO_SAFE_MAX_ROWS = int(CLOUDZERO_MAX_ROWS_PER_FILE * (1 - CLOUDZERO_SAFETY_BUFFER_PERCENT))  # 750,000 rows


class S3UploadError(Exception):
    """Custom exception for S3 upload errors."""
    pass


class S3Uploader:
    """
    S3 uploader for generated datasets.

    Extracted from cz-demo-tools S3Service and adapted for
    cz-billing-generator with multi-org support and proper error handling.
    """

    def __init__(self, s3_config: S3Config, provider_mapping_service):
        """
        Initialize S3 uploader with configuration and provider mapping.

        Args:
            s3_config: S3 configuration
            provider_mapping_service: Service for provider name transformations
        """
        self.bucket = s3_config.bucket
        self.region = s3_config.region
        self.root_folder = s3_config.root_folder
        self.provider_mapping_service = provider_mapping_service

        try:
            self.s3_client = boto3.client('s3', region_name=self.region)
        except NoCredentialsError:
            raise S3UploadError(
                "AWS credentials not found. Please set AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY environment variables."
            )
        except Exception as e:
            raise S3UploadError(f"Failed to initialize S3 client: {str(e)}")
    
    def validate_bucket_access(self) -> None:
        """
        Validate S3 bucket access and permissions.
        
        Extracted from cz-demo-tools S3Service.validate_bucket_access().
        """
        try:
            self.s3_client.head_bucket(Bucket=self.bucket)
            # Logging handled by orchestrator
        except ClientError as e:
            error_code = e.response['Error']['Code']
            if error_code == '404':
                raise S3UploadError(f"S3 bucket not found: {self.bucket}")
            elif error_code == '403':
                raise S3UploadError(f"Access denied to S3 bucket: {self.bucket}")
            else:
                raise S3UploadError(f"S3 bucket validation failed: {str(e)}")
    
    def ensure_organization_folder(self, org_folder: str) -> None:
        """
        Ensure organization folder exists in S3.
        
        Extracted from cz-demo-tools S3Service.ensure_organization_folder().
        """
        prefix = f"{self.root_folder}/{org_folder}/"
        
        try:
            # Check if any objects exist with this prefix
            response = self.s3_client.list_objects_v2(
                Bucket=self.bucket,
                Prefix=prefix,
                MaxKeys=1
            )
            
            if response.get('Contents'):
                print(f"✓ Organization folder exists: s3://{self.bucket}/{prefix}")
            else:
                print(f"✓ Organization folder ready: s3://{self.bucket}/{prefix}")
                
        except ClientError as e:
            raise S3UploadError(f"Failed to check organization folder: {str(e)}")
    
    def get_next_file_number(self, s3_prefix: str) -> int:
        """
        Get next file sequence number for upload.
        
        Extracted from cz-demo-tools S3Service file numbering logic.
        """
        try:
            response = self.s3_client.list_objects_v2(
                Bucket=self.bucket,
                Prefix=s3_prefix
            )
            
            existing_files = []
            for obj in response.get('Contents', []):
                # Extract filename from full S3 key
                filename = obj['Key'].split('/')[-1]
                if filename.startswith('data-export-') and filename.endswith('.csv.gz'):
                    existing_files.append(filename)
            
            # Calculate next file number
            latest_export_num = 0
            if existing_files:
                sorted_files = sorted(existing_files, reverse=True)
                try:
                    latest_export_num = int(sorted_files[0].split("-")[-1].split(".")[0])
                except (ValueError, IndexError):
                    print(f"    ⚠️  Warning: Could not parse existing file numbers, starting from 1")
                    latest_export_num = 0
            
            return latest_export_num + 1
            
        except ClientError as e:
            raise S3UploadError(f"Failed to get file numbering: {str(e)}")
    
    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=4, max=10),
        retry=retry_if_exception_type(ClientError),
        reraise=True
    )
    def upload_dataframe_gzip(
        self, 
        df: pl.DataFrame, 
        s3_prefix: str, 
        filename: str,
        max_rows_per_file: int = CLOUDZERO_SAFE_MAX_ROWS,
        starting_sequence: Optional[int] = None
    ) -> List[Dict[str, Any]]:
        """
        Upload Polars DataFrame as gzipped CSV to S3, splitting into multiple files if needed.
        
        Args:
            df: DataFrame to upload
            s3_prefix: S3 prefix path
            filename: Base filename (will be modified for multiple files)
            max_rows_per_file: Maximum rows per file (default 750K with 25% safety buffer for CloudZero's 1M limit)
            starting_sequence: Optional starting sequence number for integrated append workflow coordination
            
        Returns:
            List of upload results (one per file created)
        """
        # Validate input
        if df is None or df.is_empty():
            raise S3UploadError("Cannot upload empty DataFrame")
        
        # Check if splitting is needed
        if df.height <= max_rows_per_file:
            # Single file upload (existing behavior)
            return [self._upload_single_dataframe_gzip(df, s3_prefix, filename)]
        
        # Multiple file upload (new splitting behavior)
        print(f"    📊 Dataset has {df.height:,} rows, splitting into files of {max_rows_per_file:,} rows each")
        
        chunks = self._split_dataframe(df, max_rows_per_file)
        upload_results = []
        
        # Use provided starting sequence for integrated append workflow coordination,
        # or get next available file sequence number for standalone operations
        if starting_sequence is not None:
            start_sequence = starting_sequence
            print(f"    📊 Using coordinated starting sequence: {start_sequence}")
        else:
            start_sequence = self.get_next_file_number(s3_prefix)
            print(f"    📊 Calculated starting sequence from existing files: {start_sequence}")
        
        for i, chunk in enumerate(chunks):
            # Generate sequential filename starting from next available sequence number
            # Use standard naming pattern regardless of input filename for consistency
            sequence_number = start_sequence + i
            chunk_filename = f"data-export-{sequence_number:04d}.csv.gz"
            
            result = self._upload_single_dataframe_gzip(chunk, s3_prefix, chunk_filename)
            upload_results.append(result)
            
        print(f"    ✓ Split into {len(upload_results)} files, total {df.height:,} rows")
        return upload_results
    
    def _split_dataframe(self, df: pl.DataFrame, max_rows: int) -> List[pl.DataFrame]:
        """Split DataFrame into chunks of max_rows each."""
        chunks = []
        for i in range(0, df.height, max_rows):
            chunk = df.slice(i, max_rows)
            chunks.append(chunk)
        return chunks
    
    def _upload_single_dataframe_gzip(
        self, 
        df: pl.DataFrame, 
        s3_prefix: str, 
        filename: str
    ) -> Dict[str, Any]:
        """
        Upload a single Polars DataFrame as gzipped CSV to S3.
        
        Original upload_dataframe_gzip logic for single file.
        """
        s3_key = f"{s3_prefix}{filename}"
        
        try:
            # Convert Polars DataFrame to CSV and compress
            csv_data = df.write_csv()
            
            # Create gzipped CSV in memory
            import gzip
            buffer = io.BytesIO()
            with gzip.open(buffer, 'wt', encoding='utf-8') as gz_file:
                gz_file.write(csv_data)
            
            buffer.seek(0)
            
            # Get file size for progress reporting
            file_size_mb = len(buffer.getvalue()) / (1024 * 1024)
            
            # Upload to S3
            self.s3_client.put_object(
                Bucket=self.bucket,
                Key=s3_key,
                Body=buffer.getvalue(),
                ContentType='text/csv',
                ContentEncoding='gzip'
            )
            
            print(f"    ✓ Uploaded: s3://{self.bucket}/{s3_key} ({file_size_mb:.2f} MB, {df.height:,} rows)")
            
            return {
                "s3_key": s3_key,
                "filename": filename,
                "size_mb": file_size_mb,
                "rows": df.height
            }
            
        except ClientError as e:
            error_code = e.response.get('Error', {}).get('Code', 'Unknown')
            if error_code == 'AccessDenied':
                raise S3UploadError(f"Access denied uploading to S3: {s3_key}")
            elif error_code == 'NoSuchBucket':
                raise S3UploadError(f"S3 bucket not found: {self.bucket}")
            else:
                raise S3UploadError(f"Failed to upload DataFrame to S3: {str(e)}")
        except Exception as e:
            raise S3UploadError(f"Unexpected error uploading DataFrame: {str(e)}")
    
    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=4, max=10),
        retry=retry_if_exception_type(ClientError),
        reraise=True
    )
    def upload_manifest(self, manifest_data: Dict[str, Any], s3_prefix: str) -> None:
        """
        Upload manifest JSON to S3.
        
        Extracted from cz-demo-tools S3Service.upload_manifest().
        """
        s3_key = f"{s3_prefix}manifest.json"
        
        try:
            # Convert to JSON and upload
            json_data = json.dumps(manifest_data, indent=4)
            
            self.s3_client.put_object(
                Bucket=self.bucket,
                Key=s3_key,
                Body=json_data,
                ContentType='application/json'
            )
            
            print(f"    ✓ Updated manifest: s3://{self.bucket}/{s3_key}")
            
        except ClientError as e:
            raise S3UploadError(f"Failed to upload manifest: {str(e)}")
    
    def build_s3_prefix(
        self,
        org_folder: str,
        provider: str,
        billing_period: str,
        drop_id: str,
        connection_suffix: Optional[str] = None
    ) -> str:
        """
        Build S3 prefix for file upload.

        Follows cz-demo-tools S3 folder structure:
        {root_folder}/{org_folder}/{provider}/{billing_period}/{drop_id}/

        AWS Marketplace now uploads to AWS folder via provider mapping.

        If connection_suffix is provided, appends it to the provider folder name:
        {root_folder}/{org_folder}/{provider}-{suffix}/{billing_period}/{drop_id}/
        """
        folder_name = self.provider_mapping_service.get_s3_folder_name(provider)
        if connection_suffix:
            folder_name = f"{folder_name}-{connection_suffix}"
        return f"{self.root_folder}/{org_folder}/{folder_name}/{billing_period}/{drop_id}/"
    
    def build_manifest_prefix(
        self,
        org_folder: str,
        provider: str,
        billing_period: str,
        connection_suffix: Optional[str] = None
    ) -> str:
        """
        Build S3 prefix for manifest file.

        Manifest is stored at billing period level:
        {root_folder}/{org_folder}/{provider}/{billing_period}/

        AWS Marketplace now uses AWS folder via provider mapping.

        If connection_suffix is provided, appends it to the provider folder name:
        {root_folder}/{org_folder}/{provider}-{suffix}/{billing_period}/
        """
        folder_name = self.provider_mapping_service.get_s3_folder_name(provider)
        if connection_suffix:
            folder_name = f"{folder_name}-{connection_suffix}"
        return f"{self.root_folder}/{org_folder}/{folder_name}/{billing_period}/"
    
    def prune_old_drops(
        self, 
        org_folder: str, 
        provider: str, 
        billing_period: str,
        max_drops: int = 3
    ) -> List[str]:
        """
        Prune old drops for a billing period, keeping only the latest N drops.
        
        Extracted from cz-demo-tools S3Service.prune_provider_drops() logic.
        """
        errors = []
        period_prefix = f"{self.root_folder}/{org_folder}/{provider}/{billing_period}/"
        
        print(f"    🧹 Pruning old drops for {provider}/{billing_period} (keeping latest {max_drops})...")
        
        try:
            # List all drop folders in billing period
            response = self.s3_client.list_objects_v2(
                Bucket=self.bucket,
                Prefix=period_prefix,
                Delimiter='/'
            )
            
            drop_folders = []
            for prefix_info in response.get('CommonPrefixes', []):
                prefix = prefix_info['Prefix']
                # Extract drop ID from prefix
                drop_id = prefix.split('/')[-2]  # Remove trailing slash and get folder name
                if drop_id and drop_id != billing_period:  # Skip manifest.json level
                    drop_folders.append(drop_id)
            
            # Sort drop IDs (timestamp format sorts chronologically)
            drop_folders.sort(reverse=True)
            
            # Delete old drops if we exceed max_drops
            if len(drop_folders) > max_drops:
                drops_to_delete = drop_folders[max_drops:]
                print(f"      🗑️  Deleting {len(drops_to_delete)} old drops: {drops_to_delete}")
                
                for drop_id in drops_to_delete:
                    try:
                        self._delete_drop_folder(org_folder, provider, billing_period, drop_id)
                    except Exception as e:
                        error_msg = f"Failed to delete drop {drop_id}: {str(e)}"
                        errors.append(error_msg)
                        print(f"      ❌ {error_msg}")
            else:
                print(f"      ✓ No pruning needed ({len(drop_folders)} drops <= {max_drops} limit)")
        
        except Exception as e:
            error_msg = f"Failed to prune drops for {provider}/{billing_period}: {str(e)}"
            errors.append(error_msg)
            print(f"      ❌ {error_msg}")
        
        return errors
    
    def _delete_drop_folder(self, org_folder: str, provider: str, billing_period: str, drop_id: str) -> None:
        """Delete all objects in a drop folder."""
        drop_prefix = f"{self.root_folder}/{org_folder}/{provider}/{billing_period}/{drop_id}/"
        
        # List all objects in drop folder
        response = self.s3_client.list_objects_v2(
            Bucket=self.bucket,
            Prefix=drop_prefix
        )
        
        # Delete all objects
        objects_to_delete = []
        for obj in response.get('Contents', []):
            objects_to_delete.append({'Key': obj['Key']})
        
        if objects_to_delete:
            self.s3_client.delete_objects(
                Bucket=self.bucket,
                Delete={'Objects': objects_to_delete}
            )