"""
Metadata Processor - New Architecture.

Main entry point for metadata processing using the clean, straightforward new design.
Implements the complete metadata enrichment pipeline with shared key variants approach.
"""

import logging
import time
from typing import Dict, Optional

import polars as pl

from .config.models import MetadataConfig
from .config.profiles import ProfileConfig
from .services.random_service import DeterministicRandomService
from .services.universal_allocation_service import UniversalAllocationService
from .services.resource_classification_service import ResourceClassificationService
from .services.metadata_enrichment_service import MetadataEnrichmentService
from .services.provider_synchronization_service import ProviderSynchronizationService


logger = logging.getLogger(__name__)


class MetadataProcessor:
    """
    Main processor for metadata enrichment using new architecture.
    
    Processing Flow:
    1. Universal Allocation: cloud_account, region, environment (ALL resources)
    2. Resource Classification: Split based on k8s_cluster column presence  
    3. Type-Specific Processing: Exclusive 50/50 splits with enrichment
    4. Provider Synchronization: GCP and Azure special tag requirements
    
    Key Design Principles:
    - Cross-month consistency using resource fingerprint seeding
    - Shared key variants across all business themes
    - Clean separation of concerns between services
    - Memory-efficient processing with Polars
    """
    
    def __init__(self, profile_config: ProfileConfig):
        """
        Initialize metadata processor.
        
        Args:
            profile_config: Profile configuration with shared key variants
        """
        self.profile_config = profile_config
        
        # Initialize core services with dependency injection
        self.random_service = DeterministicRandomService("metadata-enrichment")
        self.universal_service = UniversalAllocationService(profile_config, self.random_service)
        self.classification_service = ResourceClassificationService(self.random_service)
        self.enrichment_service = MetadataEnrichmentService(profile_config, self.random_service)
        self.sync_service = ProviderSynchronizationService()
        
        # Processing statistics
        self.stats = {
            'total_resources_processed': 0,
            'processing_time_seconds': 0.0,
            'k8s_resources': 0,
            'non_k8s_resources': 0,
            'enrichment_types': {
                'k8s_labels': 0,
                'k8s_namespaces': 0,
                'tags': 0,
                'resource_ids': 0
            }
        }
        
        logger.info("Initialized MetadataProcessor with new architecture")
    
    def process_metadata(self, df: pl.DataFrame, config: MetadataConfig) -> pl.DataFrame:
        """
        Process complete metadata enrichment pipeline.
        
        Args:
            df: Input DataFrame with resources
            config: Metadata configuration
            
        Returns:
            DataFrame with complete metadata enrichment applied
        """
        if df.height == 0:
            logger.warning("Empty DataFrame provided for metadata processing")
            return df
        
        start_time = time.time()
        logger.info(f"Starting metadata processing for {df.height:,} resources")
        
        # DEBUG: Check input data schema
        
        # Step 1: Universal Allocation (ALL resources)
        logger.info("Step 1: Applying universal allocation (cloud_account, region, environment_variant)...")
        result = self.universal_service.apply_universal_allocation(df, config)
        
        # Step 2: Resource Classification
        logger.info("Step 2: Classifying resources based on k8s_cluster column...")
        k8s_resources, non_k8s_resources = self.classification_service.classify_resources(result)
        
        # Update statistics
        self.stats['k8s_resources'] = k8s_resources.height
        self.stats['non_k8s_resources'] = non_k8s_resources.height
        
        # Step 3: Type-Specific Processing with Exclusive 50/50 Splits
        logger.info("Step 3: Applying type-specific enrichment with exclusive splits...")
        
        processed_resources = []
        
        # Process K8s resources (if any)
        if k8s_resources.height > 0:
            enriched_k8s = self._process_k8s_resources(k8s_resources, config)
            processed_resources.append(enriched_k8s)
        else:
            logger.info("No K8s resources found - skipping K8s processing")
        
        # Process non-K8s resources (if any)
        if non_k8s_resources.height > 0:
            enriched_non_k8s = self._process_non_k8s_resources(non_k8s_resources, config)
            processed_resources.append(enriched_non_k8s)
        
        # Step 4: Combine Results
        if processed_resources:
            # Ensure all DataFrames have compatible schemas before concatenating
            if len(processed_resources) > 1:
                # Check column counts for debugging
                column_counts = [df.width for df in processed_resources]
                logger.info(f"Combining {len(processed_resources)} DataFrames with column counts: {column_counts}")
            
            final_result = pl.concat(processed_resources)
        else:
            # Fallback to universal allocation only - add null K8s columns for consistency  
            k8s_columns_to_add = []
            for col_name in ['k8s_cluster', 'k8s_namespace', 'k8s_workload', 'k8s_labels']:
                if col_name not in result.columns:
                    k8s_columns_to_add.append(pl.lit(None).alias(col_name))
            
            if k8s_columns_to_add:
                result = result.with_columns(k8s_columns_to_add)
                logger.info(f"Added {len(k8s_columns_to_add)} null K8s columns to fallback result")
            
            final_result = result
        
        # Step 5: Provider Synchronization (GCP and Azure special requirements)
        logger.info("Step 5: Applying provider-specific synchronization...")
        final_result = self.sync_service.synchronize_all_providers(final_result)
        
        # Update final statistics
        self.stats['total_resources_processed'] = final_result.height
        self.stats['processing_time_seconds'] = time.time() - start_time
        
        logger.info(f"Completed metadata processing: {final_result.height:,} resources enriched in {self.stats['processing_time_seconds']:.2f}s")
        
        return final_result
    
    def _process_k8s_resources(
        self, 
        k8s_resources: pl.DataFrame, 
        config: MetadataConfig
    ) -> pl.DataFrame:
        """
        Process K8s resources with exclusive 50/50 split (labels vs namespaces).
        
        Args:
            k8s_resources: DataFrame with K8s resources (includes environment_variant column)
            config: Metadata configuration
            
        Returns:
            DataFrame with K8s resources enriched
        """
        logger.info(f"Processing {k8s_resources.height:,} K8s resources with exclusive splits...")
        
        # Split into labels vs namespace enrichment (exclusive 50/50)
        label_resources, namespace_resources = self.classification_service.determine_k8s_enrichment_splits(
            k8s_resources
        )
        
        processed_resources = []
        
        # Process label enrichment resources (EXCLUSIVE)
        if label_resources.height > 0:
            logger.info(f"Enriching {label_resources.height:,} K8s resources with labels...")
            enriched_labels = self.enrichment_service.enrich_k8s_labels(
                label_resources, 
                config.base_allocation
            )
            processed_resources.append(enriched_labels)
            self.stats['enrichment_types']['k8s_labels'] = label_resources.height
        
        # Process namespace enrichment resources (EXCLUSIVE)
        if namespace_resources.height > 0:
            logger.info(f"Enriching {namespace_resources.height:,} K8s resources with namespaces...")
            enriched_namespaces = self.enrichment_service.enrich_k8s_namespaces(
                namespace_resources,
                config.base_allocation
            )
            processed_resources.append(enriched_namespaces)
            self.stats['enrichment_types']['k8s_namespaces'] = namespace_resources.height
        
        # Combine processed K8s resources
        if processed_resources:
            return pl.concat(processed_resources)
        else:
            return k8s_resources
    
    def _process_non_k8s_resources(
        self, 
        non_k8s_resources: pl.DataFrame, 
        config: MetadataConfig
    ) -> pl.DataFrame:
        """
        Process non-K8s resources with exclusive 50/50 split (tags vs resource IDs).
        
        Args:
            non_k8s_resources: DataFrame with non-K8s resources (includes environment_variant column)
            config: Metadata configuration
            
        Returns:
            DataFrame with non-K8s resources enriched
        """
        logger.info(f"Processing {non_k8s_resources.height:,} non-K8s resources with exclusive splits...")
        
        # Split into tags vs resource ID enrichment (exclusive 50/50)
        tag_resources, resource_id_resources = self.classification_service.determine_non_k8s_enrichment_splits(
            non_k8s_resources
        )
        
        processed_resources = []
        
        # Process tag enrichment resources (EXCLUSIVE)
        if tag_resources.height > 0:
            logger.info(f"Enriching {tag_resources.height:,} non-K8s resources with tags...")
            enriched_tags = self.enrichment_service.enrich_tags(
                tag_resources,
                config.base_allocation
            )
            processed_resources.append(enriched_tags)
            self.stats['enrichment_types']['tags'] = tag_resources.height
        
        # Process resource ID enrichment resources (EXCLUSIVE)
        if resource_id_resources.height > 0:
            logger.info(f"Enriching {resource_id_resources.height:,} non-K8s resources with resource IDs...")
            enriched_resource_ids = self.enrichment_service.enrich_resource_ids(
                resource_id_resources,
                config.base_allocation
            )
            processed_resources.append(enriched_resource_ids)
            self.stats['enrichment_types']['resource_ids'] = resource_id_resources.height
        
        # Combine processed non-K8s resources
        if processed_resources:
            result = pl.concat(processed_resources)
        else:
            result = non_k8s_resources
        
        # Ensure schema compatibility by adding missing K8s columns with null values
        k8s_columns_to_add = []
        for col_name in ['k8s_cluster', 'k8s_namespace', 'k8s_workload', 'k8s_labels']:
            if col_name not in result.columns:
                k8s_columns_to_add.append(pl.lit(None).alias(col_name))
        
        if k8s_columns_to_add:
            result = result.with_columns(k8s_columns_to_add)
            logger.info(f"Added {len(k8s_columns_to_add)} null K8s columns to non-K8s resources for schema compatibility")
        
        return result
    
    def get_processing_stats(self) -> Dict[str, any]:
        """
        Get comprehensive processing statistics.
        
        Returns:
            Dictionary with processing statistics
        """
        return {
            'total_resources_processed': self.stats['total_resources_processed'],
            'processing_time_seconds': self.stats['processing_time_seconds'],
            'resource_classification': {
                'k8s_resources': self.stats['k8s_resources'],
                'non_k8s_resources': self.stats['non_k8s_resources'],
                'k8s_percentage': (self.stats['k8s_resources'] / max(1, self.stats['total_resources_processed'])) * 100,
                'non_k8s_percentage': (self.stats['non_k8s_resources'] / max(1, self.stats['total_resources_processed'])) * 100
            },
            'enrichment_distribution': self.stats['enrichment_types'],
            'architecture': 'new_shared_key_variants'
        }
    
    def validate_processing(self, original_df: pl.DataFrame, processed_df: pl.DataFrame) -> Dict[str, bool]:
        """
        Validate metadata processing results.
        
        Args:
            original_df: Original input DataFrame
            processed_df: Processed output DataFrame
            
        Returns:
            Dictionary with validation results
        """
        validation_results = {
            'row_count_preserved': original_df.height == processed_df.height,
            'required_columns_present': True,
            'cloud_accounts_allocated': True,
            'regions_allocated': True,
            'provider_sync_valid': True,
            'exclusive_splits_maintained': True
        }
        
        if processed_df.height == 0:
            return validation_results
        
        # Check required columns are present
        required_columns = ['cloud_account', 'region']
        missing_columns = [col for col in required_columns if col not in processed_df.columns]
        validation_results['required_columns_present'] = len(missing_columns) == 0
        
        # Validate cloud account allocation
        if 'cloud_account' in processed_df.columns:
            null_accounts = processed_df.select(pl.col('cloud_account').is_null().sum()).item()
            validation_results['cloud_accounts_allocated'] = (null_accounts == 0)
        
        # Validate region allocation
        if 'region' in processed_df.columns:
            null_regions = processed_df.select(pl.col('region').is_null().sum()).item()
            validation_results['regions_allocated'] = (null_regions == 0)
        
        # Validate provider synchronization
        sync_validation = self.sync_service.validate_synchronization(processed_df)
        validation_results['provider_sync_valid'] = all(sync_validation.values())
        
        # Validate resource classification
        k8s_resources, non_k8s_resources = self.classification_service.classify_resources(processed_df)
        classification_validation = self.classification_service.validate_classification(
            k8s_resources, non_k8s_resources
        )
        validation_results['exclusive_splits_maintained'] = all(classification_validation.values())
        
        return validation_results
    
