"""Schema harmonization for metadata stage data to match usage stage format."""

import polars as pl
from typing import Dict, List
from src.log_context.core import get_logger


class SchemaHarmonizer:
    """Handles schema differences between usage and metadata stage data.
    
    Ensures metadata stage data has the same columns as usage stage data
    before applying discount processors.
    """
    
    # Usage stage columns that metadata stage might be missing
    USAGE_COLUMNS = {
        'usage_amount': '0',      # Default value for usage amount
        'pricing_unit': ''        # Default value for pricing unit
    }
    
    def __init__(self):
        self.logger = get_logger(__name__)
    
    def harmonize_metadata_schema(self, metadata_data: pl.DataFrame) -> pl.DataFrame:
        """
        Add missing usage columns to metadata stage data with appropriate defaults.
        
        Args:
            metadata_data: DataFrame from metadata stage
            
        Returns:
            DataFrame with harmonized schema matching usage stage
        """
        result = metadata_data.clone()
        schema_additions = []
        added_columns = []
        
        for column, default_value in self.USAGE_COLUMNS.items():
            if column not in result.columns:
                schema_additions.append(
                    pl.lit(default_value).alias(column)
                )
                added_columns.append(column)
        
        if schema_additions:
            result = result.with_columns(schema_additions)
            
            self.logger.debug(
                "Added missing usage columns to metadata stage data",
                added_columns=added_columns,
                original_columns=len(metadata_data.columns),
                harmonized_columns=len(result.columns)
            )
        else:
            self.logger.debug("No schema harmonization needed - all usage columns present")
        
        return result
    
    def validate_harmonized_schema(self, original_data: pl.DataFrame, harmonized_data: pl.DataFrame) -> bool:
        """
        Validate that schema harmonization was applied correctly.
        
        Args:
            original_data: Original metadata stage data
            harmonized_data: Data after harmonization
            
        Returns:
            True if harmonization is valid
        """
        # Check row count preservation
        if len(original_data) != len(harmonized_data):
            self.logger.error(
                "Schema harmonization changed row count",
                original_rows=len(original_data),
                harmonized_rows=len(harmonized_data)
            )
            return False
        
        # Check that all usage columns are present
        missing_columns = []
        for column in self.USAGE_COLUMNS.keys():
            if column not in harmonized_data.columns:
                missing_columns.append(column)
        
        if missing_columns:
            self.logger.error(
                "Schema harmonization failed to add required columns",
                missing_columns=missing_columns
            )
            return False
        
        # Check that original columns are preserved
        original_columns = set(original_data.columns)
        harmonized_columns = set(harmonized_data.columns)
        lost_columns = original_columns - harmonized_columns
        
        if lost_columns:
            self.logger.error(
                "Schema harmonization lost original columns",
                lost_columns=list(lost_columns)
            )
            return False
        
        # Validate default values for added columns
        for column, expected_default in self.USAGE_COLUMNS.items():
            if column not in original_data.columns:
                # This column was added, check default values
                unique_values = harmonized_data[column].unique().to_list()
                if len(unique_values) != 1 or unique_values[0] != expected_default:
                    self.logger.warning(
                        f"Schema harmonization: unexpected values in added column {column}",
                        expected_default=expected_default,
                        actual_values=unique_values
                    )
        
        self.logger.debug("Schema harmonization validation passed")
        return True
    
    def get_schema_differences(self, usage_data: pl.DataFrame, metadata_data: pl.DataFrame) -> Dict[str, List[str]]:
        """
        Analyze schema differences between usage and metadata stage data.
        
        Args:
            usage_data: Sample usage stage data
            metadata_data: Sample metadata stage data
            
        Returns:
            Dictionary with schema difference analysis
        """
        usage_columns = set(usage_data.columns)
        metadata_columns = set(metadata_data.columns)
        
        differences = {
            "usage_only": list(usage_columns - metadata_columns),
            "metadata_only": list(metadata_columns - usage_columns),
            "common": list(usage_columns & metadata_columns),
            "usage_columns_count": len(usage_columns),
            "metadata_columns_count": len(metadata_columns)
        }
        
        self.logger.info(
            "Schema difference analysis",
            usage_only_count=len(differences["usage_only"]),
            metadata_only_count=len(differences["metadata_only"]),
            common_count=len(differences["common"])
        )
        
        return differences
    
    def ensure_discount_columns(self, data: pl.DataFrame) -> pl.DataFrame:
        """
        Ensure discount columns are present in the data.
        
        Args:
            data: Input DataFrame
            
        Returns:
            DataFrame with discount columns added if missing
        """
        result = data.clone()
        discount_columns = {
            'discount_amount': 0.0,
            'discount_type': ''
        }
        
        schema_additions = []
        for column, default_value in discount_columns.items():
            if column not in result.columns:
                schema_additions.append(
                    pl.lit(default_value).alias(column)
                )
        
        if schema_additions:
            result = result.with_columns(schema_additions)
            self.logger.debug(
                "Added missing discount columns",
                added_columns=list(discount_columns.keys())
            )
        
        return result