"""Service for aggregating multiple k8s-only keys into single JSON expression."""

import polars as pl
from typing import List, Optional
import logging

from ..config import KeyConfig, CorrelatedKeyGroup
from .random_service import DeterministicRandomService
from .key_variant_resolver import KeyVariantResolver

logger = logging.getLogger(__name__)


class K8sLabelsAggregationService:
    """
    Service for merging multiple k8s_labels_only keys into a single k8s_labels JSON expression.

    Handles both individual keys and correlated groups that have k8s_labels_only=True,
    creating realistic messy labels with variant key names for CloudZero to clean up.
    """

    def __init__(
        self,
        random_service: DeterministicRandomService,
        key_variant_resolver: Optional[KeyVariantResolver] = None
    ):
        """Initialize with required services."""
        self.random_service = random_service
        self.key_variant_resolver = key_variant_resolver

    def build_merged_k8s_labels_expression(
        self,
        k8s_keys: List[KeyConfig],
        k8s_correlated_groups: List[CorrelatedKeyGroup]
    ) -> Optional[pl.Expr]:
        """
        Build a single k8s_labels expression that merges all k8s-only keys and groups.

        Args:
            k8s_keys: Individual KeyConfigs with k8s_labels_only=True
            k8s_correlated_groups: CorrelatedKeyGroups with k8s_labels_only=True

        Returns:
            Single Polars expression for k8s_labels column, or None if no k8s-only keys
        """
        if not k8s_keys and not k8s_correlated_groups:
            return None

        logger.debug(f"Building merged k8s_labels expression for {len(k8s_keys)} individual keys and {len(k8s_correlated_groups)} correlated groups")

        # Build JSON components for individual keys
        json_components = []

        # Add individual keys
        for key_config in k8s_keys:
            components = self._build_key_json_components(key_config)
            json_components.extend(components)

        # Add correlated group keys
        for group_config in k8s_correlated_groups:
            components = self._build_correlated_group_json_components(group_config)
            json_components.extend(components)

        if not json_components:
            logger.debug("No JSON components generated, returning empty k8s_labels")
            return self._build_empty_k8s_labels_expression()

        # Build final merged expression
        merged_expr = self._build_merged_json_expression(json_components)

        logger.debug(f"Built merged k8s_labels expression with {len(json_components)} components")
        return merged_expr

    def _build_key_json_components(self, key_config: KeyConfig) -> List[pl.Expr]:
        """Build JSON components for an individual key with variant support."""
        if self.key_variant_resolver and self.key_variant_resolver.has_variants(key_config.key_name):
            return self._build_variant_key_json_components(key_config)
        else:
            return self._build_single_key_json_component(key_config)

    def _build_single_key_json_component(self, key_config: KeyConfig) -> List[pl.Expr]:
        """Build JSON component for a key without variants."""
        # Build value selection expression
        value_expr = self.random_service.build_value_allocation_expression(
            values=key_config.values,
            context=f"k8s_only:{key_config.key_name}"
        )

        # Build JSON key-value pair: "key":"value"
        json_component = pl.concat_str([
            pl.lit('"'),
            pl.lit(key_config.key_name),  # Use canonical key name
            pl.lit('":"'),
            value_expr,
            pl.lit('"')
        ])

        return [json_component]

    def _build_variant_key_json_components(self, key_config: KeyConfig) -> List[pl.Expr]:
        """Build JSON components for a key with variants (one component per variant possibility)."""
        variants = self.key_variant_resolver.variants_config[key_config.key_name]

        # Create resource fingerprint for variant selection
        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(key_config.key_name), pl.lit(":"),
            resource_fingerprint
        ]).hash() % variant_count

        # Build value selection expression
        value_expr = self.random_service.build_value_allocation_expression(
            values=key_config.values,
            context=f"k8s_only:{key_config.key_name}"
        )

        # Build conditional expression that selects the right variant key name
        selected_variant_key = pl.when(variant_index == 0).then(pl.lit(variants[0]))
        for i, variant in enumerate(variants[1:], 1):
            selected_variant_key = selected_variant_key.when(variant_index == i).then(pl.lit(variant))
        selected_variant_key = selected_variant_key.otherwise(pl.lit(variants[0]))  # Fallback

        # Build JSON key-value pair with selected variant: "SelectedVariant":"value"
        json_component = pl.concat_str([
            pl.lit('"'),
            selected_variant_key,  # Use selected variant key name (messy labels!)
            pl.lit('":"'),
            value_expr,
            pl.lit('"')
        ])

        return [json_component]

    def _build_correlated_group_json_components(self, group_config: CorrelatedKeyGroup) -> List[pl.Expr]:
        """Build JSON components for a correlated group."""
        components = []

        # Create fingerprint 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 components 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):
                component = self._build_correlated_variant_json_component(
                    key_name, value_list, correlation_index
                )
            else:
                component = self._build_correlated_single_json_component(
                    key_name, value_list, correlation_index
                )
            components.append(component)

        return components

    def _build_correlated_single_json_component(
        self,
        key_name: str,
        value_list: List[str],
        correlation_index: pl.Expr
    ) -> pl.Expr:
        """Build JSON component for a correlated key without variants."""
        # Build value selection expression based on correlation index
        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
        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)

        # Build JSON key-value pair: "key":"value"
        return pl.concat_str([
            pl.lit('"'),
            pl.lit(key_name),  # Use canonical key name
            pl.lit('":"'),
            final_value_expr,
            pl.lit('"')
        ])

    def _build_correlated_variant_json_component(
        self,
        key_name: str,
        value_list: List[str],
        correlation_index: pl.Expr
    ) -> pl.Expr:
        """Build JSON component for a correlated key with variants."""
        variants = self.key_variant_resolver.variants_config[key_name]

        # Create resource fingerprint for variant selection
        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(key_name), pl.lit(":"),
            resource_fingerprint
        ]).hash() % variant_count

        # Build value selection expression (same as non-variant)
        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
        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)

        # Build conditional expression for variant key selection
        selected_variant_key = pl.when(variant_index == 0).then(pl.lit(variants[0]))
        for i, variant in enumerate(variants[1:], 1):
            selected_variant_key = selected_variant_key.when(variant_index == i).then(pl.lit(variant))
        selected_variant_key = selected_variant_key.otherwise(pl.lit(variants[0]))  # Fallback

        # Build JSON key-value pair with selected variant: "SelectedVariant":"value"
        return pl.concat_str([
            pl.lit('"'),
            selected_variant_key,  # Use selected variant key name (messy labels!)
            pl.lit('":"'),
            final_value_expr,
            pl.lit('"')
        ])

    def _build_merged_json_expression(self, json_components: List[pl.Expr]) -> pl.Expr:
        """Build final merged JSON expression from components."""
        # Build k8s eligibility check
        k8s_eligible_expr = self._build_k8s_eligibility_expression()

        if len(json_components) == 1:
            # Single component: {"key":"value"}
            json_content = json_components[0]
        else:
            # Multiple components: {"key1":"value1","key2":"value2"}
            json_content = json_components[0]
            for component in json_components[1:]:
                json_content = pl.concat_str([
                    json_content,
                    pl.lit(','),  # Add comma separator
                    component
                ])

        # Build final k8s_labels JSON
        k8s_labels_expr = pl.when(k8s_eligible_expr).then(
            pl.concat_str([
                pl.lit('{'),
                json_content,
                pl.lit('}')
            ])
        ).otherwise(pl.lit("{}"))

        return k8s_labels_expr.alias("k8s_labels")

    def _build_empty_k8s_labels_expression(self) -> pl.Expr:
        """Build expression for empty k8s_labels (for non-eligible resources)."""
        return pl.lit("{}").alias("k8s_labels")

    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"))
        )