"""
Data processing service implementation.

Provides data processing operations including CBF mapping, dataset loading,
and data transformation for both generated and ad-hoc uploads.
"""

from typing import List, Tuple
from pathlib import Path
import polars as pl

from ..interfaces import DataProcessingServiceInterface
from ..models import ProcessingResult, CBFProcessingConfig
from ..core.sanitizer import DataSanitizer
from ..core.cbf_mapper import load_mapping_file
from ..core.optimized_cbf_mapper import apply_mapping_and_validate_with_fallback
from ..core.dataset_loader import DatasetLoader


class DataProcessingService(DataProcessingServiceInterface):
    """
    Implementation of data processing operations.
    
    Handles data transformation, CBF mapping, and validation for both
    generated datasets and ad-hoc CSV uploads following Single Responsibility Principle.
    """
    
    def __init__(self, cbf_config: CBFProcessingConfig, provider_mapping_service):
        """
        Initialize data processing service.

        Args:
            cbf_config: CBF processing configuration
            provider_mapping_service: Service for provider name transformations
        """
        self.cbf_config = cbf_config
        self.sanitizer = DataSanitizer()
        self.dataset_loader = DatasetLoader()
        self.provider_mapping_service = provider_mapping_service

        # Load CBF mapping configuration
        self.mapping, self.required_cbf_columns = load_mapping_file(cbf_config.mapping_path)
    
    def process_generated_dataset(self, df: pl.DataFrame, provider: str) -> ProcessingResult:
        """
        Process generated dataset with CBF mapping and provider column.

        Generated datasets from the pipeline are already properly formatted
        and validated, so they only need CBF mapping and API formatting.

        Args:
            df: Generated dataset DataFrame
            provider: Provider name for cloud_provider column value

        Returns:
            ProcessingResult with processed data and metadata
        """
        if df.is_empty():
            return ProcessingResult(
                processed_data=df,
                original_row_count=0,
                processed_row_count=0,
                processing_type="generated",
                error_message="Input DataFrame is empty"
            )

        original_row_count = len(df)

        try:
            # Apply optimized CBF mapping and validation with lazy evaluation
            processed_df = apply_mapping_and_validate_with_fallback(
                df, self.mapping, self.required_cbf_columns, use_optimized=True
            )

            # Add lineitem/cloud_provider column using vectorized operation
            cloud_provider_value = self.provider_mapping_service.get_cloud_provider_value(provider)
            processed_df = processed_df.with_columns([
                pl.lit(cloud_provider_value).alias("lineitem/cloud_provider")
            ])

            # Get CBF columns that were applied
            cbf_columns_applied = list(processed_df.columns)

            return ProcessingResult(
                processed_data=processed_df,
                original_row_count=original_row_count,
                processed_row_count=len(processed_df),
                processing_type="generated",
                cbf_columns_applied=cbf_columns_applied,
                sanitization_applied=False
            )

        except Exception as e:
            return ProcessingResult(
                processed_data=pl.DataFrame(),
                original_row_count=original_row_count,
                processed_row_count=0,
                processing_type="generated",
                error_message=f"CBF processing failed: {str(e)}"
            )
    
    def process_adhoc_dataset(
        self,
        df: pl.DataFrame,
        user_mapping: List[Tuple[str, str]]
    ) -> ProcessingResult:
        """
        Process ad-hoc dataset with sanitization and user mapping.

        Ad-hoc datasets require sanitization to clean dirty user data
        before CBF mapping and validation.

        Args:
            df: Raw ad-hoc dataset DataFrame
            user_mapping: List of (source_column, target_column) tuples

        Returns:
            ProcessingResult with processed and sanitized data
        """
        if df.is_empty():
            return self._create_empty_adhoc_result()

        original_row_count = len(df)

        try:
            mapped_df = self._apply_adhoc_user_mapping(df, user_mapping)
            sanitized_df = self._sanitize_adhoc_data(mapped_df)
            return self._create_successful_adhoc_result(sanitized_df, original_row_count)

        except Exception as e:
            return self._create_failed_adhoc_result(original_row_count, str(e))

    def _create_empty_adhoc_result(self) -> ProcessingResult:
        """Create result for empty input DataFrame."""
        return ProcessingResult(
            processed_data=pl.DataFrame(),
            original_row_count=0,
            processed_row_count=0,
            processing_type="adhoc",
            error_message="Input DataFrame is empty"
        )

    def _apply_adhoc_user_mapping(self, df: pl.DataFrame, user_mapping: List[Tuple[str, str]]) -> pl.DataFrame:
        """Apply user-defined column mapping to ad-hoc data."""
        # Extract CBF columns from user mapping
        user_mapping_cbf_cols = {cbf_col for _, cbf_col in user_mapping}

        # Only require columns that are both in default CBF mapping AND user mapping
        # This allows ad-hoc uploads to work with any column structure
        required_cols = self.required_cbf_columns & user_mapping_cbf_cols

        return apply_mapping_and_validate_with_fallback(df, user_mapping, required_cols, use_optimized=True)

    def _sanitize_adhoc_data(self, df: pl.DataFrame) -> pl.DataFrame:
        """Sanitize dirty user data for ad-hoc uploads."""
        return self.sanitizer.sanitize_dataframe(df)

    def _create_successful_adhoc_result(self, sanitized_df: pl.DataFrame, original_row_count: int) -> ProcessingResult:
        """Create successful processing result for ad-hoc data."""
        return ProcessingResult(
            processed_data=sanitized_df,
            original_row_count=original_row_count,
            processed_row_count=len(sanitized_df),
            processing_type="adhoc",
            cbf_columns_applied=list(sanitized_df.columns),
            sanitization_applied=True
        )

    def _create_failed_adhoc_result(self, original_row_count: int, error_message: str) -> ProcessingResult:
        """Create failed processing result for ad-hoc data."""
        return ProcessingResult(
            processed_data=pl.DataFrame(),
            original_row_count=original_row_count,
            processed_row_count=0,
            processing_type="adhoc",
            error_message=f"Ad-hoc processing failed: {error_message}"
        )
    
    def load_dataset_from_s3(
        self,
        dataset_name: str,
        provider: str, 
        month: str,
        stage: str
    ) -> pl.DataFrame:
        """
        Load dataset chunk from S3.
        
        Args:
            dataset_name: Name of the dataset
            provider: Provider name (AWS, Azure, GCP)
            month: Month in YYYY-MM format
            stage: Processing stage (usage, metadata, etc.)
            
        Returns:
            DataFrame loaded from S3
        """
        s3_key = f"czbillgen/datasets/{dataset_name}/{stage}/{provider}/{month}.parquet"
        return self.dataset_loader.s3_loader.read_parquet_from_s3(s3_key)
    
    def load_adhoc_dataset_from_csv(self, csv_path: Path) -> pl.DataFrame:
        """
        Load ad-hoc dataset from CSV file.
        
        Args:
            csv_path: Path to CSV file
            
        Returns:
            DataFrame loaded from CSV
        """
        return self.dataset_loader.load_adhoc_dataset(csv_path)
    
    def load_user_mapping(self, mapping_path: Path) -> List[Tuple[str, str]]:
        """
        Load user-defined column mapping from file.
        
        Args:
            mapping_path: Path to mapping file
            
        Returns:
            List of (source_column, target_column) tuples
        """
        return self.dataset_loader.load_user_mapping(mapping_path)
    
    def validate_source_columns(
        self,
        df: pl.DataFrame,
        cbf_config: CBFProcessingConfig
    ) -> List[str]:
        """
        Validate that source DataFrame has required columns.
        
        Args:
            df: Source DataFrame to validate
            cbf_config: CBF processing configuration
            
        Returns:
            List of missing required columns
        """
        required_sources = self.get_required_source_columns(cbf_config)
        missing_columns = required_sources - set(df.columns)
        return sorted(missing_columns)
    
    def get_required_source_columns(self, cbf_config: CBFProcessingConfig) -> List[str]:
        """
        Get the source columns required for CBF mapping.
        
        Args:
            cbf_config: CBF processing configuration
            
        Returns:
            List of required source column names
        """
        # Use configured required columns or fall back to mapping-based detection
        if cbf_config.required_columns:
            return cbf_config.required_columns
        
        # Determine required columns from CBF mapping
        required_sources = set()
        
        for source, cbf in self.mapping:
            if cbf in self.required_cbf_columns:
                required_sources.add(source)
        
        return sorted(required_sources)