"""Variant column management service for metadata stage."""

import polars as pl
from typing import Dict, List


class VariantColumnManager:
    """Single responsibility: Manage variant column schema operations."""

    def create_variant_columns(
        self,
        df: pl.DataFrame,
        variants: Dict[str, List[str]],
        provider: str
    ) -> pl.DataFrame:
        """
        Add all possible variant columns with null values.

        Args:
            df: DataFrame to add columns to
            variants: Dictionary mapping canonical keys to variant lists
            provider: Provider name for column naming

        Returns:
            DataFrame with all variant columns added as nulls
        """
        if df.height == 0:
            return df

        # Create expressions for all possible variant columns
        new_columns = []
        for canonical_key, variant_keys in variants.items():
            for variant_key in variant_keys:
                column_name = f"resource/tag:{provider}:{variant_key}"
                # Add column with null values of string type
                new_columns.append(pl.lit(None, dtype=pl.String).alias(column_name))

        if not new_columns:
            return df

        return df.with_columns(new_columns)

    def populate_selected_column(
        self,
        df: pl.DataFrame,
        column_name: str,
        values: List[str]
    ) -> pl.DataFrame:
        """
        Populate specific variant column with values.

        Args:
            df: DataFrame to populate
            column_name: Name of column to populate
            values: List of values to use for population (one per row)

        Returns:
            DataFrame with specified column populated

        Raises:
            ValueError: If values length doesn't match DataFrame height
        """
        if df.height == 0:
            return df

        if len(values) != df.height:
            raise ValueError(f"Values length ({len(values)}) must match DataFrame height ({df.height})")

        if column_name not in df.columns:
            raise ValueError(f"Column '{column_name}' not found in DataFrame")

        # Create series from values and update the column
        value_series = pl.Series(column_name, values)
        return df.with_columns([value_series])