"""
Universal Allocation Service.

Handles the three universal allocation components that ALL resources receive:
1. Cloud Account Allocation - Provider-specific account patterns
2. Region Allocation - Geographic market distribution with provider-specific regions
3. Environment Allocation - Environment value variant selection for later enrichment

These are column-only updates (no tags created at this stage).
"""

import logging
from typing import Dict

import polars as pl

from ..config.models import MetadataConfig
from ..config.profiles import ProfileConfig
from .random_service import DeterministicRandomService


logger = logging.getLogger(__name__)


class UniversalAllocationService:
    """Service for applying universal allocation to all resources."""
    
    # Provider-specific account patterns with realistic names and weights
    PROVIDER_ACCOUNT_PATTERNS = {
        'AWS': [
            ('aws-prod-main', 0.20),
            ('aws-prod-analytics', 0.15),
            ('aws-staging-core', 0.15),
            ('aws-dev-primary', 0.10),
            ('aws-sandbox-demo', 0.07),
            ('aws-prod-secondary', 0.12),
            ('aws-staging-test', 0.08),
            ('aws-dev-secondary', 0.08),
            ('aws-sandbox-trial', 0.03),
            ('aws-prod-backup', 0.02)
        ],
        'AWS Marketplace': [  # Same as AWS
            ('aws-prod-main', 0.20),
            ('aws-prod-analytics', 0.15),
            ('aws-staging-core', 0.15),
            ('aws-dev-primary', 0.10),
            ('aws-sandbox-demo', 0.07),
            ('aws-prod-secondary', 0.12),
            ('aws-staging-test', 0.08),
            ('aws-dev-secondary', 0.08),
            ('aws-sandbox-trial', 0.03),
            ('aws-prod-backup', 0.02)
        ],
        'Azure': [
            ('Enterprise-IT-Production', 0.20),
            ('Data-Analytics-Platform', 0.15),
            ('Production-Core-Services', 0.10),
            ('Staging-Development-Env', 0.15),
            ('Dev-Sandbox-Resources', 0.08),
            ('Enterprise-Backup-Systems', 0.10),
            ('Analytics-ML-Platform', 0.12),
            ('Testing-QA-Environment', 0.05),
            ('Trial-Demo-Workspace', 0.03),
            ('DR-Recovery-Services', 0.02)
        ],
        'Azure Marketplace': [  # Same as Azure
            ('Enterprise-IT-Production', 0.20),
            ('Data-Analytics-Platform', 0.15),
            ('Production-Core-Services', 0.10),
            ('Staging-Development-Env', 0.15),
            ('Dev-Sandbox-Resources', 0.08),
            ('Enterprise-Backup-Systems', 0.10),
            ('Analytics-ML-Platform', 0.12),
            ('Testing-QA-Environment', 0.05),
            ('Trial-Demo-Workspace', 0.03),
            ('DR-Recovery-Services', 0.02)
        ],
        'GCP': [
            ('global-operations-prod', 0.18),
            ('customer-data-platform', 0.15),
            ('core-infrastructure-main', 0.12),
            ('analytics-staging-env', 0.12),
            ('ml-development-proj', 0.10),
            ('backup-recovery-systems', 0.08),
            ('testing-sandbox-env', 0.08),
            ('demo-trial-workspace', 0.07),
            ('monitoring-ops-proj', 0.06),
            ('archive-storage-proj', 0.04)
        ],
        'Anthropic': [
            ('anthropic-prod-api', 0.25),
            ('anthropic-staging-models', 0.20),
            ('anthropic-dev-research', 0.15),
            ('anthropic-testing-env', 0.12),
            ('anthropic-demo-platform', 0.10),
            ('anthropic-backup-systems', 0.08),
            ('anthropic-sandbox-trial', 0.05),
            ('anthropic-monitoring-ops', 0.03),
            ('anthropic-archive-data', 0.02)
        ],
        'OpenAI': [
            ('openai-prod-gpt', 0.25),
            ('openai-staging-dalle', 0.20),
            ('openai-dev-models', 0.15),
            ('openai-testing-api', 0.12),
            ('openai-demo-showcase', 0.10),
            ('openai-research-lab', 0.08),
            ('openai-sandbox-trial', 0.05),
            ('openai-monitoring-sys', 0.03),
            ('openai-backup-store', 0.02)
        ],
        'Databricks': [
            ('databricks-analytics-prod', 0.22),
            ('databricks-ml-staging', 0.18),
            ('databricks-data-platform', 0.15),
            ('databricks-dev-workspace', 0.12),
            ('databricks-testing-env', 0.10),
            ('databricks-demo-cluster', 0.08),
            ('databricks-research-lab', 0.07),
            ('databricks-sandbox-trial', 0.05),
            ('databricks-backup-sys', 0.03)
        ],
        'Datadog': [
            ('datadog-monitoring-prod', 0.25),
            ('datadog-logging-dev', 0.15),
            ('datadog-apm-staging', 0.15),
            ('datadog-infra-main', 0.12),
            ('datadog-security-ops', 0.10),
            ('datadog-testing-env', 0.08),
            ('datadog-demo-platform', 0.07),
            ('datadog-trial-workspace', 0.05),
            ('datadog-backup-logs', 0.03)
        ],
        'MongoDB': [
            ('mongodb-atlas-prod', 0.23),
            ('mongodb-staging-cluster', 0.17),
            ('mongodb-dev-database', 0.15),
            ('mongodb-analytics-db', 0.12),
            ('mongodb-testing-env', 0.10),
            ('mongodb-backup-replica', 0.08),
            ('mongodb-demo-cluster', 0.07),
            ('mongodb-trial-instance', 0.05),
            ('mongodb-archive-store', 0.03)
        ],
        'New Relic': [
            ('newrelic-apm-prod', 0.25),
            ('newrelic-infra-staging', 0.20),
            ('newrelic-browser-dev', 0.15),
            ('newrelic-mobile-main', 0.12),
            ('newrelic-synthetic-ops', 0.10),
            ('newrelic-testing-env', 0.08),
            ('newrelic-demo-account', 0.05),
            ('newrelic-trial-space', 0.03),
            ('newrelic-backup-mon', 0.02)
        ]
    }
    
    # Geographic market region mappings with provider-specific regions and defaults
    GEOGRAPHIC_REGIONS = {
        'na': {  # North America - 60% of resources
            'AWS': ['us-east-1', 'us-west-2', 'ca-central-1'],
            'Azure': ['eastus', 'westus2', 'canadacentral'],
            'GCP': ['us-central1', 'us-west1', 'northamerica-northeast1'],
            'default': ['us-east-1']
        },
        'emea': {  # Europe, Middle East, Africa - 20% of resources
            'AWS': ['eu-west-1', 'eu-central-1', 'me-south-1'],
            'Azure': ['westeurope', 'northeurope', 'uaenorth'],
            'GCP': ['europe-west1', 'europe-north1', 'me-west1'],
            'default': ['eu-west-1']
        },
        'latam': {  # Latin America - 15% of resources
            'AWS': ['sa-east-1'],
            'Azure': ['brazilsouth'],
            'GCP': ['southamerica-east1'],
            'default': ['sa-east-1']
        },
        'apac': {  # Asia Pacific - 5% of resources
            'AWS': ['ap-southeast-1', 'ap-northeast-1', 'ap-south-1'],
            'Azure': ['southeastasia', 'japaneast', 'centralindia'],
            'GCP': ['asia-southeast1', 'asia-northeast1', 'asia-south1'],
            'default': ['ap-southeast-1']
        }
    }
    
    # Geographic market distribution
    MARKET_DISTRIBUTION = {
        'na': 0.60,    # 60% North America
        'emea': 0.20,  # 20% Europe/Middle East/Africa
        'latam': 0.15, # 15% Latin America
        'apac': 0.05   # 5% Asia Pacific
    }
    
    def __init__(
        self, 
        profile_config: ProfileConfig, 
        random_service: DeterministicRandomService
    ):
        """
        Initialize universal allocation service.
        
        Args:
            profile_config: Profile configuration for environment variants
            random_service: Deterministic random service for consistent allocation
        """
        self.profile_config = profile_config
        self.random_service = random_service
        
        logger.info("Initialized UniversalAllocationService for cloud_account, region, and environment allocation")
    
    def apply_universal_allocation(
        self, 
        df: pl.DataFrame, 
        config: MetadataConfig
    ) -> pl.DataFrame:
        """
        Apply all universal allocation components to resources.
        
        Args:
            df: Input DataFrame with resources
            config: Metadata configuration
            
        Returns:
            DataFrame with cloud_account, region, and environment_variant columns applied
        """
        if df.height == 0:
            return df
        
        logger.info(f"Applying universal allocation to {df.height:,} resources")
        
        # DEBUG: Check input resource_id type
        
        result = df.clone()
        
        # 1. Cloud Account Allocation (column update only)
        result = self.allocate_cloud_accounts(result)
        
        # 2. Region Allocation (column update only)
        result = self.allocate_regions(result, config)
        
        # 3. Environment Allocation (vectorized column operation)
        result = self.apply_environment_allocation(result, config)
        
        logger.info("Completed universal allocation")
        
        return result
    
    def allocate_cloud_accounts(self, df: pl.DataFrame) -> pl.DataFrame:
        """
        Allocate cloud accounts with provider-specific patterns using vectorized operations.
        
        Args:
            df: DataFrame with resources
            
        Returns:
            DataFrame with cloud_account column updated
        """
        if df.height == 0:
            return df
        
        logger.info(f"Allocating cloud accounts for {df.height:,} resources")
        
        # Use vectorized cloud account allocation
        result = self._allocate_cloud_accounts_vectorized(df)
        
        logger.info("Completed cloud account allocation")
        return result
    
    def allocate_regions(self, df: pl.DataFrame, config: MetadataConfig) -> pl.DataFrame:
        """
        Allocate regions based on geographic market distribution using vectorized operations.
        
        Args:
            df: DataFrame with resources
            config: Metadata configuration with majority regions
            
        Returns:
            DataFrame with region column updated
        """
        if df.height == 0:
            return df
        
        logger.info(f"Allocating regions for {df.height:,} resources")
        
        # Use vectorized region allocation
        result = self._allocate_regions_vectorized(df, config)
        
        logger.info("Completed region allocation")
        return result
    
    def apply_environment_allocation(
        self, 
        df: pl.DataFrame, 
        config: MetadataConfig
    ) -> pl.DataFrame:
        """
        Apply environment allocation using vectorized operations.
        
        Args:
            df: DataFrame with resources
            config: Metadata configuration with majority environment
            
        Returns:
            DataFrame with environment_variant column added
        """
        if df.height == 0:
            return df
        
        logger.info(f"Applying environment allocation to {df.height:,} resources")
        
        # Get environment value variants from profile
        available_variants = self.profile_config.environments.get_variants(config.majority_environment)
        if not available_variants:
            logger.warning(f"No environment variants found for {config.majority_environment}")
            # Fallback to environment name as variant
            return df.with_columns([
                pl.lit(config.majority_environment).alias('environment_variant')
            ])
        
        # Simple case - all resources get the same variant
        if len(available_variants) == 1:
            logger.info(f"Single environment variant: {available_variants[0]}")
            return df.with_columns([
                pl.lit(available_variants[0]).alias('environment_variant')
            ])
        
        # Multiple variants - use deterministic selection
        logger.info(f"Selecting from {len(available_variants)} environment variants")
        
        # Create seed strings for deterministic hashing
        seed_column = pl.concat_str([
            pl.col('cloud_provider'),
            pl.lit(':'),
            pl.col('service'),
            pl.lit(':'),
            pl.col('resource_id'),
            pl.lit(':environment_variant')
        ])
        
        # Use modulo arithmetic with hash to select variants (vectorized)
        result = df.with_columns([
            (seed_column.hash() % len(available_variants)).alias('variant_index')
        ])
        
        # Create when/then chains to map indices to variants (vectorized)
        variant_mapping = pl.col('variant_index')
        for i, variant in enumerate(available_variants):
            if i == 0:
                variant_mapping = pl.when(pl.col('variant_index') == i).then(pl.lit(variant))
            else:
                variant_mapping = variant_mapping.when(pl.col('variant_index') == i).then(pl.lit(variant))
        
        return result.with_columns([
            variant_mapping.otherwise(pl.lit(available_variants[0])).alias('environment_variant')
        ]).drop('variant_index')
    
    def _allocate_cloud_accounts_vectorized(self, df: pl.DataFrame) -> pl.DataFrame:
        """Allocate cloud accounts using vectorized operations."""
        result = df.clone()
        
        # DEBUG: Check before creating seed
        
        # Create resource fingerprint column for seeding
        result = result.with_columns([
            pl.concat_str([
                pl.col('cloud_provider'),
                pl.lit(':'),
                pl.col('service'),
                pl.lit(':'),
                pl.col('resource_id'),
                pl.lit(':cloud_account')
            ]).alias('_seed')
        ])
        
        # DEBUG: Check after creating seed
        
        # Process each provider separately for efficient vectorized selection
        provider_results = []
        
        for provider in result.select('cloud_provider').unique().to_series().to_list():
            provider_df = result.filter(pl.col('cloud_provider') == provider)
            
            if provider_df.height == 0:
                continue
                
            # Get account patterns for this provider
            account_patterns = self.PROVIDER_ACCOUNT_PATTERNS.get(provider)
            
            if not account_patterns:
                # Default pattern for unlisted providers
                provider_df = provider_df.with_columns([
                    pl.lit(f"{provider.lower()}-prod-main").alias('cloud_account')
                ])
            else:
                # Extract accounts and weights
                accounts = [pattern[0] for pattern in account_patterns]
                weights = [pattern[1] for pattern in account_patterns]
                
                # Use hash-based deterministic selection (vectorized)
                seed_hash = provider_df.select(pl.col('_seed').hash()).to_series()
                
                # Convert to cumulative weights for selection
                import numpy as np
                cumulative_weights = np.cumsum(weights)
                total_weight = cumulative_weights[-1]
                
                # Normalize hash values to [0, total_weight) range
                normalized_values = (seed_hash % 1000000) / 1000000.0 * total_weight
                
                # Select accounts based on cumulative weights
                selected_indices = np.searchsorted(cumulative_weights, normalized_values)
                selected_accounts = [accounts[idx] for idx in selected_indices]
                
                # Add selected accounts to the provider DataFrame
                provider_df = provider_df.with_columns([
                    pl.Series('cloud_account', selected_accounts)
                ])
            
            provider_results.append(provider_df)
        
        # Combine results from all providers
        if provider_results:
            result = pl.concat(provider_results)
        
        # Clean up temporary columns and maintain original order
        result = result.drop(['_seed']).sort('cloud_provider', 'service', 'resource_id')
        
        return result
    
    def _allocate_regions_vectorized(self, df: pl.DataFrame, config: MetadataConfig) -> pl.DataFrame:
        """Allocate regions using vectorized operations."""
        result = df.clone()
        
        # Step 1: Select geographic markets (vectorized)
        # Create seed for market selection
        result = result.with_columns([
            pl.concat_str([
                pl.col('cloud_provider'),
                pl.lit(':'),
                pl.col('service'),
                pl.lit(':'),
                pl.col('resource_id'),
                pl.lit(':market')
            ]).alias('_market_seed')
        ])
        
        # Vectorized market selection based on MARKET_DISTRIBUTION
        import numpy as np
        markets = list(self.MARKET_DISTRIBUTION.keys())
        weights = list(self.MARKET_DISTRIBUTION.values())
        cumulative_weights = np.cumsum(weights)
        total_weight = cumulative_weights[-1]
        
        # Hash-based market selection (vectorized)
        market_hash = result.select(pl.col('_market_seed').hash()).to_series()
        normalized_values = (market_hash % 1000000) / 1000000.0 * total_weight
        selected_market_indices = np.searchsorted(cumulative_weights, normalized_values)
        selected_markets = [markets[idx] for idx in selected_market_indices]
        
        result = result.with_columns([
            pl.Series('_selected_market', selected_markets)
        ])
        
        # Step 2: Select regions within markets (grouped by provider and market)
        provider_market_results = []
        
        for provider in result.select('cloud_provider').unique().to_series().to_list():
            provider_df = result.filter(pl.col('cloud_provider') == provider)
            
            for market in provider_df.select('_selected_market').unique().to_series().to_list():
                market_df = provider_df.filter(pl.col('_selected_market') == market)
                
                if market_df.height == 0:
                    continue
                
                # Get regions for this provider/market combination
                market_regions = self.GEOGRAPHIC_REGIONS[market]
                provider_regions = market_regions.get(provider, market_regions['default'])
                
                if len(provider_regions) == 1:
                    # Single region - no selection needed
                    market_df = market_df.with_columns([
                        pl.lit(provider_regions[0]).alias('region')
                    ])
                else:
                    # Multiple regions - apply weighting
                    region_weights = []
                    
                    # Check for majority region configuration
                    if (hasattr(config, 'majority_regions') and 
                        market in config.majority_regions and 
                        hasattr(config.majority_regions[market], 'majority')):
                        
                        majority_region = config.majority_regions[market].majority
                        if majority_region in provider_regions:
                            # 70% majority, 30% others
                            remaining_weight = 0.30 / max(1, len(provider_regions) - 1)
                            for region in provider_regions:
                                if region == majority_region:
                                    region_weights.append((region, 0.70))
                                else:
                                    region_weights.append((region, remaining_weight))
                        else:
                            # Equal distribution if majority not found
                            equal_weight = 1.0 / len(provider_regions)
                            region_weights = [(region, equal_weight) for region in provider_regions]
                    else:
                        # Equal distribution
                        equal_weight = 1.0 / len(provider_regions)
                        region_weights = [(region, equal_weight) for region in provider_regions]
                    
                    # Vectorized region selection within market
                    regions = [rw[0] for rw in region_weights]
                    weights = [rw[1] for rw in region_weights]
                    cumulative_weights = np.cumsum(weights)
                    total_weight = cumulative_weights[-1]
                    
                    # Create region seed and select
                    region_seed_col = pl.concat_str([
                        pl.col('cloud_provider'),
                        pl.lit(':'),
                        pl.col('service'),
                        pl.lit(':'),
                        pl.col('resource_id'),
                        pl.lit(':region')
                    ])
                    region_seeds = market_df.select(region_seed_col.hash()).to_series()
                    
                    normalized_values = (region_seeds % 1000000) / 1000000.0 * total_weight
                    selected_region_indices = np.searchsorted(cumulative_weights, normalized_values)
                    selected_regions = [regions[idx] for idx in selected_region_indices]
                    
                    market_df = market_df.with_columns([
                        pl.Series('region', selected_regions)
                    ])
                
                provider_market_results.append(market_df)
        
        # Combine all results
        if provider_market_results:
            result = pl.concat(provider_market_results)
        
        # Clean up temporary columns and maintain original order
        result = result.drop(['_market_seed', '_selected_market']).sort('cloud_provider', 'service', 'resource_id')
        
        return result
    
    def get_allocation_stats(self, df: pl.DataFrame) -> Dict[str, any]:
        """
        Get universal allocation statistics.
        
        Args:
            df: DataFrame with allocated resources
            
        Returns:
            Dictionary with allocation statistics
        """
        if df.height == 0:
            return {
                'total_resources': 0,
                'unique_cloud_accounts': 0,
                'unique_regions': 0,
                'providers': []
            }
        
        stats = {
            'total_resources': df.height,
            'unique_cloud_accounts': 0,
            'unique_regions': 0,
            'providers': []
        }
        
        if 'cloud_account' in df.columns:
            stats['unique_cloud_accounts'] = df.select('cloud_account').n_unique()
        
        if 'region' in df.columns:
            stats['unique_regions'] = df.select('region').n_unique()
        
        if 'cloud_provider' in df.columns:
            stats['providers'] = df.select('cloud_provider').unique().to_series().to_list()
        
        return stats
    
