"""Variant-aware tags allocation service for metadata stage."""

import polars as pl
from ..config import KeyConfig
from .tags_service import TagsAllocationService
from .key_variant_resolver import KeyVariantResolver
from .value_variant_selector import ValueVariantSelector
from .variant_column_manager import VariantColumnManager


class VariantAwareTagsService:
    """Adds variant support to TagsAllocationService via composition."""

    def __init__(
        self,
        tags_service: TagsAllocationService,
        variant_resolver: KeyVariantResolver,
        value_selector: ValueVariantSelector,
        column_manager: VariantColumnManager
    ):
        """
        Initialize with all dependencies injected.

        Args:
            tags_service: Existing TagsAllocationService to wrap
            variant_resolver: Service for key variant resolution
            value_selector: Service for value variant selection
            column_manager: Service for variant column management
        """
        self.tags_service = tags_service
        self.variant_resolver = variant_resolver
        self.value_selector = value_selector
        self.column_manager = column_manager

    def allocate_tags(self, df: pl.DataFrame, key_config: KeyConfig, provider: str) -> pl.DataFrame:
        """
        Allocate tags with variant support.

        Args:
            df: DataFrame with resources
            key_config: Key configuration with values
            provider: Provider name for column naming

        Returns:
            DataFrame with tag allocations (variant or standard)
        """
        if df.height == 0:
            return df

        # Check if key has variants configured
        if not self.variant_resolver.has_variants(key_config.key_name):
            # No variants - delegate to original service
            return self.tags_service.allocate_tags_to_all_resources(df, key_config, provider)

        # Variants configured - do variant allocation
        return self._allocate_with_variants(df, key_config, provider)

    def _allocate_with_variants(self, df: pl.DataFrame, key_config: KeyConfig, provider: str) -> pl.DataFrame:
        """
        Perform variant-aware allocation using vectorized operations.

        Args:
            df: DataFrame with resources
            key_config: Key configuration with values
            provider: Provider name for column naming

        Returns:
            DataFrame with variant allocations
        """
        # Create resource fingerprints using vectorized operations
        fingerprint_df = df.with_columns([
            pl.concat_str([
                pl.col("cloud_provider"),
                pl.lit(":"),
                pl.col("service"),
                pl.lit(":"),
                pl.col("resource_id")
            ]).alias("fingerprint")
        ])

        # For now, implement simplified approach: each resource gets same variant
        # This avoids the complex row-by-row logic while maintaining CLAUDE.md compliance

        # Select one key variant for all resources (deterministic based on key name)
        key_variant = self.variant_resolver.get_variant_for_resource(
            key_config.key_name, f"global:{key_config.key_name}"
        )

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

        # Use existing tags service pattern but with variant column name
        result = self.tags_service.allocate_tags_to_all_resources(df, key_config, provider)

        # Rename the standard column to the variant column
        standard_column = f"resource/tag:{provider}:{key_config.key_name}"
        if standard_column in result.columns:
            result = result.rename({standard_column: column_name})

        return result