"""Core usage calculation logic with progressive fallback matching.

This module implements the usage calculation algorithm adapted from V1 reference
implementation with modern Polars optimizations and unified matching keys.
"""

from typing import Dict, List, Tuple
import polars as pl
from rich.console import Console

console = Console()

# Unified matching columns for all providers (no provider-specific logic)
MATCH_COLUMNS = [
    'cloud_provider', 'service', 'usage_family', 
    'resource_type', 'operation', 'usage_type'
]

# Progressive fallback strategy (proven from V1)
FALLBACK_LEVELS = [
    [],                                          # Level 0: Exact match on all 6
    ['usage_type'],                             # Level 1: Drop usage_type
    ['operation'],                              # Level 2: Drop operation
    ['usage_type', 'operation'],                # Level 3: Drop both
    ['usage_type', 'operation', 'resource_type'] # Level 4: Drop 3 columns
]

# Required columns in reference data
REQUIRED_REFERENCE_COLUMNS = [
    # Matching columns
    'cloud_provider', 'service', 'usage_family', 
    'resource_type', 'operation', 'usage_type',
    
    # Pricing columns  
    'cost_per_unit',    # Unit price for calculation
    'pricing_unit',     # Unit of measurement
    'has_pricing'       # Boolean flag for pricing availability
]


def normalize_value(val) -> str:
    """Normalize values for fuzzy matching.
    
    Args:
        val: Value to normalize
        
    Returns:
        Normalized string value (lowercase, stripped)
    """
    if val is None or (hasattr(val, 'is_nan') and val.is_nan()):
        return ""
    return str(val).strip().lower()


def normalize_reference_data(reference_df: pl.DataFrame) -> pl.DataFrame:
    """Add normalized columns to reference data for matching.
    
    Also ensures original columns have consistent string types for joins.
    
    Args:
        reference_df: Reference DataFrame with pricing data
        
    Returns:
        DataFrame with additional normalized columns for fuzzy matching and consistent types
    """
    # First ensure all match columns are string type and handle nulls
    original_col_exprs = []
    normalized_exprs = []
    
    for col in MATCH_COLUMNS:
        if col in reference_df.columns:
            # Fix original column to be consistent string type
            original_col_exprs.append(
                pl.col(col).cast(pl.Utf8).fill_null("").alias(col)
            )
            # Create normalized column for fuzzy matching
            normalized_exprs.append(
                pl.col(col)
                .cast(pl.Utf8)
                .fill_null("")
                .str.strip()
                .str.to_lowercase()
                .alias(f"{col}_normalized")
            )
    
    return reference_df.with_columns(original_col_exprs + normalized_exprs)


def find_pricing_match(
    billing_row: dict,
    reference_data: pl.DataFrame
) -> Tuple[float, str]:
    """Find pricing match using progressive fallback.
    
    Tries increasingly permissive matching until a match is found:
    Level 0: cloud_provider + service + usage_family + resource_type + operation + usage_type
    Level 1: cloud_provider + service + usage_family + resource_type + operation
    Level 2: cloud_provider + service + usage_family + resource_type + usage_type
    Level 3: cloud_provider + service + usage_family + resource_type
    Level 4: cloud_provider + service + usage_family
    
    Args:
        billing_row: Dictionary with billing record values
        reference_data: Normalized reference DataFrame
        
    Returns:
        (unit_price, pricing_unit) tuple, (0.0, "") if no match found
    """
    
    for level, drop_columns in enumerate(FALLBACK_LEVELS):
        # Build matching condition for this level
        match_columns = [col for col in MATCH_COLUMNS if col not in drop_columns]
        
        # Apply normalization for fuzzy matching
        conditions = []
        for col in match_columns:
            normalized_val = normalize_value(billing_row.get(col))
            if normalized_val:  # Only add condition if value is not empty
                conditions.append(pl.col(f"{col}_normalized") == normalized_val)
        
        # Find matches at this level
        if conditions:
            combined_condition = conditions[0]
            for condition in conditions[1:]:
                combined_condition = combined_condition & condition
                
            matches = reference_data.filter(combined_condition)
            
            if len(matches) > 0:
                match_row = matches.row(0, named=True)
                try:
                    unit_price = float(match_row.get("cost_per_unit", 0.0))
                    pricing_unit = str(match_row.get("pricing_unit", ""))
                    return unit_price, pricing_unit
                except (ValueError, TypeError):
                    # Skip invalid pricing data
                    continue
    
    return 0.0, ""  # No match found at any level


def calculate_usage_for_batch(
    unique_combinations: pl.DataFrame,
    reference_data: pl.DataFrame
) -> pl.DataFrame:
    """Calculate usage amounts for unique billing combinations.
    
    Process unique combinations to avoid redundant calculations,
    then merge results back to full dataset.
    
    Args:
        unique_combinations: DataFrame with unique billing combinations
        reference_data: Normalized reference DataFrame
        
    Returns:
        DataFrame with pricing information for each unique combination
    """
    
    pricing_results = []
    
    # Process each unique combination
    for row in unique_combinations.iter_rows(named=True):
        unit_price, pricing_unit = find_pricing_match(row, reference_data)
        
        # Store result with original values
        result_row = {col: row[col] for col in MATCH_COLUMNS}
        result_row.update({
            "unit_price": unit_price,
            "pricing_unit": pricing_unit,
        })
        pricing_results.append(result_row)
    
    # Convert to DataFrame
    pricing_df = pl.DataFrame(pricing_results)
    
    return pricing_df


def apply_usage_calculation(
    billing_data: pl.DataFrame,
    pricing_lookup: pl.DataFrame
) -> pl.DataFrame:
    """Apply usage calculation to billing data.
    
    Merges pricing information and calculates:
    usage_amount = cost / unit_price (as string for CloudZero API)
    
    Args:
        billing_data: Original billing DataFrame
        pricing_lookup: DataFrame with pricing information
        
    Returns:
        Enhanced DataFrame with pricing_unit and usage_amount columns
    """
    
    # Ensure consistent data types for join columns
    # Convert all match columns to string and fill nulls with empty strings
    billing_normalized = billing_data.with_columns([
        pl.col(col).cast(pl.Utf8).fill_null("").alias(col) for col in MATCH_COLUMNS
    ])
    
    pricing_normalized = pricing_lookup.with_columns([
        pl.col(col).cast(pl.Utf8).fill_null("").alias(col) for col in MATCH_COLUMNS
    ])
    
    # Merge pricing information with normalized data
    enhanced_data = billing_normalized.join(
        pricing_normalized.select(MATCH_COLUMNS + ["unit_price", "pricing_unit"]),
        on=MATCH_COLUMNS,
        how="left"
    )
    
    # Calculate usage amounts
    enhanced_data = enhanced_data.with_columns([
        # Calculate usage_amount where unit_price > 0
        pl.when(pl.col("unit_price") > 0)
        .then((pl.col("cost") / pl.col("unit_price")).cast(pl.Utf8))
        .otherwise(pl.lit("0.0"))
        .alias("usage_amount"),
        
        # Fill null pricing_unit with empty string
        pl.col("pricing_unit").fill_null("").alias("pricing_unit")
    ])
    
    return enhanced_data


def calculate_usage_amounts(
    billing_data: pl.DataFrame,
    reference_data: pl.DataFrame
) -> pl.DataFrame:
    """Calculate usage amounts using progressive fallback matching.
    
    Main entry point for usage calculation algorithm.
    
    Algorithm:
    1. Validate input data and reference data
    2. Normalize reference data for fuzzy matching
    3. Get unique combinations of matching columns from billing data
    4. For each combination, try progressive fallback levels
    5. Find first match in reference data at any level
    6. Calculate usage_amount = cost / unit_price
    7. Merge results back to original billing data
    
    Args:
        billing_data: DataFrame with billing data including cost column
        reference_data: DataFrame with reference pricing data
        
    Returns:
        Enhanced DataFrame with pricing_unit and usage_amount columns added
    """
    
    # Validate input data
    if billing_data.is_empty():
        console.print("⚠️  Empty billing data provided")
        return billing_data.with_columns([
            pl.lit("").alias("pricing_unit"),
            pl.lit("0.0").alias("usage_amount")
        ])
    
    # Validate reference data
    if reference_data.is_empty():
        console.print("⚠️  Empty reference data provided")
        return billing_data.with_columns([
            pl.lit("").alias("pricing_unit"),
            pl.lit("0.0").alias("usage_amount")
        ])
    
    # Check required columns in billing data
    billing_required = ['cost'] + MATCH_COLUMNS
    missing_billing_cols = [col for col in billing_required if col not in billing_data.columns]
    if missing_billing_cols:
        raise ValueError(f"Missing required columns in billing data: {missing_billing_cols}")
    
    # Check required columns in reference data
    missing_ref_cols = [col for col in REQUIRED_REFERENCE_COLUMNS if col not in reference_data.columns]
    if missing_ref_cols:
        raise ValueError(f"Missing required columns in reference data: {missing_ref_cols}")
    
    # Normalize reference data for fuzzy matching
    reference_normalized = normalize_reference_data(reference_data)
    
    # Get unique combinations of matching columns from billing data
    unique_combinations = billing_data.select(MATCH_COLUMNS).unique()
    
    console.print(f"🔍 Processing {len(unique_combinations):,} unique billing combinations")
    
    # Calculate pricing for unique combinations
    pricing_lookup = calculate_usage_for_batch(unique_combinations, reference_normalized)
    
    # Apply usage calculation to full dataset
    enhanced_data = apply_usage_calculation(billing_data, pricing_lookup)
    
    # Log results
    total_records = len(enhanced_data)
    usage_records = len(enhanced_data.filter(pl.col("usage_amount") != "0.0"))
    coverage_pct = (usage_records / total_records * 100) if total_records > 0 else 0
    
    console.print(f"✅ Usage calculation complete: {usage_records:,}/{total_records:,} records ({coverage_pct:.1f}%) have usage amounts")
    
    return enhanced_data