"""
K8s processing service that orchestrates M:N mapping and cost splitting.

Coordinates the complete K8s allocation process:
1. Filter compute resources
2. Generate K8s topology
3. Perform M:N resource-to-workload mapping
4. Split costs using Dirichlet distribution
5. Transform records with K8s metadata
"""

from typing import Dict, List, Any, Optional
import polars as pl
from collections import defaultdict

from src.log_context import get_logger
from ..config import K8sConfig
from .mapping_service import ResourceToWorkloadMapper, K8sTopology
from .compute_filter_service import ComputeFilterService
from .column_mapper import column_mapper
from ..validation.cluster_isolation_validator import ClusterIsolationValidator


class K8sProcessorService:
    """
    Main service for processing K8s allocation with M:N mapping.
    
    Orchestrates the complete flow from raw billing data to
    K8s-allocated records with workload metadata.
    """
    
    def __init__(self, config: K8sConfig):
        self.config = config
        self.logger = get_logger(__name__)
        
        # Initialize service components
        self.compute_filter = ComputeFilterService(config)
        self.resource_mapper = ResourceToWorkloadMapper(config)
        self.cluster_validator = ClusterIsolationValidator()
        
        # Topology will be initialized during processing
        self.topology: Optional[K8sTopology] = None
        self.resource_mappings: Dict[str, List[str]] = {}
        
        self.logger.info(
            "Initialized K8sProcessorService",
            dataset_name=config.dataset_name
        )
    
    def process_month_data(self, df: pl.DataFrame, provider: str) -> pl.DataFrame:
        """
        Process month data with K8s allocation.

        Args:
            df: DataFrame with billing data for one month

        Returns:
            DataFrame with K8s-allocated records
        """
        if df.is_empty():
            self.logger.warning("Empty DataFrame provided for K8s processing")
            return df

        # Preserve original resource_id for metadata stage fingerprinting
        resource_id_col = column_mapper.get_column_name(df, "ResourceID")
        if resource_id_col:
            df = df.with_columns([
                pl.col(resource_id_col).alias("_original_resource_id")
            ])

        # Find cost column using ColumnMapper
        cost_column = column_mapper.get_column_name(df, "Cost")
        if not cost_column:
            raise ValueError("Cost column not found in DataFrame")
        
        self.logger.info(
            "Starting K8s processing for month data",
            input_rows=len(df),
            input_cost=df.get_column(cost_column).sum()
        )
        
        # Step 1: Filter compute resources for K8s allocation
        self.logger.debug("Starting compute resource filtering")
        k8s_compute_df = self.compute_filter.get_k8s_compute_data(df)
        self.logger.debug(f"Filtered to {len(k8s_compute_df)} compute resources")
        
        if k8s_compute_df.is_empty():
            self.logger.warning("No compute resources available for K8s allocation")
            return df
        
        # Step 2: Get resource ID column for sampling
        resource_id_col = column_mapper.get_column_name(k8s_compute_df, "ResourceID")
        
        # Step 3: Calculate resource allocation with capacity limits
        # CRITICAL: Sort resource IDs for deterministic cross-month consistency
        # Without sorting, unique().to_list() can return different orders each month
        # causing different resource sampling despite identical resource sets
        eligible_resource_ids = sorted(k8s_compute_df[resource_id_col].unique().to_list())
        self.logger.debug(f"Calculating resource allocation for {len(eligible_resource_ids)} eligible resources (deterministically sorted)")
        
        sampled_resource_ids, workloads_per_namespace = self.resource_mapper.calculate_resource_allocation(eligible_resource_ids)
        self.logger.info(
            "Resource allocation calculated",
            eligible_resources=len(eligible_resource_ids),
            sampled_resources=len(sampled_resource_ids),
            workloads_per_namespace=workloads_per_namespace
        )
        
        # Step 4: Generate K8s topology with calculated workload distribution
        self.logger.debug(f"Generating topology for {len(sampled_resource_ids)} sampled resources")
        self.topology = self.resource_mapper.generate_topology(sampled_resource_ids, provider, workloads_per_namespace)
        self.logger.debug("K8s topology generated successfully")
        
        # Step 5: Distribute resources to workloads (dual-constraint allocation)
        self.logger.debug("Starting dual-constraint resource allocation")
        self.resource_mappings = self.resource_mapper.distribute_resources_dual_constraint(
            self.topology, 
            sampled_resource_ids
        )
        self.logger.debug(f"Dual-constraint allocation completed for {len(self.resource_mappings)} resources")
        
        # Step 6: Filter compute records to only sampled resources and transform with K8s metadata
        sampled_compute_df = k8s_compute_df.filter(pl.col(resource_id_col).is_in(sampled_resource_ids))
        self.logger.debug(f"Starting record transformation for {len(sampled_compute_df)} sampled compute records")
        k8s_records = self._transform_compute_records(sampled_compute_df)
        self.logger.debug(f"Record transformation completed - {len(k8s_records)} output records")
        
        # Step 7: Combine K8s and non-K8s data
        # Get non-K8s data (everything except sampled resources)
        non_k8s_df = self.compute_filter.get_non_k8s_data(df, sampled_resource_ids)
        
        if k8s_records:
            k8s_df = pl.DataFrame(k8s_records)

            # Add K8s columns to non-K8s data with empty strings
            # Ensure _original_resource_id is present if not already
            if "_original_resource_id" not in non_k8s_df.columns and resource_id_col:
                non_k8s_df = non_k8s_df.with_columns([
                    pl.col(resource_id_col).alias("_original_resource_id")
                ])

            non_k8s_df = non_k8s_df.with_columns([
                pl.lit("").alias("k8s_cluster"),
                pl.lit("").alias("k8s_namespace"),
                pl.lit("").alias("k8s_workload"),
                pl.lit("").alias("k8s_labels")
            ])
            
            # Ensure schema compatibility before concatenation
            self.logger.debug(
                "Schema alignment for DataFrame concatenation",
                non_k8s_columns=non_k8s_df.columns,
                k8s_columns=k8s_df.columns,
                non_k8s_shape=non_k8s_df.shape,
                k8s_shape=k8s_df.shape
            )
            
            # Align schemas: ensure k8s_df has the same column order as non_k8s_df
            k8s_df = k8s_df.select(non_k8s_df.columns)
            
            # Use vertical_relaxed to handle minor schema differences gracefully
            result_df = pl.concat([non_k8s_df, k8s_df], how="vertical_relaxed")
            
            # Validate cluster isolation to ensure cross-cloud clusters don't occur
            if not k8s_df.is_empty():
                self.logger.debug("Validating cluster isolation for generated K8s data")
                k8s_validation_df = k8s_df.with_columns([
                    pl.lit(provider).alias("cloud_provider")
                ])
                validation_results = self.cluster_validator.validate_provider_isolation(k8s_validation_df)
                
                if not validation_results["validation_passed"]:
                    self.logger.error(
                        "Cluster isolation validation failed - cross-cloud clusters detected!",
                        violations=validation_results["isolation_violations"],
                        provider=provider
                    )
                    raise ValueError(f"Cross-cloud cluster isolation failure for provider {provider}. "
                                   f"Found {validation_results['shared_clusters']} shared clusters.")
                else:
                    self.logger.info(
                        "Cluster isolation validation passed",
                        provider=provider,
                        clusters_validated=validation_results["total_clusters"]
                    )
        else:
            # No K8s records, just add empty K8s columns to original data
            # Ensure _original_resource_id is present if not already
            if "_original_resource_id" not in non_k8s_df.columns and resource_id_col:
                non_k8s_df = non_k8s_df.with_columns([
                    pl.col(resource_id_col).alias("_original_resource_id")
                ])

            result_df = non_k8s_df.with_columns([
                pl.lit("").alias("k8s_cluster"),
                pl.lit("").alias("k8s_namespace"),
                pl.lit("").alias("k8s_workload"),
                pl.lit("").alias("k8s_labels")
            ])
        
        # Step 8: Validate result
        validation_report = self._validate_processing_result(df, result_df)
        
        # Find cost column in result using ColumnMapper
        result_cost_column = column_mapper.get_column_name(result_df, "Cost")
        if not result_cost_column:
            raise ValueError("Cost column not found in result DataFrame")
        
        self.logger.info(
            "Completed K8s processing for month data",
            output_rows=len(result_df),
            output_cost=result_df.get_column(result_cost_column).sum(),
            k8s_allocated_rows=len(k8s_df),
            validation_passed=validation_report["cost_conservation"]
        )
        
        return result_df
    
    def _transform_compute_records(self, compute_df: pl.DataFrame) -> List[Dict[str, Any]]:
        """
        Transform compute records with K8s metadata using M:N mapping.
        
        Args:
            compute_df: DataFrame with compute resources
            
        Returns:
            List of records with K8s metadata
        """
        if not self.topology or not self.resource_mappings:
            raise ValueError("Topology and resource mappings must be initialized before transformation")
        
        # Precompute reverse mappings for faster lookup
        workload_to_namespace = {}
        namespace_to_cluster = {}
        
        for namespace, workloads in self.topology.workloads.items():
            for workload in workloads:
                workload_to_namespace[workload] = namespace
        
        for cluster, namespaces in self.topology.namespaces.items():
            for namespace in namespaces:
                namespace_to_cluster[namespace] = cluster
        
        records = []
        total_original_cost = 0
        total_allocated_cost = 0
        
        # Get column names for this DataFrame
        resource_id_col = column_mapper.get_column_name(compute_df, "ResourceID")
        cost_col = column_mapper.get_column_name(compute_df, "Cost")
        
        # Process each row and create multiple records based on M:N mapping
        total_rows = len(compute_df)
        for i, row in enumerate(compute_df.iter_rows(named=True)):
            if i % 10000 == 0:  # Log progress every 10k rows
                self.logger.debug(f"Processing row {i}/{total_rows} ({i/total_rows*100:.1f}%)")
            
            resource_id = row[resource_id_col]
            original_cost = row[cost_col]
            total_original_cost += original_cost
            
            # Get workloads for this resource
            workloads = self.resource_mappings.get(resource_id, [])
            
            if not workloads:
                # Resource has empty workload mapping - "free" it by keeping as non-K8s resource
                self.logger.debug(
                    "Freeing resource with empty workload mapping - keeping as non-K8s resource",
                    resource_id=resource_id
                )
                
                # Keep original resource data without K8s allocation (add empty K8s columns for schema consistency)
                freed_resource = {
                    resource_id_col: resource_id,
                    cost_col: row[cost_col],
                    # Add empty K8s metadata columns to match schema of allocated resources
                    "k8s_cluster": "",
                    "k8s_namespace": "",
                    "k8s_workload": "",
                    "k8s_labels": "",
                    # Copy all other original columns
                    **{col: row[col] for col in row.keys() if col not in [resource_id_col, cost_col]}
                }
                
                records.append(freed_resource)
                total_allocated_cost += row[cost_col]  # Count freed resource cost towards conservation
                continue
            
            # Split cost among workloads using Dirichlet distribution
            cost_allocation = self.resource_mapper.split_costs_deterministically(
                resource_id, 
                original_cost, 
                workloads,
                self.topology  # Pass topology for environment-based cost adjustments
            )
            
            # Create record for each workload
            for workload, allocated_cost in cost_allocation.items():
                namespace = workload_to_namespace.get(workload)
                cluster = namespace_to_cluster.get(namespace)
                
                if not namespace or not cluster:
                    self.logger.warning(
                        "Missing namespace or cluster for workload",
                        workload=workload,
                        namespace=namespace,
                        cluster=cluster
                    )
                    continue
                
                # Create new record with K8s metadata
                new_record = dict(row)
                new_record.update({
                    "k8s_cluster": cluster,
                    "k8s_namespace": namespace,
                    "k8s_workload": workload,
                    "k8s_labels": "",  # Empty string for no K8s labels
                    cost_col: allocated_cost  # Use actual column name
                })
                
                records.append(new_record)
                total_allocated_cost += allocated_cost
        
        # Validate cost conservation
        cost_difference = abs(total_original_cost - total_allocated_cost)
        cost_tolerance = min(total_original_cost * 0.001, 50000.0)  # 0.1% or $50K cap
        if cost_difference > cost_tolerance:
            self.logger.error(
                "Cost conservation violated in record transformation",
                original_cost=total_original_cost,
                allocated_cost=total_allocated_cost,
                difference=cost_difference,
                tolerance=cost_tolerance
            )
        
        self.logger.info(
            "Transformed compute records",
            original_records=len(compute_df),
            k8s_records=len(records),
            original_cost=total_original_cost,
            allocated_cost=total_allocated_cost,
            cost_difference=cost_difference
        )
        
        return records
    
    def _validate_processing_result(self, original_df: pl.DataFrame, result_df: pl.DataFrame) -> Dict[str, Any]:
        """
        Validate that processing maintains cost conservation.
        
        Args:
            original_df: Original DataFrame
            result_df: Processed DataFrame
            
        Returns:
            Validation report
        """
        # Find cost columns using ColumnMapper
        original_cost_column = column_mapper.get_column_name(original_df, "Cost")
        result_cost_column = column_mapper.get_column_name(result_df, "Cost")
        
        if not original_cost_column or not result_cost_column:
            raise ValueError("Cost column not found in original or result DataFrame")
        
        original_cost = original_df.get_column(original_cost_column).sum()
        result_cost = result_df.get_column(result_cost_column).sum()
        cost_difference = abs(original_cost - result_cost)
        cost_tolerance = min(original_cost * 0.001, 50000.0)  # 0.1% or $50K cap
        cost_matches = cost_difference < cost_tolerance
        
        # Count K8s vs non-K8s records
        k8s_records = result_df.filter(pl.col("k8s_cluster") != "")
        non_k8s_records = result_df.filter(pl.col("k8s_cluster") == "")
        
        validation_report = {
            "cost_conservation": cost_matches,
            "original_cost": float(original_cost),
            "result_cost": float(result_cost),
            "cost_difference": float(cost_difference),
            "original_rows": len(original_df),
            "result_rows": len(result_df),
            "k8s_rows": len(k8s_records),
            "non_k8s_rows": len(non_k8s_records),
            "k8s_cost": float(k8s_records.get_column(result_cost_column).sum()) if len(k8s_records) > 0 else 0,
            "non_k8s_cost": float(non_k8s_records.get_column(result_cost_column).sum()) if len(non_k8s_records) > 0 else 0
        }
        
        if cost_matches:
            self.logger.info(
                "K8s processing validation passed",
                **validation_report
            )
        else:
            self.logger.error(
                "K8s processing validation failed - cost conservation violated",
                **validation_report
            )
        
        return validation_report
    
    def get_topology_summary(self) -> Dict[str, Any]:
        """
        Get summary of generated topology.
        
        Returns:
            Dictionary with topology statistics
        """
        if not self.topology:
            return {"error": "Topology not initialized"}
        
        total_namespaces = sum(len(ns_list) for ns_list in self.topology.namespaces.values())
        total_workloads = sum(len(wl_list) for wl_list in self.topology.workloads.values())
        
        # Calculate mapping statistics
        if self.resource_mappings:
            workloads_per_resource = [len(wls) for wls in self.resource_mappings.values()]
            workload_to_resources = defaultdict(int)
            for workloads in self.resource_mappings.values():
                for workload in workloads:
                    workload_to_resources[workload] += 1
            
            avg_workloads_per_resource = sum(workloads_per_resource) / len(workloads_per_resource)
            avg_resources_per_workload = sum(workload_to_resources.values()) / len(workload_to_resources)
        else:
            avg_workloads_per_resource = 0
            avg_resources_per_workload = 0
        
        return {
            "clusters": len(self.topology.clusters),
            "total_namespaces": total_namespaces,
            "total_workloads": total_workloads,
            "mapped_resources": len(self.resource_mappings),
            "avg_workloads_per_resource": round(avg_workloads_per_resource, 2),
            "avg_resources_per_workload": round(avg_resources_per_workload, 2)
        }