"""Correlated allocation service for metadata stage."""

import polars as pl
from typing import Dict, List, Optional
from ..config import CorrelatedKeyGroup
from .random_service import DeterministicRandomService
from .tags_service import TagsAllocationService
from .four_group_service import FourGroupAllocationService
from .azure_resource_id_service import AzureResourceIdService
from .key_variant_resolver import KeyVariantResolver


class CorrelatedAllocationService:
    """Service for correlated key group allocation - multiple keys allocated together."""

    def __init__(
        self,
        random_service: DeterministicRandomService,
        azure_resource_id_service: Optional[AzureResourceIdService] = None,
        key_variant_resolver: Optional[KeyVariantResolver] = None
    ):
        """Initialize with deterministic random service and optional Azure resource ID service."""
        self.random_service = random_service
        self.azure_resource_id_service = azure_resource_id_service
        self.key_variant_resolver = key_variant_resolver
        self.tags_service = TagsAllocationService(random_service)
        self.four_group_service = FourGroupAllocationService(random_service)

    def allocate_correlated_group(
        self,
        df: pl.DataFrame,
        group_config: CorrelatedKeyGroup,
        provider: str
    ) -> pl.DataFrame:
        """
        Allocate correlated values to resources using index-based correlation.

        Args:
            df: DataFrame with resources to allocate
            group_config: Correlated key group configuration
            provider: Provider name for column naming

        Returns:
            DataFrame with correlated allocations applied
        """
        if df.height == 0:
            return df

        # Get correlated value sets for each resource
        correlated_assignments = self._get_correlated_assignments(df, group_config)

        # Apply assignments based on tags_only setting
        if group_config.tags_only:
            result_df = self._apply_tags_only_allocations(df, correlated_assignments, provider)
        else:
            result_df = self._apply_four_group_allocations(df, correlated_assignments, provider)

        # Azure special case: Update resource group in resource IDs for resourceGroupName
        if (self.azure_resource_id_service and
            provider.lower() == "azure" and
            "resource_id" in result_df.columns):

            # Check if any key in the correlated group is resourceGroupName
            for key_name in group_config.keys:
                if key_name.lower() == "resourcegroupname":
                    # Use the metadata tag column as the source for resource group names
                    tag_column = f"resource/tag:{provider}:{key_name}"
                    if tag_column in result_df.columns:
                        result_df = self.azure_resource_id_service.update_resource_group_in_resource_id(
                            result_df, tag_column
                        )
                    break

        return result_df

    def _get_correlated_assignments(
        self,
        df: pl.DataFrame,
        group_config: CorrelatedKeyGroup
    ) -> pl.DataFrame:
        """
        Get correlated value assignments for each resource using vectorized operations.

        Returns:
            DataFrame with resource identifiers and their correlated values
        """
        # Create fingerprints vectorized using string concatenation
        fingerprints_df = df.select([
            "resource_id",
            # Create deterministic fingerprint for each resource
            pl.concat_str([
                pl.lit(self.random_service.dataset_name),
                pl.lit(":"),
                pl.col("cloud_provider"),
                pl.lit(":"),
                pl.col("service"),
                pl.lit(":"),
                pl.col("resource_id"),
                pl.lit(":correlated")
            ]).alias("fingerprint")
        ])

        # Calculate correlation indices vectorized
        first_key = next(iter(group_config.keys))
        num_sets = len(group_config.keys[first_key])

        indices_df = fingerprints_df.with_columns([
            # Hash fingerprint and calculate modulo for index selection
            (pl.col("fingerprint").hash() % num_sets).cast(pl.UInt32).alias("correlation_index")
        ])

        # For each key in the correlated group, select values using vectorized operations
        result_exprs = [pl.col("resource_id")]

        for key_name, value_list in group_config.keys.items():
            # Create a mapping DataFrame for this key's values with consistent data types
            value_mapping = pl.DataFrame({
                "correlation_index": list(range(len(value_list))),
                key_name: value_list
            }).cast({"correlation_index": pl.UInt32})

            # Join to get the value for each resource's correlation index
            indices_df = indices_df.join(
                value_mapping,
                on="correlation_index",
                how="left"
            )

            # Only include non-null, non-empty values
            result_exprs.append(
                pl.when(
                    (pl.col(key_name).is_not_null()) &
                    (pl.col(key_name) != "")
                ).then(pl.col(key_name)).otherwise(None).alias(key_name)
            )

        # Select final columns, filtering out null-only columns
        result = indices_df.select(result_exprs)

        return result

    def _select_correlated_values(
        self,
        resource_fingerprint: str,
        group_config: CorrelatedKeyGroup
    ) -> Dict[str, str]:
        """
        Select correlated values using index-based correlation.

        Args:
            resource_fingerprint: Unique fingerprint for resource
            group_config: Correlated key group configuration

        Returns:
            Dictionary of key_name -> selected_value (excluding nulls)
        """
        # Get number of correlation sets
        first_key = next(iter(group_config.keys))
        num_sets = len(group_config.keys[first_key])

        # Deterministic index selection
        selected_index = hash(resource_fingerprint) % num_sets

        # Build correlated values, handling nulls
        correlated_values = {}
        for key_name, value_list in group_config.keys.items():
            value = value_list[selected_index]
            if value is not None and value != "":
                correlated_values[key_name] = value

        return correlated_values

    def _apply_tags_only_allocations(
        self,
        df: pl.DataFrame,
        assignments: pl.DataFrame,
        provider: str
    ) -> pl.DataFrame:
        """Apply correlated allocations as resource/tag columns (tags_only=true)."""
        result = df.clone()

        # Get all keys that have assignments
        tag_columns = [col for col in assignments.columns if col != "resource_id"]

        for key_name in tag_columns:
            # Create column name
            normalized_provider = self._normalize_provider_name(provider, key_name)
            column_name = f"resource/tag:{normalized_provider}:{key_name}"

            # Join assignments to get the correlated value for each resource
            result = result.join(
                assignments.select(["resource_id", key_name]),
                on="resource_id",
                how="left"
            ).rename({key_name: column_name})

        return result

    def _apply_four_group_allocations(
        self,
        df: pl.DataFrame,
        assignments: pl.DataFrame,
        provider: str
    ) -> pl.DataFrame:
        """Apply correlated allocations using four-group distribution (tags_only=false)."""
        # For four-group allocation, we need to simulate individual key configs
        # and apply them using the existing four-group service
        result = df.clone()

        # Get all keys that have assignments
        tag_columns = [col for col in assignments.columns if col != "resource_id"]

        for key_name in tag_columns:
            # Extract unique values for this key from assignments
            unique_values = assignments[key_name].drop_nulls().unique().to_list()
            if not unique_values:
                continue

            # Create a temporary KeyConfig for this correlated key
            from ..config import KeyConfig
            temp_key_config = KeyConfig(
                key_name=key_name,
                values=unique_values,
                tags_only=False  # This will be processed by four-group service
            )

            # Apply four-group allocation for this key
            # The four-group service will handle the complexity of K8s vs non-K8s distribution
            result = self.four_group_service.allocate_four_groups(
                result, temp_key_config, provider
            )

        return result

    def _normalize_provider_name(self, provider: str, key_name: str = None) -> str:
        """
        Normalize provider name for column naming.

        Special case: Azure subscriptionName and resourceGroupName always use 'azure:' prefix.
        """
        provider_lower = provider.lower()

        # Special handling for Azure subscription and resource group tags
        if provider_lower == "azure" and key_name:
            key_lower = key_name.lower()
            if key_lower in ["subscriptionname", "resourcegroupname"]:
                return "azure"

        provider_mapping = {
            "aws": "aws",
            "azure": "az",
            "gcp": "gcp",
            "anthropic": "anthropic_cz"
        }
        return provider_mapping.get(provider_lower, provider_lower)

    def build_correlated_expressions(
        self,
        group_config: CorrelatedKeyGroup,
        provider: str
    ) -> List[pl.Expr]:
        """
        Build Polars expressions for correlated allocations with key variant support.

        Args:
            group_config: Correlated key group configuration
            provider: Provider name for column naming

        Returns:
            List of Polars expressions that can be used in with_columns()
        """
        expressions = []

        # Create fingerprint base for correlation
        fingerprint_base = pl.concat_str([
            pl.lit(self.random_service.dataset_name),
            pl.lit(":"),
            pl.col("cloud_provider"),
            pl.lit(":"),
            pl.col("service"),
            pl.lit(":"),
            pl.col("resource_id"),
            pl.lit(":correlated")
        ])

        # Calculate correlation index
        first_key = next(iter(group_config.keys))
        num_sets = len(group_config.keys[first_key])
        correlation_index = fingerprint_base.hash() % num_sets

        # Build expression for each key in the correlated group
        for key_name, value_list in group_config.keys.items():
            if self.key_variant_resolver and self.key_variant_resolver.has_variants(key_name):
                # Build variant-aware expressions for this key
                variant_exprs = self._build_variant_correlated_expressions(
                    key_name, value_list, correlation_index, provider,
                    group_config.tags_only, group_config.k8s_labels_only
                )
                expressions.extend(variant_exprs)
            else:
                # Standard single expression (backward compatibility)
                single_expr = self._build_single_correlated_expression(
                    key_name, value_list, correlation_index, provider,
                    group_config.tags_only, group_config.k8s_labels_only
                )
                expressions.append(single_expr)

        return expressions

    def _build_variant_correlated_expressions(
        self,
        canonical_key: str,
        value_list: List[str],
        correlation_index: pl.Expr,
        provider: str,
        tags_only: bool,
        k8s_labels_only: bool = False
    ) -> List[pl.Expr]:
        """
        Build variant-aware expressions for a correlated key.

        Creates multiple columns (one per variant) where each resource gets assigned
        to exactly one variant column based on deterministic selection, but all
        resources in the same correlated group get the same VALUES.
        """
        expressions = []
        variants = self.key_variant_resolver.variants_config[canonical_key]
        normalized_provider = self._normalize_provider_name(provider, canonical_key)

        # Create resource fingerprint for variant selection (different from correlation fingerprint)
        resource_fingerprint = pl.concat_str([
            pl.col("cloud_provider"),
            pl.lit(":"),
            pl.col("service"),
            pl.lit(":"),
            pl.col("resource_id")
        ])

        # Calculate which variant this resource should use
        variant_count = len(variants)
        variant_index = pl.concat_str([
            pl.lit(canonical_key),
            pl.lit(":"),
            resource_fingerprint
        ]).hash() % variant_count

        # Build value selection expression (same for all variants)
        value_expr = pl.when(correlation_index == 0).then(
            pl.lit(value_list[0]) if value_list[0] is not None and value_list[0] != "" else None
        )

        for i, value in enumerate(value_list[1:], 1):
            condition_value = pl.lit(value) if value is not None and value != "" else None
            value_expr = value_expr.when(correlation_index == i).then(condition_value)

        # Use first non-null value as default, or None if all are null/empty
        default_value = None
        for value in value_list:
            if value is not None and value != "":
                default_value = pl.lit(value)
                break

        final_value_expr = value_expr.otherwise(default_value)

        # Create expressions for each variant column
        if k8s_labels_only:
            # K8s-only mode: Create single k8s_labels JSON with variant key
            k8s_eligible_expr = self._build_k8s_eligibility_expression()

            # Build single k8s_labels expression that uses the selected variant
            k8s_labels_expr = pl.when(k8s_eligible_expr).then(
                pl.concat_str([
                    pl.lit('{"'),
                    # Select the variant key based on variant_index
                    pl.when(variant_index == 0).then(pl.lit(variants[0]))
                    .when(variant_index == 1).then(pl.lit(variants[1]) if len(variants) > 1 else pl.lit(variants[0]))
                    .when(variant_index == 2).then(pl.lit(variants[2]) if len(variants) > 2 else pl.lit(variants[0]))
                    .otherwise(pl.lit(variants[0])),  # Fallback to first variant
                    pl.lit('":"'),
                    final_value_expr,
                    pl.lit('"}')
                ])
            ).otherwise(pl.lit("{}"))

            expressions.append(k8s_labels_expr.alias("k8s_labels"))
        else:
            # Tags-only or four-group mode
            for i, variant in enumerate(variants):
                if tags_only:
                    column_name = f"resource/tag:{normalized_provider}:{variant}"
                else:
                    # For four-group allocation, use tags approach as fallback
                    column_name = f"resource/tag:{normalized_provider}:{variant}"

                # This resource gets the value only if it's assigned to this variant
                variant_expr = pl.when(variant_index == i).then(final_value_expr).otherwise(None)
                expressions.append(variant_expr.alias(column_name))

        return expressions

    def _build_single_correlated_expression(
        self,
        key_name: str,
        value_list: List[str],
        correlation_index: pl.Expr,
        provider: str,
        tags_only: bool,
        k8s_labels_only: bool = False
    ) -> pl.Expr:
        """
        Build a single expression for a correlated key without variants.
        """
        # Build value selection expression
        value_expr = pl.when(correlation_index == 0).then(
            pl.lit(value_list[0]) if value_list[0] is not None and value_list[0] != "" else None
        )

        for i, value in enumerate(value_list[1:], 1):
            condition_value = pl.lit(value) if value is not None and value != "" else None
            value_expr = value_expr.when(correlation_index == i).then(condition_value)

        # Use first non-null value as default, or None if all are null/empty
        default_value = None
        for value in value_list:
            if value is not None and value != "":
                default_value = pl.lit(value)
                break

        final_expr = value_expr.otherwise(default_value)

        # Apply based on allocation mode
        if k8s_labels_only:
            # K8s-only mode: Create k8s_labels JSON
            k8s_eligible_expr = self._build_k8s_eligibility_expression()

            k8s_labels_expr = pl.when(k8s_eligible_expr).then(
                pl.concat_str([
                    pl.lit('{"'),
                    pl.lit(key_name),
                    pl.lit('":"'),
                    final_expr,
                    pl.lit('"}')
                ])
            ).otherwise(pl.lit("{}"))

            return k8s_labels_expr.alias("k8s_labels")
        elif tags_only:
            normalized_provider = self._normalize_provider_name(provider, key_name)
            column_name = f"resource/tag:{normalized_provider}:{key_name}"
            return final_expr.alias(column_name)
        else:
            # For four-group allocation, we'd need more complex expressions
            # For now, use tags approach as fallback
            normalized_provider = self._normalize_provider_name(provider, key_name)
            column_name = f"resource/tag:{normalized_provider}:{key_name}"
            return final_expr.alias(column_name)

    def _build_k8s_eligibility_expression(self) -> pl.Expr:
        """
        Build expression to determine if a resource is eligible for k8s allocation.

        Based on CLAUDE.md requirements:
        - AWS: Service = "AmazonEC2" AND ResourceType = "instance"
        - Azure: Service = "microsoft.compute" AND ResourceType = "virtualmachines"
        - GCP: Service = "Compute Engine" AND ResourceType = "instances"
        """
        return (
            # AWS eligibility
            ((pl.col("cloud_provider") == "AWS") &
             (pl.col("service") == "AmazonEC2") &
             (pl.col("resource_type") == "instance")) |
            # Azure eligibility
            ((pl.col("cloud_provider") == "Azure") &
             (pl.col("service") == "microsoft.compute") &
             (pl.col("resource_type") == "virtualmachines")) |
            # GCP eligibility
            ((pl.col("cloud_provider") == "GCP") &
             (pl.col("service") == "Compute Engine") &
             (pl.col("resource_type") == "instances"))
        )