"""
Unified data processing with CBF mapping for upload operations.

Handles both generated and ad-hoc datasets with proper CloudZero Billing Format
(CBF) mapping, validation, and sanitization extracted from the existing upload package.
"""

import polars as pl
from pathlib import Path
from typing import List, Set, Tuple, Dict, Any, Optional

from .sanitizer import DataSanitizer
from .cbf_mapper import apply_mapping_and_validate, load_mapping_file


class CBFProcessor:
    """
    CloudZero Billing Format processor for upload operations.
    
    Uses the existing CBF mapping logic from pipeline.upload.cbf_mapper
    to provide consistent CBF processing for both generated and ad-hoc datasets.
    """
    
    def __init__(self, cbf_mapping_path: Optional[Path] = None):
        """Initialize CBF processor with optional mapping configuration."""
        self.cbf_mapping_path = cbf_mapping_path
        self.sanitizer = DataSanitizer()
        
        # Load mapping configuration
        self.mapping, self.required_cbf_columns = load_mapping_file(cbf_mapping_path)
    
    def process_generated_dataset(self, df: pl.DataFrame) -> pl.DataFrame:
        """
        Process generated dataset (already clean, no sanitization needed).
        
        Generated datasets from the pipeline are already properly formatted
        and validated, so they only need CBF mapping and API formatting.
        """
        # Apply CBF mapping and validation using existing logic
        cbf_df = apply_mapping_and_validate(df, self.mapping, self.required_cbf_columns)
        
        return cbf_df
    
    def process_adhoc_dataset(
        self,
        df: pl.DataFrame,
        user_mapping: List[Tuple[str, str]]
    ) -> pl.DataFrame:
        """
        Process ad-hoc dataset with sanitization and user mapping.

        Ad-hoc datasets require sanitization to clean dirty user data
        before CBF mapping and validation.
        """
        # Extract required columns from user mapping
        # The user mapping determines which CBF columns are required
        user_mapping_dict = {cbf_col for _, cbf_col in user_mapping}

        # Use actual CBF required columns from the default mapping
        # Only require columns that are both in default mapping AND user mapping
        required_cols = self.required_cbf_columns & user_mapping_dict

        # Apply user-defined mapping first
        mapped_df = apply_mapping_and_validate(df, user_mapping, required_cols)

        # CRITICAL: Sanitize dirty user data
        sanitized_df = self.sanitizer.sanitize_dataframe(mapped_df)

        return sanitized_df
    
    def get_required_source_columns(self) -> Set[str]:
        """
        Get the source columns required for mapping.
        
        Returns the source columns that must be present in the input data
        to satisfy the required CBF columns.
        """
        required_sources = set()
        
        for source, cbf in self.mapping:
            if cbf in self.required_cbf_columns:
                required_sources.add(source)
        
        return required_sources
    
    def validate_source_columns(self, df: pl.DataFrame) -> List[str]:
        """
        Validate that the source DataFrame has the required columns.
        
        Returns a list of missing required columns.
        """
        required_sources = self.get_required_source_columns()
        missing_columns = required_sources - set(df.columns)
        
        return sorted(missing_columns)


class DataProcessor:
    """
    Unified data processor for upload operations.
    
    Provides a single interface for processing both generated and ad-hoc
    datasets with appropriate CBF mapping and sanitization.
    """
    
    def __init__(self, cbf_mapping_path: Optional[str] = None):
        """Initialize data processor with CBF mapping configuration."""
        self.cbf_processor = CBFProcessor(cbf_mapping_path)
    
    def process_generated_dataset(self, df: pl.DataFrame) -> pl.DataFrame:
        """Process generated dataset for upload."""
        return self.cbf_processor.process_generated_dataset(df)
    
    def process_adhoc_dataset(
        self, 
        df: pl.DataFrame, 
        user_mapping: List[Tuple[str, str]]
    ) -> pl.DataFrame:
        """Process ad-hoc dataset with sanitization for upload."""
        return self.cbf_processor.process_adhoc_dataset(df, user_mapping)