"""Resource ID Only allocation service for metadata stage."""

import polars as pl
from typing import List
from ..config import KeyConfig, CorrelatedKeyGroup
from .random_service import DeterministicRandomService
from .resource_id_service import ResourceIdInjectionService


class ResourceIdOnlyService:
    """
    Service for resource_id_only allocation mode.

    Injects metadata values into resource_id column for 100% of resources
    (including K8s resources). No tag columns or K8s label modifications.

    Uses composition with ResourceIdInjectionService for injection logic.
    """

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

        Args:
            random_service: Service for deterministic value allocation
        """
        self.random_service = random_service
        self.injection_service = ResourceIdInjectionService(random_service)

    def build_resource_id_only_expression(
        self,
        key_config: KeyConfig
    ) -> pl.Expr:
        """
        Build expression to prepend value to resource ID part only.

        Preserves type/id format by only modifying the ID part after '/'.
        Preprocessing creates type/id format, this expression maintains it.

        Args:
            key_config: Key configuration with values to inject

        Returns:
            Polars expression that modifies resource_id column

        Example:
            Expression transforms: service-usage/abc123 → service-usage/value-abc123
        """
        # Build expression that allocates value to all resources
        value_expr = self.random_service.build_value_allocation_expression(
            values=key_config.values,
            context=f"resource_id_only:{key_config.key_name}"
        )

        # Split on both : and / to handle AWS (: separator) and other providers (/ separator)
        segments_colon = pl.col("resource_id").str.split(":")
        segments_slash = pl.col("resource_id").str.split("/")

        # Use : split for AWS, / split for others
        type_part = pl.when(pl.col("cloud_provider") == "AWS") \
            .then(segments_colon.list.get(0)) \
            .otherwise(segments_slash.list.get(0))

        id_part = pl.when(pl.col("cloud_provider") == "AWS") \
            .then(segments_colon.list.get(-1)) \
            .otherwise(segments_slash.list.get(-1))

        # Reconstruct with appropriate separator
        separator = pl.when(pl.col("cloud_provider") == "AWS") \
            .then(pl.lit(":")) \
            .otherwise(pl.lit("/"))

        # Prepend value to ID part only, reconstruct with appropriate separator
        return (
            pl.concat_str([
                type_part,
                separator,
                value_expr,
                pl.lit("-"),
                id_part
            ])
            .alias("resource_id")
        )

    def build_correlated_resource_id_expressions(
        self,
        group_config: CorrelatedKeyGroup
    ) -> List[pl.Expr]:
        """
        Build expressions for chain injection of correlated values.

        Preserves type/id format by only modifying the ID part after '/'.
        Multiple correlated values are prepended sequentially into ID part:
        1. Preprocessed: service-usage/abc123
        2. After org: service-usage/org-x-abc123
        3. After project: service-usage/proj-a-org-x-abc123

        Keys are processed in sorted order for deterministic behavior.
        Result shows values in reverse injection order (most recent first).

        Args:
            group_config: Correlated key group configuration

        Returns:
            List of expressions (one per correlated key in sorted order)

        Example:
            For keys {org: [x, y], project: [a, b]}, returns 2 expressions
            that sequentially prepend org then project values to ID part only.
        """
        expressions = []

        # Sort keys for deterministic ordering
        sorted_keys = sorted(group_config.keys.keys())

        for key_name in sorted_keys:
            values = group_config.keys[key_name]

            # Build expression that allocates value from correlated set
            value_expr = self.random_service.build_value_allocation_expression(
                values=values,
                context=f"resource_id_only:correlated:{key_name}"
            )

            # Split on both : and / to handle AWS and other providers
            segments_colon = pl.col("resource_id").str.split(":")
            segments_slash = pl.col("resource_id").str.split("/")

            type_part = pl.when(pl.col("cloud_provider") == "AWS") \
                .then(segments_colon.list.get(0)) \
                .otherwise(segments_slash.list.get(0))

            id_part = pl.when(pl.col("cloud_provider") == "AWS") \
                .then(segments_colon.list.get(-1)) \
                .otherwise(segments_slash.list.get(-1))

            # Reconstruct with appropriate separator
            separator = pl.when(pl.col("cloud_provider") == "AWS") \
                .then(pl.lit(":")) \
                .otherwise(pl.lit("/"))

            # Prepend value to ID part only (chain injection)
            chain_expr = (
                pl.concat_str([
                    type_part,
                    separator,
                    value_expr,
                    pl.lit("-"),
                    id_part
                ])
                .alias("resource_id")
            )

            expressions.append(chain_expr)

        return expressions
