"""
Data sanitization for ad-hoc upload processing.

Extracted from the existing upload package sanitizer.py to clean
dirty user-provided CSV data before CBF processing.
"""

import re
from collections import defaultdict
from typing import Dict, List

import polars as pl


class DataSanitizer:
    """
    Data sanitizer for cleaning ad-hoc user-provided datasets.
    
    Extracted from the existing upload package to provide consistent
    data cleaning for user CSV files that contain dirty data.
    """
    
    # Cost fields that need cleaning and validation
    COST_FIELDS = [
        "cost/cost",
        "cost/discounted_cost", 
        "cost/amortized_cost",
        "cost/discounted_amortized_cost",
        "cost/on_demand_cost",
    ]
    
    # Time fields that need timestamp standardization
    TIME_FIELDS = [
        "time/usage_start",
        "time/usage_end",
    ]
    
    def __init__(self):
        """Initialize sanitizer with tracking for invalid values."""
        self.invalid_counts: Dict[str, int] = defaultdict(int)
        self.invalid_examples: Dict[str, List[str]] = defaultdict(list)
    
    def sanitize_dataframe(self, df: pl.DataFrame) -> pl.DataFrame:
        """
        Sanitize user-provided data before CBF processing.
        
        Extracted from sanitizer.sanitize_dataframe() to clean:
        - Cost fields (remove currency symbols, standardize decimals)
        - Timestamp fields (convert to ISO 8601 format)
        - Null values (convert to empty strings)
        - Data types (ensure consistent string types)
        """
        # Replace NaNs and nulls
        df = df.fill_null("")
        
        # Clean cost fields
        for col in self.COST_FIELDS:
            if col in df.columns:
                df = df.with_columns(
                    pl.col(col)
                    .cast(pl.Utf8)
                    .map_elements(
                        lambda x: self._clean_cost_string(x, col), 
                        return_dtype=pl.Float64
                    )
                    .cast(pl.Utf8)
                    .alias(col)
                )
        
        # Report invalid cost values
        if self.invalid_counts:
            print("[!] Some cost values were invalid and replaced with '0.01':")
            for col, count in self.invalid_counts.items():
                print(f"    • {col}: {count} values replaced")
                for ex in self.invalid_examples[col]:
                    print(f"       - example: '{ex}'")
        
        # Standardize timestamp columns to ISO 8601
        for col in self.TIME_FIELDS:
            if col in df.columns:
                df = df.with_columns(
                    pl.col(col)
                    .str.strip_chars()
                    .str.to_datetime(strict=False)
                    .dt.strftime("%Y-%m-%dT%H:%M:%S")
                    .alias(col)
                )
        
        # Convert other columns to strings
        other_cols = [c for c in df.columns if c not in self.COST_FIELDS + self.TIME_FIELDS]
        df = df.with_columns(
            [pl.col(c).fill_null("").cast(pl.Utf8).alias(c) for c in other_cols]
        )
        
        return df
    
    def _clean_cost_string(self, value: str, col: str) -> float:
        """
        Clean cost string values.
        
        Extracted from sanitizer.clean_cost_string() to handle:
        - Currency symbols: '$10.99' -> 10.99
        - Comma separators: '1,000.50' -> 1000.50
        - Invalid formats: 'N/A' -> 0.01
        """
        try:
            raw = value.strip() if isinstance(value, str) else str(value)
            # Remove non-numeric characters except decimal and negative
            clean = re.sub(r"[^\d.\-]", "", raw)
            val = float(clean)
            # Ensure minimum 0.01, round to 2 decimals
            return max(0.01, round(val, 2))
        except Exception:
            # Track invalid values for reporting
            self.invalid_counts[col] += 1
            if len(self.invalid_examples[col]) < 3:
                self.invalid_examples[col].append(str(value))
            return 0.01