"""
S3 formatting utilities for upload operations.

Provides drop ID generation, manifest management, and S3-specific
formatting extracted from cz-demo-tools patterns.
"""

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


class S3Formatter:
    """
    S3-specific formatting utilities for upload operations.
    
    Extracted from cz-demo-tools to provide consistent S3 formatting
    including drop ID generation, file naming, and manifest management.
    """
    
    def generate_drop_id(self) -> str:
        """
        Generate drop ID timestamp.
        
        Extracted from cz-demo-tools to provide consistent drop ID
        format: YYYYMMDDTHHMMSSZ (UTC timestamp).
        """
        return datetime.utcnow().strftime("%Y%m%dT%H%M%SZ")
    
    def split_by_month(self, df: pl.DataFrame) -> Dict[str, pl.DataFrame]:
        """
        Split DataFrame by month using time/usage_start column.
        
        Extracted from cz-demo-tools month-based splitting logic
        to organize data by billing periods.
        """
        if "time/usage_start" not in df.columns:
            raise ValueError("DataFrame must contain 'time/usage_start' column for month splitting")
        
        # Convert to datetime if needed
        df = df.with_columns(
            pl.col("time/usage_start").str.strptime(pl.Datetime, "%Y-%m-%dT%H:%M:%S")
        )
        
        # Group by month and create billing period keys
        monthly_data = {}
        for month_group in df.group_by(pl.col("time/usage_start").dt.month()):
            month_df = month_group[1]
            
            # Extract month for billing period calculation
            sample_date = month_df["time/usage_start"][0]
            billing_period = self._calculate_billing_period(sample_date)
            
            monthly_data[billing_period] = month_df
        
        return monthly_data
    
    def filter_by_providers(self, df: pl.DataFrame, providers: List[str]) -> pl.DataFrame:
        """
        Filter DataFrame by specified cloud providers.
        
        Looks for provider information in common columns like
        'cloud_provider', 'resource/service', etc.
        """
        if df.empty:
            return df
        
        # Check for cloud_provider column first
        if "cloud_provider" in df.columns:
            return df.filter(pl.col("cloud_provider").is_in(providers))
        
        # Check for service-based provider detection
        if "resource/service" in df.columns:
            # This is a simplified provider detection - would need more sophisticated logic
            # based on actual service patterns in the data
            provider_patterns = {
                "AWS": ["aws", "amazon", "ec2", "s3", "rds"],
                "GCP": ["gcp", "google", "compute", "storage"],
                "Azure": ["azure", "microsoft"]
            }
            
            filter_conditions = []
            for provider in providers:
                if provider in provider_patterns:
                    patterns = provider_patterns[provider]
                    provider_condition = pl.col("resource/service").str.to_lowercase().str.contains(
                        "|".join(patterns)
                    )
                    filter_conditions.append(provider_condition)
            
            if filter_conditions:
                combined_filter = filter_conditions[0]
                for condition in filter_conditions[1:]:
                    combined_filter = combined_filter | condition
                return df.filter(combined_filter)
        
        # If no provider detection possible, return all data
        return df
    
    def generate_filename(self, sequence_number: int) -> str:
        """
        Generate S3 filename with sequence number.
        
        Extracted from cz-demo-tools file naming pattern:
        data-export-NNNN.csv.gz (4-digit zero-padded).
        """
        return f"data-export-{sequence_number:04d}.csv.gz"
    
    def create_manifest_data(self, drop_id: str) -> Dict[str, Any]:
        """
        Create manifest data dictionary.
        
        Extracted from cz-demo-tools manifest format.
        """
        return {
            "version": "1.3.0",
            "current_drop_id": drop_id
        }
    
    def _calculate_billing_period(self, date_value) -> str:
        """
        Calculate billing period from date.
        
        Extracted from cz-demo-tools billing period calculation.
        Format: YYYYMMDD-YYYYMMDD (first day of month to first day of next month).
        """
        # Extract year and month from the date
        if hasattr(date_value, 'year'):  # datetime object
            year = date_value.year
            month = date_value.month
        else:  # string parsing
            date_str = str(date_value)[:10]  # YYYY-MM-DD part
            year = int(date_str[:4])
            month = int(date_str[5:7])
        
        # Calculate billing period
        start_date = f"{year:04d}{month:02d}01"
        
        # Calculate next month
        if month == 12:
            next_year = year + 1
            next_month = 1
        else:
            next_year = year
            next_month = month + 1
        
        end_date = f"{next_year:04d}{next_month:02d}01"
        
        return f"{start_date}-{end_date}"