"""Four-group allocation service orchestrator for metadata stage."""

import polars as pl
from typing import List, Optional
from ..config import KeyConfig
from .random_service import DeterministicRandomService
from .classification_service import ResourceClassificationService
from .tags_service import TagsAllocationService
from .resource_id_service import ResourceIdInjectionService
from .k8s_labels_service import K8sLabelsService
from .k8s_namespace_service import K8sNamespaceInjectionService
from .gcp_tags_service import GcpTagsService
from .key_variant_resolver import KeyVariantResolver
from .value_sanitizer import K8sValueSanitizer
from src.k8s.validation.name_validator import K8sNameValidator


class FourGroupAllocationService:
    """Orchestrates four-group allocation (tags_only=false scenario)."""

    def __init__(
        self,
        random_service: DeterministicRandomService,
        key_variant_resolver: Optional[KeyVariantResolver] = None
    ):
        """Initialize with all required services."""
        self.random_service = random_service
        self.key_variant_resolver = key_variant_resolver
        self.classification_service = ResourceClassificationService(random_service)
        self.tags_service = TagsAllocationService(random_service)
        self.resource_id_service = ResourceIdInjectionService(random_service)
        self.gcp_tags_service = GcpTagsService()
        self.k8s_labels_service = K8sLabelsService(random_service)

        # Create validator and sanitizer for k8s naming compliance
        name_validator = K8sNameValidator()
        value_sanitizer = K8sValueSanitizer(name_validator)

        self.k8s_namespace_service = K8sNamespaceInjectionService(random_service, value_sanitizer)

    def allocate_four_groups(
        self,
        df: pl.DataFrame,
        key_config: KeyConfig,
        provider: str
    ) -> pl.DataFrame:
        """
        Apply four-group allocation to resources.

        Process:
        1. Split into K8s vs Non-K8s resources
        2. K8s resources: Split custom namespaces into labels vs namespace injection groups
        3. Non-K8s resources: Split into tags vs resource ID injection groups
        4. Apply appropriate enrichment to each group
        5. Combine results

        Args:
            df: DataFrame with provider-month resources
            key_config: Key configuration (with tags_only=false)
            provider: Provider name for tag column naming

        Returns:
            DataFrame with four-group allocation applied
        """
        if df.height == 0:
            return df

        # Step 1: Classify into K8s vs Non-K8s
        k8s_resources, non_k8s_resources = self.classification_service.classify_k8s_vs_non_k8s(df)

        processed_groups = []

        # Step 2: Process K8s resources if present
        if k8s_resources.height > 0:
            k8s_labels_group, k8s_namespace_group = self.classification_service.split_k8s_namespaces(k8s_resources)

            # Preserve default/Idle namespaces (pass through unchanged)
            default_namespaces = self.classification_service.get_default_namespaces()
            k8s_default_idle = k8s_resources.filter(
                pl.col("k8s_namespace").is_in(default_namespaces)
            )

            if k8s_default_idle.height > 0:
                processed_groups.append(k8s_default_idle)

            # Process K8s labels group (custom namespaces → labels)
            if k8s_labels_group.height > 0:
                # Ensure k8s_labels column exists
                if "k8s_labels" not in k8s_labels_group.columns:
                    k8s_labels_group = k8s_labels_group.with_columns([
                        pl.lit("{}").alias("k8s_labels")
                    ])

                processed_labels = self.k8s_labels_service.add_k8s_labels(k8s_labels_group, key_config)
                processed_groups.append(processed_labels)

            # Process K8s namespace injection group (custom namespaces → namespace modification)
            if k8s_namespace_group.height > 0:
                processed_namespaces = self.k8s_namespace_service.inject_into_namespaces(k8s_namespace_group, key_config)
                processed_groups.append(processed_namespaces)

        # Step 3: Process Non-K8s resources if present
        if non_k8s_resources.height > 0:
            tags_group, resource_id_group = self.classification_service.split_non_k8s_resources(non_k8s_resources)

            # Process tags group (non-k8s → resource/tag columns)
            if tags_group.height > 0:
                processed_tags = self.tags_service.allocate_tags_to_all_resources(tags_group, key_config, provider)
                processed_groups.append(processed_tags)

            # Process resource ID injection group (non-k8s → resource_id modification)
            if resource_id_group.height > 0:
                processed_resource_ids = self.resource_id_service.inject_into_resource_ids(resource_id_group, key_config)
                # Sync GCP resource_name tag if this is GCP provider
                if provider.lower() in ['gcp', 'google cloud platform', 'google cloud', 'google']:
                    processed_resource_ids = self.gcp_tags_service.sync_resource_name_with_resource_id(processed_resource_ids)
                processed_groups.append(processed_resource_ids)

        # Step 4: Combine all processed groups
        if processed_groups:
            # Ensure all DataFrames have the same columns by adding missing columns
            result = self._harmonize_and_combine_groups(processed_groups)
        else:
            result = df.clone()

        return result

    def _harmonize_and_combine_groups(self, groups: list[pl.DataFrame]) -> pl.DataFrame:
        """
        Harmonize schemas and combine DataFrames from different processing groups.

        Args:
            groups: List of processed DataFrames from different groups

        Returns:
            Combined DataFrame with harmonized schema
        """
        if not groups:
            return pl.DataFrame()

        if len(groups) == 1:
            return groups[0]

        # Get all unique columns across all groups
        all_columns = set()
        for group in groups:
            all_columns.update(group.columns)

        # Add missing columns to each group
        harmonized_groups = []
        for group in groups:
            harmonized = group.clone()

            for col in all_columns:
                if col not in harmonized.columns:
                    # Add missing column with appropriate default value
                    if col == "k8s_labels":
                        harmonized = harmonized.with_columns([pl.lit("{}").alias(col)])
                    elif col.startswith("resource/tag:"):
                        harmonized = harmonized.with_columns([pl.lit(None, dtype=pl.Utf8).alias(col)])
                    else:
                        # Try to infer dtype from existing groups
                        existing_dtype = None
                        for other_group in groups:
                            if col in other_group.columns:
                                existing_dtype = other_group.schema[col]
                                break

                        if existing_dtype:
                            harmonized = harmonized.with_columns([pl.lit(None, dtype=existing_dtype).alias(col)])
                        else:
                            harmonized = harmonized.with_columns([pl.lit(None, dtype=pl.Utf8).alias(col)])

            harmonized_groups.append(harmonized)

        # Combine all harmonized groups
        return pl.concat(harmonized_groups, how="vertical")

    def get_allocation_summary(self, df: pl.DataFrame, key_config: KeyConfig) -> dict:
        """
        Get summary of four-group allocation for debugging/logging.

        Args:
            df: Original DataFrame
            key_config: Key configuration being processed

        Returns:
            Dictionary with allocation statistics
        """
        classification_summary = self.classification_service.get_processing_summary(df)

        return {
            "key_name": key_config.key_name,
            "values_count": len(key_config.values),
            "allocation_type": "four_group",
            "classification": classification_summary
        }

    def build_four_group_expressions(
        self,
        df: pl.DataFrame,
        key_config: KeyConfig,
        provider: str
    ) -> List[pl.Expr]:
        """
        Build Polars expressions for allocation based on key configuration mode.

        Args:
            df: DataFrame for analysis (used for potential classification)
            key_config: Key configuration with values and allocation mode
            provider: Provider name for column naming

        Returns:
            List of Polars expressions that can be used in with_columns()

        Note: For tags_only=false (four-group), returns expressions that should
        be applied sequentially to support chaining (k8s_namespace modifications).
        """
        if key_config.k8s_labels_only:
            return self._build_k8s_only_expressions(df, key_config)
        elif key_config.tags_only:
            return self._build_tags_only_expressions(key_config, provider)
        else:
            # Default four-group behavior (both tags and k8s)
            # Returns all expressions as a list (will be applied sequentially)
            return self._build_both_expressions(key_config, provider)

    def _build_k8s_only_expressions(self, df: pl.DataFrame, key_config: KeyConfig) -> List[pl.Expr]:
        """
        Build expressions that create k8s labels only (no resource tags).

        Only k8s-eligible resources get k8s labels; others get empty k8s_labels.
        """
        if self.key_variant_resolver and self.key_variant_resolver.has_variants(key_config.key_name):
            return self._build_k8s_variant_expressions(key_config)
        else:
            return self._build_single_k8s_expression(key_config)

    def _build_k8s_variant_expressions(self, key_config: KeyConfig) -> List[pl.Expr]:
        """Build k8s-only expressions with key variants."""
        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 k8s eligibility check
        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(
            # Create JSON with the selected variant key
            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('":"'),
                value_expr,
                pl.lit('"}')
            ])
        ).otherwise(pl.lit("{}"))

        return [k8s_labels_expr.alias("k8s_labels")]

    def _build_single_k8s_expression(self, key_config: KeyConfig) -> List[pl.Expr]:
        """Build a single k8s-only expression 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 k8s eligibility check
        k8s_eligible_expr = self._build_k8s_eligibility_expression()

        # Create k8s_labels JSON for eligible resources only
        k8s_labels_expr = pl.when(k8s_eligible_expr).then(
            pl.concat_str([
                pl.lit('{"'),
                pl.lit(key_config.key_name),
                pl.lit('":"'),
                value_expr,
                pl.lit('"}')
            ])
        ).otherwise(pl.lit("{}"))

        return [k8s_labels_expr.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"))
        )

    def _build_tags_only_expressions(self, key_config: KeyConfig, provider: str) -> List[pl.Expr]:
        """Build tags-only expressions (existing behavior)."""
        if self.key_variant_resolver and self.key_variant_resolver.has_variants(key_config.key_name):
            return self._build_tags_variant_expressions(key_config, provider)
        else:
            return self._build_single_tag_expression(key_config, provider)

    def _build_tags_variant_expressions(self, key_config: KeyConfig, provider: str) -> List[pl.Expr]:
        """Build tags-only expressions with key variants."""
        expressions = []
        variants = self.key_variant_resolver.variants_config[key_config.key_name]
        normalized_provider = self.tags_service._normalize_provider_name(provider)

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

        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"tags_only:{key_config.key_name}"
        )

        # Create expressions for each variant column
        for i, variant in enumerate(variants):
            column_name = f"resource/tag:{normalized_provider}:{variant}"
            variant_expr = pl.when(variant_index == i).then(value_expr).otherwise(None)
            expressions.append(variant_expr.alias(column_name))

        return expressions

    def _build_single_tag_expression(self, key_config: KeyConfig, provider: str) -> List[pl.Expr]:
        """Build a single tags-only expression without variants."""
        normalized_provider = self.tags_service._normalize_provider_name(provider)
        column_name = f"resource/tag:{normalized_provider}:{key_config.key_name}"

        tag_expr = self.random_service.build_value_allocation_expression(
            values=key_config.values,
            context=f"tags_only:{key_config.key_name}"
        )

        return [tag_expr.alias(column_name)]

    def _build_both_expressions(self, key_config: KeyConfig, provider: str) -> List[pl.Expr]:
        """
        Build expressions for both resource tags AND k8s labels (default four-group behavior).

        Implements four-group allocation:
        1. K8s resources with custom namespaces: Split 50/50 between k8s labels and namespace injection
        2. Non-k8s resources: Split 50/50 between tags and resource ID injection

        Note: Resource ID injection is handled separately in resource_id_exprs (not here).
        This method handles: k8s labels, namespace injection, and tags.

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

        Returns:
            List of expressions for k8s_labels, k8s_namespace, and resource/tag columns
        """
        expressions = []

        # Sanitize values for k8s namespace injection
        from .value_sanitizer import K8sValueSanitizer
        from src.k8s.validation.name_validator import K8sNameValidator

        sanitizer = K8sValueSanitizer(K8sNameValidator())
        sanitized_values = sanitizer.sanitize_values(key_config.values)

        # Classification expressions
        # K8s eligibility: Has non-null, non-empty, non-UNALLOCATED cluster
        is_k8s = (
            pl.col("k8s_cluster").is_not_null() &
            (pl.col("k8s_cluster") != "") &
            (pl.col("k8s_cluster") != "__UNALLOCATED__")
        )

        # Custom namespace check (exclude default/system namespaces)
        default_namespaces = ["default", "kube-system", "kube-public", "kube-node-lease", "__UNALLOCATED__"]
        is_custom_namespace = (
            pl.col("k8s_namespace").is_not_null() &
            (pl.col("k8s_namespace") != "") &
            ~pl.col("k8s_namespace").is_in(default_namespaces)
        )

        # Deterministic 50/50 split for k8s custom namespaces (by namespace, not resource)
        namespace_hash = pl.concat_str([
            pl.lit(f"{self.random_service.dataset_name}:k8s_namespace_split:"),
            pl.col("k8s_namespace")
        ]).hash() % 2

        # Group assignments for k8s resources
        is_k8s_labels_group = is_k8s & is_custom_namespace & (namespace_hash == 0)
        is_k8s_namespace_group = is_k8s & is_custom_namespace & (namespace_hash == 1)

        # Deterministic 50/50 split for non-k8s resources (by resource)
        is_non_k8s = ~is_k8s
        resource_hash = pl.concat_str([
            pl.lit(f"{self.random_service.dataset_name}:non_k8s_split:"),
            pl.col("cloud_provider"), pl.lit(":"),
            pl.col("service"), pl.lit(":"),
            pl.col("resource_id")
        ]).hash() % 2

        # Group assignments for non-k8s resources
        is_tags_group = is_non_k8s & (resource_hash == 0)
        # is_resource_id_group = is_non_k8s & (resource_hash == 1)  # Handled separately

        # Build value allocation expressions
        # For k8s labels and tags: per-resource allocation
        resource_value_expr = self.random_service.build_value_allocation_expression(
            values=key_config.values,
            context=f"four_group:{key_config.key_name}"
        )

        # For namespace injection: per-namespace allocation (use namespace as seed)
        # Create fingerprint that includes namespace for consistency
        namespace_fingerprint = pl.concat_str([
            pl.lit(f"{self.random_service.dataset_name}:k8s_namespace_injection:{key_config.key_name}:"),
            pl.col("k8s_namespace")
        ])

        # Select value based on namespace fingerprint hash
        namespace_value_index = namespace_fingerprint.hash() % len(sanitized_values)

        # Build case expression for value selection
        namespace_value_expr = pl.when(namespace_value_index == 0).then(pl.lit(sanitized_values[0]))
        for i in range(1, len(sanitized_values)):
            namespace_value_expr = namespace_value_expr.when(namespace_value_index == i).then(pl.lit(sanitized_values[i]))
        namespace_value_expr = namespace_value_expr.otherwise(pl.lit(sanitized_values[0]))

        # Expression 1: K8s labels (for k8s labels group only)
        # Merge with existing k8s_labels JSON if present
        # Use fill_null to handle case where column might not exist
        existing_labels = pl.col("k8s_labels").fill_null("{}")
        has_existing = (
            (existing_labels != "") &
            (existing_labels != "{}")
        )

        k8s_labels_expr = pl.when(is_k8s_labels_group).then(
            pl.when(has_existing).then(
                # Merge: Remove trailing }, add comma, append new key-value, close }
                pl.concat_str([
                    existing_labels.str.slice(0, existing_labels.str.len_chars() - 1),
                    pl.lit(',"'),
                    pl.lit(key_config.key_name),
                    pl.lit('":"'),
                    resource_value_expr,
                    pl.lit('"}')
                ])
            ).otherwise(
                # New: Create fresh JSON
                pl.concat_str([
                    pl.lit('{"'),
                    pl.lit(key_config.key_name),
                    pl.lit('":"'),
                    resource_value_expr,
                    pl.lit('"}')
                ])
            )
        ).otherwise(existing_labels)  # Keep existing for non-labels-group
        expressions.append(k8s_labels_expr.alias("k8s_labels"))

        # Expression 2: Namespace injection (for k8s namespace injection group only)
        # Inject sanitized value into namespace name
        namespace_expr = pl.when(is_k8s_namespace_group).then(
            pl.concat_str([
                pl.col("k8s_namespace"),
                pl.lit("-"),
                namespace_value_expr
            ])
        ).otherwise(pl.col("k8s_namespace"))
        expressions.append(namespace_expr.alias("k8s_namespace"))

        # Expression 3: Tags (for non-k8s tags group only)
        normalized_provider = self.tags_service._normalize_provider_name(provider)
        column_name = f"resource/tag:{normalized_provider}:{key_config.key_name}"
        tag_expr = pl.when(is_tags_group).then(resource_value_expr).otherwise(None)
        expressions.append(tag_expr.alias(column_name))

        # Note: Resource ID injection for is_resource_id_group is handled separately
        # in resource_id_exprs (requires sequential application, not batch)

        return expressions