"""
Column name mapping utilities for K8s processing.

Handles mapping between different column naming conventions used in
the generate stage output and the expected K8s processing column names.
"""

import polars as pl
from typing import Dict, Optional

from src.log_context import get_logger


class ColumnMapper:
    """
    Maps column names between different naming conventions.
    
    The generate stage uses lowercase column names (resource_id, cost, service)
    while some K8s processing expects PascalCase (ResourceID, Cost, Service).
    """
    
    # Mapping from expected K8s column names to possible actual column names
    COLUMN_MAPPINGS = {
        "ResourceID": ["resource_id", "ResourceID"],
        "Cost": ["cost", "Cost"],
        "Service": ["service", "Service"],
        "ResourceType": ["resource_type", "ResourceType"]
    }
    
    def __init__(self):
        self.logger = get_logger(__name__)
    
    def get_column_name(self, df: pl.DataFrame, expected_name: str) -> Optional[str]:
        """
        Get the actual column name in the DataFrame for an expected column.
        
        Args:
            df: DataFrame to check
            expected_name: Expected column name (e.g., "ResourceID")
            
        Returns:
            Actual column name if found, None otherwise
        """
        possible_names = self.COLUMN_MAPPINGS.get(expected_name, [expected_name])
        
        for name in possible_names:
            if name in df.columns:
                return name
        
        return None
    
    def validate_required_columns(self, df: pl.DataFrame, required_columns: list) -> Dict[str, str]:
        """
        Validate that all required columns exist and return the mapping.
        
        Args:
            df: DataFrame to validate
            required_columns: List of expected column names
            
        Returns:
            Dictionary mapping expected_name -> actual_name
            
        Raises:
            ValueError: If any required column is missing
        """
        column_mapping = {}
        missing_columns = []
        
        for expected_name in required_columns:
            actual_name = self.get_column_name(df, expected_name)
            if actual_name is None:
                missing_columns.append(expected_name)
            else:
                column_mapping[expected_name] = actual_name
        
        if missing_columns:
            self.logger.error(
                "Required columns missing from DataFrame",
                missing_columns=missing_columns,
                available_columns=df.columns
            )
            raise ValueError(f"Missing required columns: {missing_columns}")
        
        return column_mapping
    
    def standardize_column_names(self, df: pl.DataFrame) -> pl.DataFrame:
        """
        Standardize column names to K8s expected format.
        
        Args:
            df: DataFrame with potentially different column names
            
        Returns:
            DataFrame with standardized column names
        """
        rename_mapping = {}
        
        for expected_name, possible_names in self.COLUMN_MAPPINGS.items():
            for possible_name in possible_names:
                if possible_name in df.columns and possible_name != expected_name:
                    rename_mapping[possible_name] = expected_name
                    break
        
        if rename_mapping:
            self.logger.debug(
                "Standardizing column names",
                rename_mapping=rename_mapping
            )
            df = df.rename(rename_mapping)
        
        return df


# Global instance for easy access
column_mapper = ColumnMapper()