"""K8s labels service for metadata stage."""

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


class K8sLabelsService:
    """Service for adding labels to k8s_labels JSON (tags_only=false scenario)."""

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

    def add_k8s_labels(self, df: pl.DataFrame, key_config: KeyConfig) -> pl.DataFrame:
        """
        Add key-value pairs to k8s_labels JSON column.

        Args:
            df: DataFrame with K8s resources (must have k8s_labels column)
            key_config: Key configuration with values to add

        Returns:
            DataFrame with updated k8s_labels JSON
        """
        if df.height == 0:
            return df

        if "k8s_labels" not in df.columns:
            raise ValueError("DataFrame must have 'k8s_labels' column for K8s labels injection")

        # Use deterministic random service to select values
        df_with_values = self.random_service.allocate_values_to_resources(
            df=df,
            values=key_config.values,
            context=f"k8s_labels:{key_config.key_name}",
            column_name=f"_k8s_value_{key_config.key_name}"
        )

        # Update k8s_labels JSON with new key-value pairs
        value_column = f"_k8s_value_{key_config.key_name}"

        result = df_with_values.with_columns([
            self._merge_k8s_labels_expression(key_config.key_name, value_column).alias("k8s_labels")
        ]).drop(value_column)

        return result

    def _merge_k8s_labels_expression(self, key_name: str, value_column: str) -> pl.Expr:
        """
        Create Polars expression for merging new key-value into k8s_labels JSON.

        Args:
            key_name: Label key to add
            value_column: Column containing the value to add

        Returns:
            Polars expression for updated k8s_labels JSON
        """
        def merge_labels_udf(existing_json: str, new_value: str) -> str:
            """Merge new label into existing k8s_labels JSON."""
            try:
                if not existing_json or existing_json.strip() == "":
                    labels = {}
                else:
                    labels = json.loads(existing_json)
            except (json.JSONDecodeError, TypeError):
                # If existing JSON is invalid, start fresh
                labels = {}

            # Add new label
            labels[key_name] = new_value

            # Return sorted JSON for consistency
            return json.dumps(labels, sort_keys=True)

        # Use map_elements for the JSON manipulation
        return pl.struct([pl.col("k8s_labels"), pl.col(value_column)]).map_elements(
            lambda x: merge_labels_udf(x["k8s_labels"], x[value_column]),
            return_dtype=pl.Utf8
        )