"""Account allocation service for metadata stage."""

import polars as pl
from typing import List, Optional
from ..config import KeyConfig
from .random_service import DeterministicRandomService
from .azure_resource_id_service import AzureResourceIdService


class AccountAllocationService:
    """Service for account special case (updates cloud_account column + provider-specific tags)."""

    def __init__(
        self,
        random_service: DeterministicRandomService,
        azure_resource_id_service: Optional[AzureResourceIdService] = None
    ):
        """Initialize with deterministic random service and optional Azure resource ID service."""
        self.random_service = random_service
        self.azure_resource_id_service = azure_resource_id_service

    def allocate_account_column(self, df: pl.DataFrame, key_config: KeyConfig, provider: str) -> pl.DataFrame:
        """
        Update cloud_account column and provider-specific account tags.
        Preserves resource_id for cross-month consistency.

        All providers support account special case and update cloud_account column.
        Azure and GCP have additional tag requirements.

        Provider-specific behavior:
        - All providers: Updates cloud_account column
        - Azure: Additionally updates resource/tag:azure:subscriptionName
        - GCP: Additionally updates resource/tag:gcp_cz:project_name

        Args:
            df: DataFrame with resources
            key_config: Account key configuration with values from config
            provider: Provider name from config (normalized: aws, azure, gcp, etc.)

        Returns:
            DataFrame with updated cloud_account column and provider-specific tags
        """
        if df.height == 0:
            return df

        # Step 1: Update cloud_account column for all providers
        result_df = self.random_service.allocate_values_to_resources(
            df=df,
            values=key_config.values,
            context=f"account_allocation:{key_config.key_name}:{provider}",
            column_name="cloud_account"
        )

        # Step 2: Handle provider-specific tag requirements
        provider_lower = provider.lower()

        if provider_lower == "azure":
            # Azure requires resource/tag:azure:subscriptionName to match cloud_account
            result_df = self._update_azure_subscription_tag(result_df, key_config, provider)

            # Azure: Additionally update subscription ID in resource_id column
            if self.azure_resource_id_service and "resource_id" in result_df.columns:
                result_df = self.azure_resource_id_service.update_subscription_in_resource_id(result_df)

        elif provider_lower == "gcp":
            # GCP requires resource/tag:gcp_cz:project_name to match cloud_account
            result_df = self._update_gcp_project_tag(result_df, key_config, provider)

        # AWS, Anthropic, OpenAI, and other providers only need cloud_account column updated

        return result_df

    def _update_azure_subscription_tag(self, df: pl.DataFrame, key_config: KeyConfig, provider: str) -> pl.DataFrame:
        """Update Azure subscriptionName tag to match cloud_account column."""
        # Use same context as cloud_account to ensure identical allocation
        return self.random_service.allocate_values_to_resources(
            df=df,
            values=key_config.values,
            context=f"account_allocation:{key_config.key_name}:{provider}",
            column_name="resource/tag:azure:subscriptionName"
        )

    def _update_gcp_project_tag(self, df: pl.DataFrame, key_config: KeyConfig, provider: str) -> pl.DataFrame:
        """Update GCP project_name tag to match cloud_account column."""
        # Use same context as cloud_account to ensure identical allocation
        return self.random_service.allocate_values_to_resources(
            df=df,
            values=key_config.values,
            context=f"account_allocation:{key_config.key_name}:{provider}",
            column_name="resource/tag:gcp_cz:project_name"
        )

    def is_account_key(self, key_name: str) -> bool:
        """
        Check if a key name represents an account configuration.

        Args:
            key_name: Key name to check

        Returns:
            True if this is an account key, False otherwise
        """
        return key_name.lower() == "account"

    def build_account_expressions(self, key_config: 'KeyConfig', provider: str) -> List[pl.Expr]:
        """
        Build Polars expressions for account allocation without creating intermediate DataFrames.

        Args:
            key_config: Key configuration with account values
            provider: Provider name for provider-specific tag handling

        Returns:
            List of Polars expressions that can be used in with_columns()
        """
        expressions = []
        context = f"account_allocation:{key_config.key_name}:{provider}"

        # Build cloud_account column expression
        account_expr = self.random_service.build_value_allocation_expression(
            values=key_config.values,
            context=context
        )
        expressions.append(account_expr.alias("cloud_account"))

        # Handle provider-specific tag requirements
        provider_lower = provider.lower()

        if provider_lower == "azure":
            # Azure requires resource/tag:azure:subscriptionName to match cloud_account
            expressions.append(account_expr.alias("resource/tag:azure:subscriptionName"))

        elif provider_lower == "gcp":
            # GCP requires resource/tag:gcp_cz:project_name to match cloud_account
            expressions.append(account_expr.alias("resource/tag:gcp_cz:project_name"))

        # AWS, Anthropic, OpenAI, and other providers only need cloud_account column

        return expressions