"""
Add CloudZero export files to the unified reference dataset.
"""

import typer
from pathlib import Path
from typing import Optional
import polars as pl
from ..common import BaseCommand

app = typer.Typer(help="Add export files to reference dataset")


class DataAddCommand(BaseCommand):
    """Add export files to unified reference dataset."""
    
    def add_export_file(
        self,
        export_file: str,
        reference_file: str = "data/reference.parquet",
        verbose: bool = False
    ):
        """Add CloudZero export file to the reference dataset."""
        try:
            export_path = Path(export_file)
            reference_path = Path(reference_file)
            
            self.validate_file_exists(export_path, "export file")
            
            if verbose:
                self.console.print(f"[blue]Adding export: {export_path}[/blue]")
                self.console.print(f"[blue]To reference: {reference_path}[/blue]")
            
            # Load export file
            self.console.print("📥 Loading export file...")
            export_df = self._load_export_file(export_path)
            
            # Load existing reference data
            if reference_path.exists():
                self.console.print("📥 Loading existing reference data...")
                reference_df = pl.read_parquet(reference_path)
            else:
                self.console.print("📝 Creating new reference dataset...")
                # Create empty reference with proper schema
                reference_df = pl.DataFrame(schema={
                    "cloud_provider": pl.Utf8,
                    "service": pl.Utf8,
                    "usage_family": pl.Utf8,
                    "resource_type": pl.Utf8,
                    "operation": pl.Utf8,
                    "usage_type": pl.Utf8,
                    "instance_type": pl.Utf8,
                    "pricing_unit": pl.Utf8,
                    "total_usage": pl.Float64,
                    "total_cost": pl.Float64,
                    "cost_per_unit": pl.Float64,
                    "has_pricing": pl.Boolean,
                    "record_count": pl.Int32
                })
            
            # Filter export to only new combinations, then merge
            self.console.print("🔍 Analyzing combinations...")
            filtered_export_df, new_count, existing_count = self._filter_new_combinations(reference_df, export_df)
            
            if len(filtered_export_df) > 0:
                self.console.print("🔗 Merging new combinations...")
                merged_df = self._merge_datasets(reference_df, filtered_export_df)
            else:
                self.console.print("✨ No new combinations to add")
                merged_df = reference_df
            
            # Save updated reference
            self.console.print(f"💾 Saving to: {reference_path}")
            reference_path.parent.mkdir(parents=True, exist_ok=True)
            merged_df.write_parquet(reference_path)
            
            # Display comprehensive summary
            self._display_comprehensive_summary(export_df, filtered_export_df, reference_df, merged_df, new_count, existing_count)
            
        except Exception as e:
            self.handle_exception(e, "adding export file")
    
    def _load_export_file(self, export_path: Path) -> pl.DataFrame:
        """Load and normalize export file."""
        # Load export CSV
        export_df = pl.read_csv(export_path)
        
        # Drop unnamed columns (often index columns from exports)
        unnamed_cols = [col for col in export_df.columns if col.strip() == "" or col.startswith("Unnamed")]
        if unnamed_cols:
            export_df = export_df.drop(unnamed_cols)
        
        # Normalize column names (exports may have different formats)
        column_mapping = {
            "Cloud Provider": "cloud_provider",
            "Service": "service", 
            "Usage Family": "usage_family",
            "Resource Type": "resource_type",
            "Operation": "operation",
            "Usage Type": "usage_type",
            "Instance Type": "instance_type",
            "Pricing Unit": "pricing_unit",
            "Usage Amount": "total_usage",
            "Real Cost": "total_cost",
            # Handle variations
            "CloudProvider": "cloud_provider",
            "UsageFamily": "usage_family",
            "ResourceType": "resource_type",
            "UsageType": "usage_type",
            "InstanceType": "instance_type",
            "PricingUnit": "pricing_unit"
        }
        
        # Rename columns
        for old_name, new_name in column_mapping.items():
            if old_name in export_df.columns:
                export_df = export_df.rename({old_name: new_name})
        
        # Ensure required columns exist
        required_cols = [
            "cloud_provider", "service", "usage_family", 
            "resource_type", "operation", "usage_type"
        ]
        
        for col in required_cols:
            if col not in export_df.columns:
                raise ValueError(f"Required column '{col}' not found in export file")
        
        # Add missing optional columns
        if "instance_type" not in export_df.columns:
            export_df = export_df.with_columns(pl.lit(None, dtype=pl.Utf8).alias("instance_type"))
        if "pricing_unit" not in export_df.columns:
            export_df = export_df.with_columns(pl.lit(None, dtype=pl.Utf8).alias("pricing_unit"))
        if "total_usage" not in export_df.columns:
            export_df = export_df.with_columns(pl.lit(None, dtype=pl.Float64).alias("total_usage"))
        if "total_cost" not in export_df.columns:
            export_df = export_df.with_columns(pl.lit(None, dtype=pl.Float64).alias("total_cost"))
        
        # Clean resource_type column first (always safe as string operation)
        export_df = export_df.with_columns([
            # Clean resource_type column (remove "EC2: " prefix from CloudZero exports)
            pl.when(pl.col("resource_type").is_null())
            .then(pl.lit(None))
            .otherwise(
                pl.col("resource_type")
                .str.replace(r"^[^:]+:\s*", "")  # Remove "prefix: " pattern
                .str.strip_chars()  # Remove any extra whitespace
            ).alias("resource_type")
        ])
        
        # Handle cost and usage columns based on their data types
        cost_dtype = export_df["total_cost"].dtype
        usage_dtype = export_df["total_usage"].dtype
        
        cost_expr = (
            pl.col("total_cost").cast(pl.Float64) if cost_dtype.is_numeric()
            else pl.when(pl.col("total_cost").is_null() | (pl.col("total_cost") == ""))
                 .then(pl.lit(0.0))
                 .otherwise(
                     pl.col("total_cost")
                     .str.replace_all(r"[\$,]", "")  # Remove $ and commas
                     .cast(pl.Float64, strict=False)
                 )
        ).alias("total_cost")
        
        usage_expr = (
            pl.col("total_usage").cast(pl.Float64) if usage_dtype.is_numeric()
            else pl.when(pl.col("total_usage").is_null() | (pl.col("total_usage") == ""))
                 .then(pl.lit(0.0))
                 .otherwise(
                     pl.col("total_usage")
                     .str.replace_all(r"[,]", "")  # Remove commas
                     .cast(pl.Float64, strict=False)
                 )
        ).alias("total_usage")
        
        export_df = export_df.with_columns([cost_expr, usage_expr])
        
        # Calculate cost per unit and other fields
        export_df = export_df.with_columns([
            pl.when((pl.col("total_usage") > 0) & (pl.col("total_cost") > 0))
            .then(pl.col("total_cost") / pl.col("total_usage"))
            .otherwise(pl.lit(None))
            .alias("cost_per_unit"),
            
            (pl.col("total_cost") > 0).alias("has_pricing"),
            pl.lit(1).alias("record_count")
        ])
        
        # Ensure column order matches reference schema
        final_columns = [
            "cloud_provider", "service", "usage_family", "resource_type",
            "operation", "usage_type", "instance_type", "pricing_unit",
            "total_usage", "total_cost", "cost_per_unit", "has_pricing", "record_count"
        ]
        
        export_df = export_df.select(final_columns)
        
        return export_df
    
    def _filter_new_combinations(self, reference_df: pl.DataFrame, export_df: pl.DataFrame):
        """Filter export data to only include combinations not in reference data."""
        
        # Key fields for uniqueness analysis
        key_fields = [
            "cloud_provider", "service", "usage_family", 
            "resource_type", "operation", "usage_type"
        ]
        
        if reference_df.is_empty():
            # No existing reference data, all export combinations are new
            export_unique = export_df.select(key_fields).unique()
            return export_df, export_unique.height, 0
        
        # Get unique combinations from each dataset
        export_combinations = export_df.select(key_fields).unique()
        existing_combinations = reference_df.select(key_fields).unique()
        
        # Handle null values for join compatibility
        # Replace None with a unique string temporarily for join purposes
        null_placeholder = "<<NULL_PLACEHOLDER>>"
        
        export_for_join = export_combinations.with_columns([
            pl.col("operation").fill_null(null_placeholder),
            pl.col("usage_type").fill_null(null_placeholder)
        ])
        
        existing_for_join = existing_combinations.with_columns([
            pl.col("operation").fill_null(null_placeholder),
            pl.col("usage_type").fill_null(null_placeholder)
        ])
        
        # Find which export combinations are truly new
        export_status = export_for_join.join(
            existing_for_join.with_columns(pl.lit(True).alias("exists_in_reference")),
            on=key_fields,
            how="left",
            coalesce=True
        )
        
        # Convert back to original null values for final result
        export_status = export_status.with_columns([
            pl.when(pl.col("operation") == null_placeholder).then(None).otherwise(pl.col("operation")).alias("operation"),
            pl.when(pl.col("usage_type") == null_placeholder).then(None).otherwise(pl.col("usage_type")).alias("usage_type")
        ])
        
        # Separate new vs existing
        new_combinations = export_status.filter(pl.col("exists_in_reference").is_null()).select(key_fields)
        existing_found = export_status.filter(pl.col("exists_in_reference").is_not_null()).select(key_fields)
        
        # Filter export data to only include records with new combinations
        if new_combinations.is_empty():
            # No new combinations, return empty DataFrame with same schema
            empty_export = export_df.head(0)  # Preserves schema
            return empty_export, 0, existing_found.height
        
        # Filter export data to only rows that match new combinations
        # Use the same null placeholder approach for the final filtering
        export_for_final_join = export_df.with_columns([
            pl.col("operation").fill_null(null_placeholder),
            pl.col("usage_type").fill_null(null_placeholder)
        ])
        
        new_for_final_join = new_combinations.with_columns([
            pl.col("operation").fill_null(null_placeholder),
            pl.col("usage_type").fill_null(null_placeholder)
        ])
        
        filtered_export = export_for_final_join.join(
            new_for_final_join.with_columns(pl.lit(True).alias("is_new")),
            on=key_fields,
            how="inner"
        ).drop("is_new")
        
        # Convert back to original null values
        filtered_export = filtered_export.with_columns([
            pl.when(pl.col("operation") == null_placeholder).then(None).otherwise(pl.col("operation")).alias("operation"),
            pl.when(pl.col("usage_type") == null_placeholder).then(None).otherwise(pl.col("usage_type")).alias("usage_type")
        ])
        
        return filtered_export, new_combinations.height, existing_found.height
    
    def _merge_datasets(self, reference_df: pl.DataFrame, export_df: pl.DataFrame) -> pl.DataFrame:
        """Merge export data into reference dataset with deduplication."""
        
        # Metadata combination columns for grouping
        metadata_cols = [
            "cloud_provider", "service", "usage_family",
            "resource_type", "operation", "usage_type", "instance_type"
        ]
        
        # Combine datasets
        combined_df = pl.concat([reference_df, export_df])
        
        # Group by metadata combination and aggregate
        merged_df = combined_df.group_by(metadata_cols).agg([
            # Take first non-null pricing_unit
            pl.col("pricing_unit").filter(pl.col("pricing_unit").is_not_null()).first(),
            # Sum usage and costs
            pl.col("total_usage").sum(),
            pl.col("total_cost").sum(),
            # Average cost per unit (recalculated below)
            pl.col("cost_per_unit").filter(pl.col("cost_per_unit").is_not_null()).mean(),
            # Any record has pricing
            pl.col("has_pricing").any(),
            # Sum record counts
            pl.col("record_count").sum()
        ])
        
        # Recalculate cost per unit from aggregated values
        merged_df = merged_df.with_columns([
            pl.when(pl.col("total_usage") > 0)
            .then(pl.col("total_cost") / pl.col("total_usage"))
            .otherwise(pl.col("cost_per_unit"))
            .alias("cost_per_unit")
        ])
        
        # Update has_pricing based on aggregated cost
        merged_df = merged_df.with_columns([
            (pl.col("total_cost").is_not_null() & (pl.col("total_cost") > 0)).alias("has_pricing")
        ])
        
        # Sort for consistent ordering
        merged_df = merged_df.sort(["cloud_provider", "service", "usage_family", "resource_type"])
        
        return merged_df
    
    def _display_comprehensive_summary(self, original_export_df: pl.DataFrame, filtered_export_df: pl.DataFrame, reference_df_before: pl.DataFrame, merged_df: pl.DataFrame, new_count: int, existing_count: int):
        """Display comprehensive addition summary with full transparency."""
        
        # Export analysis (original export data)
        export_records = original_export_df.height
        export_with_pricing = original_export_df.filter(pl.col("has_pricing")).height if not original_export_df.is_empty() else 0
        export_unique_count = new_count + existing_count
        
        # Filtered export analysis (what actually got merged)
        filtered_records = filtered_export_df.height
        
        # Reference data before analysis
        before_total = reference_df_before.height
        before_with_pricing = reference_df_before.filter(pl.col("has_pricing")).height
        before_coverage = (before_with_pricing / before_total * 100) if before_total > 0 else 0
        
        # Reference data after analysis
        after_total = merged_df.height
        after_with_pricing = merged_df.filter(pl.col("has_pricing")).height
        after_coverage = (after_with_pricing / after_total * 100) if after_total > 0 else 0
        
        # Net changes
        net_combinations = after_total - before_total
        net_with_pricing = after_with_pricing - before_with_pricing
        coverage_change = after_coverage - before_coverage
        
        # Display comprehensive summary
        self.console.print(f"\n✅ Export processed successfully!")
        
        self.console.print(f"   📥 Export Analysis:")
        self.console.print(f"      • Total records: {export_records:,}")
        self.console.print(f"      • Unique combinations: {export_unique_count:,}")
        self.console.print(f"      • Records with pricing: {export_with_pricing:,}")
        
        if export_records != export_unique_count:
            avg_aggregation = export_records / export_unique_count
            self.console.print(f"      • Average records per combination: {avg_aggregation:.1f}")
        
        self.console.print(f"\n   🔄 Import Impact:")
        self.console.print(f"      • New combinations: {new_count:,}")
        self.console.print(f"      • Existing combinations: {existing_count:,}")
        self.console.print(f"      • Records merged: {filtered_records:,} (only new combinations)")
        
        # Show service breakdown for new combinations if reasonable count
        if new_count > 0 and new_count <= 10 and not filtered_export_df.is_empty():
            services = filtered_export_df.select("service").unique().to_series().to_list()
            if len(services) <= 5:
                self.console.print(f"      • New services: {', '.join(services)}")
        
        self.console.print(f"\n   📊 Reference Data Changes:")
        self.console.print(f"      • Before: {before_total:,} combinations ({before_with_pricing:,} with pricing)")
        self.console.print(f"      • After:  {after_total:,} combinations ({after_with_pricing:,} with pricing)")
        self.console.print(f"      • Net:    {net_combinations:+,} combinations ({net_with_pricing:+,} with pricing)")
        
        self.console.print(f"\n   📈 Coverage:")
        self.console.print(f"      • Before: {before_coverage:.1f}%")
        self.console.print(f"      • After:  {after_coverage:.1f}%")
        self.console.print(f"      • Change: {coverage_change:+.1f}%")
        
        # Data quality insights
        self._display_quality_insights(original_export_df, export_unique_count, export_records)
    
    def _display_quality_insights(self, export_df: pl.DataFrame, unique_count: int, total_records: int):
        """Display data quality insights and warnings."""
        
        insights = []
        
        # Perfect pricing coverage check
        with_pricing = export_df.filter(pl.col("has_pricing")).height
        if with_pricing == total_records and total_records > 100:
            insights.append("Perfect pricing coverage (100%) - verify data source")
        
        # High aggregation check
        if total_records > unique_count:
            avg_aggregation = total_records / unique_count
            if avg_aggregation > 100:
                insights.append(f"High aggregation detected (avg: {avg_aggregation:.0f} records/combination)")
            
            # Find highest aggregation combination
            key_fields = ["cloud_provider", "service", "usage_family", "resource_type", "operation", "usage_type"]
            aggregation_analysis = export_df.select(key_fields).group_by(key_fields).agg([
                pl.len().alias('record_count')
            ]).sort('record_count', descending=True)
            
            highest = aggregation_analysis.head(1)
            if len(highest) > 0:
                max_count = highest.select('record_count').item()
                if max_count > 1000:
                    service = highest.select('service').item()
                    insights.append(f"Highest aggregation: {max_count:,} records ({service})")
        
        # Low diversity check  
        if unique_count < 10 and total_records > 1000:
            insights.append(f"Low service diversity ({unique_count} types for {total_records:,} records)")
        
        # Display insights if any
        if insights:
            self.console.print(f"\n   ⚠️  Quality Notes:")
            for insight in insights:
                self.console.print(f"      • {insight}")


# Create command instance
add_cmd = DataAddCommand()


@app.command()
def export(
    export_file: str = typer.Argument(..., help="Path to CloudZero export CSV file"),
    reference_file: str = typer.Option(
        "data/reference.parquet", 
        help="Path to reference dataset"
    ),
    verbose: bool = typer.Option(False, "--verbose", "-v", help="Enable verbose output")
):
    """Add CloudZero export file to the reference dataset."""
    add_cmd.add_export_file(export_file, reference_file, verbose)