"""Reference data loading with lazy loading and provider pre-filtering optimization.

This module provides optimized reference data loading following the design decisions:
- Use same reference.parquet as generate package
- Lazy loading with provider pre-filtering for performance
- Caching per provider for reuse across months
"""

from pathlib import Path
from typing import Dict, Optional
import polars as pl
from rich.console import Console

from ...log_context.core import get_logger

console = Console()
logger = get_logger(__name__)

# Default reference data path - same as generate package
DEFAULT_REFERENCE_PATH = Path.home() / ".czbg" / "cache" / "reference_data" / "reference.parquet"

# Reference data cache - keyed by provider
_reference_cache: Dict[str, pl.DataFrame] = {}


def validate_reference_data(reference_path: Path) -> bool:
    """Validate that reference file exists and has required structure.
    
    Args:
        reference_path: Path to unified reference parquet file
        
    Returns:
        True if valid, False otherwise
    """
    if not reference_path.exists():
        console.print(f"❌ Reference file not found: {reference_path}")
        return False

    try:
        # Use lazy loading to check schema without loading full file
        lazy_df = pl.scan_parquet(reference_path)
        schema = lazy_df.schema
        
        # Check required columns
        required_columns = [
            # Matching columns
            "cloud_provider", "service", "usage_family", 
            "resource_type", "operation", "usage_type",
            # Pricing columns
            "cost_per_unit", "pricing_unit", "has_pricing"
        ]
        
        missing_columns = [col for col in required_columns if col not in schema]
        if missing_columns:
            console.print(f"❌ Missing required columns in reference data: {missing_columns}")
            return False

        # Count records with pricing data using lazy evaluation
        total_records = lazy_df.select(pl.len()).collect().item()
        pricing_records = (
            lazy_df
            .filter(pl.col("has_pricing") == True)
            .select(pl.len())
            .collect()
            .item()
        )

        console.print(f"✅ Reference file validated: {pricing_records:,} pricing records from {total_records:,} total records")
        return True

    except Exception as e:
        console.print(f"❌ Error reading reference file: {e}")
        return False


def load_reference_data_for_provider(
    provider: str, 
    reference_path: Optional[Path] = None,
    use_cache: bool = True
) -> pl.DataFrame:
    """Load reference data filtered by provider using lazy loading.
    
    Uses provider pre-filtering for memory efficiency:
    - Only loads records for specific provider
    - Only loads records with has_pricing=true
    - Caches results per provider for reuse across months
    
    Args:
        provider: Cloud provider (AWS, GCP, Azure, etc.)
        reference_path: Optional custom reference data path
        use_cache: Whether to use caching (default True)
        
    Returns:
        DataFrame with reference data for the provider
        
    Raises:
        FileNotFoundError: If reference file doesn't exist
        ValueError: If reference data is invalid
    """
    
    # Use cache if available and requested
    if use_cache and provider in _reference_cache:
        logger.info(f"Using cached reference data for {provider}")
        return _reference_cache[provider]
    
    # Use default path if not provided
    if reference_path is None:
        reference_path = DEFAULT_REFERENCE_PATH
    
    # Validate reference data exists and is valid
    if not validate_reference_data(reference_path):
        raise FileNotFoundError(f"Invalid or missing reference data: {reference_path}")
    
    try:
        logger.info(f"Loading reference data for provider: {provider}")
        
        # Use lazy loading with provider pre-filtering for memory efficiency
        reference_data = (
            pl.scan_parquet(reference_path)
            .filter(pl.col('cloud_provider') == provider)
            .filter(pl.col('has_pricing') == True)
            .collect()
        )
        
        if reference_data.is_empty():
            logger.warning(f"No pricing records found for provider: {provider}")
            console.print(f"⚠️  No pricing data available for {provider}")
        else:
            logger.info(f"Loaded {len(reference_data):,} pricing records for {provider}")
            console.print(f"📊 Loaded {len(reference_data):,} pricing records for {provider}")
        
        # Cache the result for reuse across months
        if use_cache:
            _reference_cache[provider] = reference_data
            logger.info(f"Cached reference data for {provider}")
        
        return reference_data
        
    except Exception as e:
        logger.error(f"Failed to load reference data for {provider}: {e}")
        raise ValueError(f"Error loading reference data for {provider}: {e}") from e


def load_all_reference_data(
    reference_path: Optional[Path] = None
) -> pl.DataFrame:
    """Load all reference data (not provider-filtered).
    
    Used for cases where we need to see all available data.
    
    Args:
        reference_path: Optional custom reference data path
        
    Returns:
        DataFrame with all reference data
        
    Raises:
        FileNotFoundError: If reference file doesn't exist
        ValueError: If reference data is invalid
    """
    
    # Use default path if not provided
    if reference_path is None:
        reference_path = DEFAULT_REFERENCE_PATH
    
    # Validate reference data exists and is valid
    if not validate_reference_data(reference_path):
        raise FileNotFoundError(f"Invalid or missing reference data: {reference_path}")
    
    try:
        logger.info("Loading all reference data")
        
        # Load all data with pricing
        reference_data = (
            pl.scan_parquet(reference_path)
            .filter(pl.col('has_pricing') == True)
            .collect()
        )
        
        logger.info(f"Loaded {len(reference_data):,} total pricing records")
        console.print(f"📊 Loaded {len(reference_data):,} total pricing records")
        
        return reference_data
        
    except Exception as e:
        logger.error(f"Failed to load reference data: {e}")
        raise ValueError(f"Error loading reference data: {e}") from e


def clear_reference_cache() -> None:
    """Clear the reference data cache.
    
    Useful for memory management or when reference data is updated.
    """
    global _reference_cache
    cleared_providers = list(_reference_cache.keys())
    _reference_cache.clear()
    
    if cleared_providers:
        logger.info(f"Cleared reference cache for providers: {cleared_providers}")
        console.print(f"🧹 Cleared reference cache for {len(cleared_providers)} providers")


def get_cached_providers() -> list[str]:
    """Get list of providers currently cached.
    
    Returns:
        List of provider names with cached reference data
    """
    return list(_reference_cache.keys())


def get_cache_stats() -> Dict[str, int]:
    """Get cache statistics.
    
    Returns:
        Dictionary with cache statistics per provider
    """
    return {
        provider: len(data) for provider, data in _reference_cache.items()
    }


def preload_reference_data_for_providers(
    providers: list[str],
    reference_path: Optional[Path] = None
) -> Dict[str, bool]:
    """Preload reference data for multiple providers.
    
    Useful for batch processing to avoid repeated loading.
    
    Args:
        providers: List of provider names to preload
        reference_path: Optional custom reference data path
        
    Returns:
        Dictionary mapping provider -> success status
    """
    results = {}
    
    for provider in providers:
        try:
            load_reference_data_for_provider(provider, reference_path, use_cache=True)
            results[provider] = True
        except Exception as e:
            logger.error(f"Failed to preload reference data for {provider}: {e}")
            results[provider] = False
    
    successful = sum(results.values())
    console.print(f"📦 Preloaded reference data for {successful}/{len(providers)} providers")
    
    return results