"""Expression building service for metadata stage."""

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

from ..config import ProviderConfig

logger = logging.getLogger(__name__)


class ExpressionBuilder:
    """
    Builds Polars expressions from metadata configuration.

    Single responsibility: Convert metadata config into Polars expressions.
    Does NOT cache expressions - that's ExpressionCache's job.
    Does NOT apply expressions - that's MetadataProcessor's job.
    """

    def __init__(
        self,
        random_service,
        tags_service,
        four_group_service,
        correlated_service,
        k8s_aggregation_service,
        region_service,
        account_service,
        variant_expression_service,
        resource_id_only_service
    ):
        """
        Initialize with all services needed for expression building.

        Args:
            random_service: DeterministicRandomService for consistent allocation
            tags_service: TagsAllocationService for tag expressions
            four_group_service: FourGroupAllocationService for four-group allocation
            correlated_service: CorrelatedAllocationService for correlated groups
            k8s_aggregation_service: K8sLabelsAggregationService for k8s labels
            region_service: RegionAllocationService for region allocation
            account_service: AccountAllocationService for account allocation
            variant_expression_service: VariantAwareExpressionService for variants
            resource_id_only_service: ResourceIdOnlyService for resource_id_only mode
        """
        self.random_service = random_service
        self.tags_service = tags_service
        self.four_group_service = four_group_service
        self.correlated_service = correlated_service
        self.k8s_aggregation_service = k8s_aggregation_service
        self.region_service = region_service
        self.account_service = account_service
        self.variant_expression_service = variant_expression_service
        self.resource_id_only_service = resource_id_only_service

    def build_expressions_for_provider(
        self,
        provider_config: ProviderConfig,
        provider: str,
        df: pl.DataFrame = None
    ) -> Dict[str, List[pl.Expr]]:
        """
        Build all expressions for a provider, grouped by application strategy.

        Returns expressions grouped into:
        - 'resource_id': Must be applied sequentially (chain injection)
        - 'batch': Can be applied together in one operation

        Args:
            provider_config: Provider configuration with keys and groups
            provider: Provider name (e.g., "aws", "azure")
            df: Optional DataFrame (required for four-group expressions)

        Returns:
            Dictionary with 'resource_id' and 'batch' expression lists
        """
        resource_id_exprs = []
        batch_exprs = []

        # Step 1: Separate k8s-only keys from others
        k8s_keys, other_keys = self._separate_k8s_only_keys(
            provider_config.keys
        )
        k8s_groups, other_groups = self._separate_k8s_only_groups(
            provider_config.correlated_groups
        )

        # Step 2: Build k8s labels expression (if any k8s-only keys)
        if k8s_keys or k8s_groups:
            k8s_expr = self._build_k8s_labels_expression(
                k8s_keys, k8s_groups
            )
            if k8s_expr is not None:
                batch_exprs.append(k8s_expr)
                logger.debug(
                    f"Built k8s_labels expression for {len(k8s_keys)} keys "
                    f"and {len(k8s_groups)} groups"
                )

        # Step 3: Build individual key expressions (separated by type)
        key_resource_id_exprs, key_batch_exprs = self._build_individual_key_expressions(
            df, other_keys, provider
        )
        resource_id_exprs.extend(key_resource_id_exprs)
        batch_exprs.extend(key_batch_exprs)

        # Step 4: Build correlated group expressions (separated by type)
        group_resource_id_exprs, group_batch_exprs = self._build_correlated_group_expressions(
            other_groups, provider
        )
        resource_id_exprs.extend(group_resource_id_exprs)
        batch_exprs.extend(group_batch_exprs)

        logger.debug(
            f"Built {len(resource_id_exprs)} resource_id + {len(batch_exprs)} batch "
            f"expressions for {provider}"
        )

        return {
            'resource_id': resource_id_exprs,
            'batch': batch_exprs
        }

    def _separate_k8s_only_keys(self, keys):
        """
        Separate k8s_labels_only keys from others.

        Args:
            keys: List of KeyConfig objects

        Returns:
            Tuple of (k8s_keys, other_keys)
        """
        k8s_keys = [
            k for k in keys
            if getattr(k, 'k8s_labels_only', False)
        ]
        other_keys = [
            k for k in keys
            if not getattr(k, 'k8s_labels_only', False)
        ]
        return k8s_keys, other_keys

    def _separate_k8s_only_groups(self, groups):
        """
        Separate k8s_labels_only groups from others.

        Args:
            groups: List of CorrelatedKeyGroup objects

        Returns:
            Tuple of (k8s_groups, other_groups)
        """
        k8s_groups = [
            g for g in groups
            if getattr(g, 'k8s_labels_only', False)
        ]
        other_groups = [
            g for g in groups
            if not getattr(g, 'k8s_labels_only', False)
        ]
        return k8s_groups, other_groups

    def _build_k8s_labels_expression(self, k8s_keys, k8s_groups):
        """
        Build single merged k8s_labels expression.

        Args:
            k8s_keys: List of k8s-only KeyConfig objects
            k8s_groups: List of k8s-only CorrelatedKeyGroup objects

        Returns:
            Polars expression or None if no k8s keys/groups
        """
        if not k8s_keys and not k8s_groups:
            return None

        return self.k8s_aggregation_service.build_merged_k8s_labels_expression(
            k8s_keys, k8s_groups
        )

    def _build_individual_key_expressions(self, df, keys, provider):
        """
        Build expressions for individual keys (non-k8s-only).

        Handles region, account, tags_only, resource_id_only, and four-group keys.
        Separates resource_id expressions from batch expressions.

        Args:
            df: DataFrame (required for four-group classification)
            keys: List of KeyConfig objects
            provider: Provider name

        Returns:
            Tuple of (resource_id_expressions, batch_expressions)
        """
        resource_id_exprs = []
        batch_exprs = []

        for key_config in keys:
            if self.region_service.is_region_key(key_config.key_name):
                # Region special case
                region_expr = self.region_service.build_region_expression(
                    key_config
                )
                batch_exprs.append(region_expr.alias("region"))

            elif self.account_service.is_account_key(key_config.key_name):
                # Account special case
                account_exprs = self.account_service.build_account_expressions(
                    key_config, provider
                )
                batch_exprs.extend(account_exprs)

            elif key_config.resource_id_only:
                # Resource ID only scenario - must apply sequentially
                resource_id_expr = self.resource_id_only_service.build_resource_id_only_expression(
                    key_config
                )
                resource_id_exprs.append(resource_id_expr)

            elif key_config.tags_only:
                # Tags-only scenario (cacheable)
                tag_exprs = self.variant_expression_service.build_variant_expressions(
                    key_config, provider
                )
                batch_exprs.extend(tag_exprs)

            else:
                # Four-group scenario (requires DataFrame)
                # Apply sequentially to support chaining (k8s_namespace modifications)
                if df is not None:
                    four_group_exprs = self.four_group_service.build_four_group_expressions(
                        df, key_config, provider
                    )
                    # Add to resource_id_exprs for sequential application
                    resource_id_exprs.extend(four_group_exprs)

        return resource_id_exprs, batch_exprs

    def _build_correlated_group_expressions(self, groups, provider):
        """
        Build expressions for correlated groups (non-k8s-only).

        Separates resource_id expressions from batch expressions.

        Args:
            groups: List of CorrelatedKeyGroup objects
            provider: Provider name

        Returns:
            Tuple of (resource_id_expressions, batch_expressions)
        """
        resource_id_exprs = []
        batch_exprs = []

        for group_config in groups:
            if group_config.resource_id_only:
                # Resource ID only scenario - chain injection, must apply sequentially
                resource_id_group_exprs = self.resource_id_only_service.build_correlated_resource_id_expressions(
                    group_config
                )
                resource_id_exprs.extend(resource_id_group_exprs)
            else:
                # Tags-only or four-group scenario
                correlated_exprs = self.correlated_service.build_correlated_expressions(
                    group_config, provider
                )
                batch_exprs.extend(correlated_exprs)

        return resource_id_exprs, batch_exprs
