"""
Compute resource filtering service for K8s allocation.

Identifies which resources should be allocated to K8s based on:
- Service types (EC2, Compute Engine, etc.)
- Resource types (instances, virtual machines, etc.)
- Sampling percentage for allocation
"""

from typing import Set, List, Dict, Any
import polars as pl
import numpy as np

from src.log_context import get_logger
from ..config import K8sConfig
from .column_mapper import column_mapper


class ComputeFilterService:
    """
    Service for filtering and sampling compute resources for K8s allocation.
    
    Identifies compute resources that should be allocated to K8s workloads
    based on service types, resource types, and configuration settings.
    """
    
    # Specific service-to-resource-type mappings for K8s allocation
    COMPUTE_ELIGIBILITY_CRITERIA = {
        "AmazonEC2": "instance",
        "microsoft.compute": "virtualmachines", 
        "Compute Engine": "instances"
    }
    
    def __init__(self, config: K8sConfig):
        self.config = config
        self.logger = get_logger(__name__)
        
        # Initialize deterministic sampling RNG
        self.sampling_rng = np.random.default_rng(42)
        
        self.logger.info(
            "Initialized ComputeFilterService",
            dataset_name=config.dataset_name,
            compute_eligibility_criteria=self.COMPUTE_ELIGIBILITY_CRITERIA
        )
    
    def create_compute_mask(self, df: pl.DataFrame) -> pl.Series:
        """
        Create a mask for filtering compute resources.
        
        Args:
            df: DataFrame with billing data
            
        Returns:
            Boolean mask for compute resources
        """
        # Get actual column names
        service_col = column_mapper.get_column_name(df, "Service")
        resource_type_col = column_mapper.get_column_name(df, "ResourceType")
        
        if not service_col or not resource_type_col:
            self.logger.error(
                "Required columns missing for compute filtering",
                available_columns=df.columns,
                service_col=service_col,
                resource_type_col=resource_type_col
            )
            return pl.Series([False] * len(df))
        
        # Create mask using specific service-to-resource-type mapping
        compute_mask = pl.Series([False] * len(df))  # Start with all False
        
        for service, expected_resource_type in self.COMPUTE_ELIGIBILITY_CRITERIA.items():
            service_resource_mask = (
                (df[service_col] == service) & 
                (df[resource_type_col] == expected_resource_type)
            )
            compute_mask = compute_mask | service_resource_mask
        
        compute_count = compute_mask.sum()
        total_count = len(df)
        
        self.logger.debug(
            "Created compute mask",
            compute_resources=compute_count,
            total_resources=total_count,
            compute_percentage=f"{compute_count/total_count:.1%}" if total_count > 0 else "0.0%"
        )
        
        return compute_mask
    
    def filter_compute_resources(self, df: pl.DataFrame) -> pl.DataFrame:
        """
        Filter DataFrame to only compute resources.
        
        Args:
            df: DataFrame with billing data
            
        Returns:
            DataFrame containing only compute resources
        """
        compute_mask = self.create_compute_mask(df)
        compute_df = df.filter(compute_mask)
        
        self.logger.info(
            "Filtered compute resources",
            original_rows=len(df),
            compute_rows=len(compute_df),
            filtered_percentage=f"{len(compute_df)/len(df):.1%}" if len(df) > 0 else "0.0%"
        )
        
        return compute_df
    
    def sample_resources_for_k8s(self, compute_df: pl.DataFrame) -> List[str]:
        """
        Sample compute resources for K8s allocation.
        
        Args:
            compute_df: DataFrame with compute resources
            
        Returns:
            List of resource IDs selected for K8s allocation
        """
        if compute_df.is_empty():
            self.logger.warning("No compute resources available for K8s sampling")
            return []
        
        # Get unique resource IDs per service
        sampled_resource_ids = []
        
        # Get actual column names
        service_col = column_mapper.get_column_name(compute_df, "Service")
        resource_id_col = column_mapper.get_column_name(compute_df, "ResourceID")
        
        for service in compute_df[service_col].unique():
            service_df = compute_df.filter(pl.col(service_col) == service)
            unique_resource_ids = service_df[resource_id_col].unique().to_list()
            
            # Calculate sample size
            # Dual-constraint algorithm uses deterministic resource allocation
            # Sample size determined by optimal topology requirements, not percentage
            sample_size = len(unique_resource_ids)  # Use all eligible resources for dual-constraint
            
            if sample_size > 0:
                # Deterministic sampling using numpy
                self.sampling_rng.shuffle(unique_resource_ids)
                sampled_ids = unique_resource_ids[:sample_size]
                sampled_resource_ids.extend(sampled_ids)
                
                self.logger.debug(
                    "Sampled resources for service",
                    service=service,
                    total_resources=len(unique_resource_ids),
                    sampled_resources=len(sampled_ids),
                    sample_percentage=f"{len(sampled_ids)/len(unique_resource_ids):.1%}"
                )
        
        self.logger.info(
            "Completed resource sampling for K8s",
            total_sampled_resources=len(sampled_resource_ids)
        )
        
        return sampled_resource_ids
    
    def get_k8s_compute_data(self, df: pl.DataFrame) -> pl.DataFrame:
        """
        Get compute data that should be allocated to K8s.
        
        Args:
            df: DataFrame with billing data
            
        Returns:
            DataFrame with compute data for K8s allocation
        """
        # Step 1: Filter to compute resources
        compute_df = self.filter_compute_resources(df)
        
        if compute_df.is_empty():
            self.logger.warning("No compute resources found for K8s allocation")
            return pl.DataFrame()
        
        # Step 2: Sample resources for K8s allocation
        sampled_resource_ids = self.sample_resources_for_k8s(compute_df)
        
        if not sampled_resource_ids:
            self.logger.warning("No resources sampled for K8s allocation")
            return pl.DataFrame()
        
        # Step 3: Filter to sampled resources
        resource_id_col = column_mapper.get_column_name(compute_df, "ResourceID")
        cost_col = column_mapper.get_column_name(compute_df, "Cost")
        
        k8s_compute_df = compute_df.filter(
            pl.col(resource_id_col).is_in(sampled_resource_ids)
        )
        
        self.logger.info(
            "Selected compute data for K8s allocation",
            k8s_compute_rows=len(k8s_compute_df),
            unique_resources=k8s_compute_df.get_column(resource_id_col).n_unique(),
            total_cost=k8s_compute_df.get_column(cost_col).sum()
        )
        
        self.logger.debug(
            "K8s compute sampling details",
            original_compute_rows=len(compute_df),
            k8s_compute_rows=len(k8s_compute_df),
            sampling_percentage=f"{len(k8s_compute_df)/len(compute_df):.1%}" if len(compute_df) > 0 else "0%",
            unique_resources_selected=k8s_compute_df.get_column(resource_id_col).n_unique()
        )
        
        return k8s_compute_df
    
    def get_non_k8s_data(self, df: pl.DataFrame, k8s_resource_ids: List[str]) -> pl.DataFrame:
        """
        Get data that should NOT be allocated to K8s.
        
        Args:
            df: Original DataFrame with billing data
            k8s_resource_ids: List of resource IDs allocated to K8s
            
        Returns:
            DataFrame with non-K8s data
        """
        if not k8s_resource_ids:
            return df
        
        # Filter out K8s-allocated resources
        resource_id_col = column_mapper.get_column_name(df, "ResourceID")
        non_k8s_df = df.filter(
            ~pl.col(resource_id_col).is_in(k8s_resource_ids)
        )
        
        self.logger.info(
            "Selected non-K8s data",
            non_k8s_rows=len(non_k8s_df),
            original_rows=len(df),
            k8s_resource_count=len(k8s_resource_ids),
            non_k8s_percentage=f"{len(non_k8s_df)/len(df):.1%}" if len(df) > 0 else "0.0%"
        )
        
        return non_k8s_df
    
    def validate_compute_allocation(self, original_df: pl.DataFrame, k8s_df: pl.DataFrame, non_k8s_df: pl.DataFrame) -> Dict[str, Any]:
        """
        Validate that compute allocation preserves cost totals.
        
        Args:
            original_df: Original DataFrame
            k8s_df: K8s-allocated DataFrame
            non_k8s_df: Non-K8s DataFrame
            
        Returns:
            Validation report
        """
        # Get cost column name using ColumnMapper
        cost_col = column_mapper.get_column_name(original_df, "Cost")
        if not cost_col:
            raise ValueError("Cost column not found in DataFrame")
        
        original_cost = original_df[cost_col].sum()
        k8s_cost = k8s_df[cost_col].sum() if not k8s_df.is_empty() else 0
        non_k8s_cost = non_k8s_df[cost_col].sum() if not non_k8s_df.is_empty() else 0
        total_allocated_cost = k8s_cost + non_k8s_cost
        
        cost_difference = abs(original_cost - total_allocated_cost)
        cost_tolerance = min(original_cost * 0.001, 50000.0)  # 0.1% or $50K cap
        cost_matches = cost_difference < cost_tolerance
        
        validation_report = {
            "cost_conservation": cost_matches,
            "original_cost": float(original_cost),
            "k8s_cost": float(k8s_cost),
            "non_k8s_cost": float(non_k8s_cost),
            "total_allocated_cost": float(total_allocated_cost),
            "cost_difference": float(cost_difference),
            "original_rows": len(original_df),
            "k8s_rows": len(k8s_df),
            "non_k8s_rows": len(non_k8s_df),
            "allocation_method": "dual_constraint",
            "actual_k8s_cost_percentage": f"{k8s_cost/original_cost:.1%}" if original_cost > 0 else "0.0%"
        }
        
        if cost_matches:
            self.logger.info(
                "Compute allocation validation passed",
                **validation_report
            )
        else:
            self.logger.error(
                "Compute allocation validation failed - cost conservation violated",
                **validation_report
            )
        
        return validation_report