"""Region allocation service for metadata stage."""

import polars as pl
from ..config import KeyConfig
from .random_service import DeterministicRandomService


class RegionAllocationService:
    """Service for region special case (updates region column directly)."""

    def __init__(self, random_service: DeterministicRandomService):
        """Initialize with deterministic random service for consistent allocation."""
        self.random_service = random_service

    def allocate_region_column(self, df: pl.DataFrame, key_config: KeyConfig) -> pl.DataFrame:
        """
        Update or add the region column directly with allocated values.

        Special case: Region keys (key_name == "region" or "Region") ignore the tags_only flag
        and always update the region column directly instead of creating tag columns.

        Args:
            df: DataFrame with resources
            key_config: Key configuration with region values (tags_only flag is ignored)

        Returns:
            DataFrame with updated/added region column
        """
        if df.height == 0:
            return df

        # Use deterministic random service to allocate region values
        return self.random_service.allocate_values_to_resources(
            df=df,
            values=key_config.values,
            context=f"region_allocation:{key_config.key_name}",
            column_name="region"
        )

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

        Args:
            key_name: Key name to check

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

    def build_region_expression(self, key_config: 'KeyConfig') -> pl.Expr:
        """
        Build Polars expression for region allocation without creating intermediate DataFrames.

        Args:
            key_config: Key configuration with region values

        Returns:
            Polars expression that can be used in with_columns()
        """
        # Build expression using the same logic as random service
        return self.random_service.build_value_allocation_expression(
            values=key_config.values,
            context=f"region_allocation:{key_config.key_name}"
        )