"""Variant-aware expression service for metadata stage."""

import polars as pl
from typing import List
from ..config import KeyConfig
from .tags_service import TagsAllocationService
from .key_variant_resolver import KeyVariantResolver


class VariantAwareExpressionService:
    """Provides variant-aware expressions for metadata processing."""

    def __init__(
        self,
        tags_service: TagsAllocationService,
        variant_resolver: KeyVariantResolver
    ):
        """
        Initialize with dependencies.

        Args:
            tags_service: Base tags service for non-variant expressions
            variant_resolver: Service for key variant resolution
        """
        self.tags_service = tags_service
        self.variant_resolver = variant_resolver

    def build_tags_expression(self, key_config: KeyConfig, provider: str) -> pl.Expr:
        """
        Build Polars expression for tags allocation with variant support.

        Note: For single-column compatibility. Use build_variant_expressions()
        for multi-column variant support.

        Args:
            key_config: Key configuration with values to allocate
            provider: Provider name for column naming

        Returns:
            Polars expression that can be used in with_columns()
        """
        if not self.variant_resolver.has_variants(key_config.key_name):
            # No variants - delegate to original service
            return self.tags_service.build_tags_expression(key_config, provider)

        # Variants configured - create single variant expression (legacy behavior)
        return self._build_variant_expression(key_config, provider)

    def build_variant_expressions(self, key_config: KeyConfig, provider: str) -> List[pl.Expr]:
        """
        Build multiple Polars expressions for multi-column variant support.

        Creates one column per variant with distributed values across resources.
        Each resource gets a value in exactly one variant column.

        Args:
            key_config: Key configuration with values to allocate
            provider: Provider name for column naming

        Returns:
            List of Polars expressions, one per variant column
        """
        if not self.variant_resolver.has_variants(key_config.key_name):
            # No variants - return single expression
            return [self.tags_service.build_tags_expression(key_config, provider)]

        # Variants configured - create multiple variant expressions
        return self._build_multi_variant_expressions(key_config, provider)

    def _build_variant_expression(self, key_config: KeyConfig, provider: str) -> pl.Expr:
        """
        Build expression using variant key name.

        Args:
            key_config: Key configuration
            provider: Provider name

        Returns:
            Polars expression with variant column name
        """
        # Normalize provider name for consistent column naming
        normalized_provider = self.tags_service._normalize_provider_name(provider)

        # Get deterministic variant for this key (global scope for now)
        key_variant = self.variant_resolver.get_variant_for_resource(
            key_config.key_name, f"global:{key_config.key_name}"
        )

        # Create variant column name
        column_name = f"resource/tag:{normalized_provider}:{key_variant}"

        # Build expression using same allocation logic but with variant column name
        return self.tags_service.random_service.build_value_allocation_expression(
            values=key_config.values,
            context=f"tags_all:{key_config.key_name}"
        ).alias(column_name)

    def _build_multi_variant_expressions(self, key_config: KeyConfig, provider: str) -> List[pl.Expr]:
        """
        Build multiple expressions for variant columns with distributed values.

        Each resource gets a value in exactly one variant column based on
        deterministic resource fingerprinting for cross-month consistency.

        Args:
            key_config: Key configuration with values to allocate
            provider: Provider name for column naming

        Returns:
            List of expressions, one per variant column
        """
        # Normalize provider name for consistent column naming
        normalized_provider = self.tags_service._normalize_provider_name(provider)

        # Get all variants for this key
        variants = self.variant_resolver.variants_config[key_config.key_name]

        # Build base value allocation expression (same for all variants)
        base_value_expr = self.tags_service.random_service.build_value_allocation_expression(
            values=key_config.values,
            context=f"tags_all:{key_config.key_name}"
        )

        expressions = []

        for i, variant in enumerate(variants):
            # Create deterministic condition for this variant
            # Use resource_id + canonical_key for cross-month consistency
            variant_condition = self._build_variant_selection_condition(
                key_config.key_name, i, len(variants)
            )

            # Create conditional expression: value if selected for this variant, null otherwise
            variant_expr = pl.when(variant_condition).then(base_value_expr).otherwise(None)

            # Create column name for this variant
            column_name = f"resource/tag:{normalized_provider}:{variant}"

            expressions.append(variant_expr.alias(column_name))

        return expressions

    def _build_variant_selection_condition(self, canonical_key: str, variant_index: int, total_variants: int) -> pl.Expr:
        """
        Build condition that deterministically selects resources for a specific variant.

        Uses resource_id hashing for cross-month consistency.

        Args:
            canonical_key: The canonical key name (e.g., "cost_center")
            variant_index: Index of this variant (0-based)
            total_variants: Total number of variants

        Returns:
            Polars expression that evaluates to True for resources assigned to this variant
        """
        # Create deterministic selection based on resource_id + canonical_key
        # This ensures same resource gets same variant across months
        fingerprint_expr = pl.col("resource_id").cast(pl.String) + f":{canonical_key}"

        # Hash the fingerprint and take modulo to select variant
        hash_expr = fingerprint_expr.hash()
        variant_selection_expr = hash_expr % total_variants

        # Return condition: this resource is assigned to this variant index
        return variant_selection_expr == variant_index