# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

from __future__ import annotations

import json
from enum import Enum
from typing import Any, Literal

from pydantic import Field, field_validator

from data_designer.config.base import ProcessorConfig
from data_designer.config.errors import InvalidConfigError


class ProcessorType(str, Enum):
    """Enumeration of available processor types.

    Attributes:
        DROP_COLUMNS: Processor that removes specified columns from the output dataset.
        SCHEMA_TRANSFORM: Processor that creates a new dataset with a transformed schema using Jinja2 templates.
    """

    DROP_COLUMNS = "drop_columns"
    SCHEMA_TRANSFORM = "schema_transform"


def get_processor_config_from_kwargs(processor_type: ProcessorType, **kwargs: Any) -> ProcessorConfig:
    """Create a processor configuration from a processor type and keyword arguments.

    Args:
        processor_type: The type of processor to create.
        **kwargs: Additional keyword arguments passed to the processor constructor.

    Returns:
        A processor configuration object of the specified type.
    """
    if processor_type == ProcessorType.DROP_COLUMNS:
        return DropColumnsProcessorConfig(**kwargs)
    elif processor_type == ProcessorType.SCHEMA_TRANSFORM:
        return SchemaTransformProcessorConfig(**kwargs)


class DropColumnsProcessorConfig(ProcessorConfig):
    """Drop columns from the output dataset (prefer ``drop=True`` in the column config).

    This processor removes specified columns from the generated dataset. The dropped
    columns are saved separately in a `dropped-columns` directory for reference.
    When this processor is added via the config builder, the corresponding column
    configs are automatically marked with `drop = True`.

    Attributes:
        column_names (required): List of column names to remove from the output dataset.

    Inherited Attributes:
        name (required): Name of the processor.
    """

    column_names: list[str] = Field(description="List of column names to drop from the output dataset.")
    processor_type: Literal[ProcessorType.DROP_COLUMNS] = ProcessorType.DROP_COLUMNS


class SchemaTransformProcessorConfig(ProcessorConfig):
    """Configuration for transforming the dataset schema using Jinja2 templates.

    This processor creates a new dataset with a transformed schema. Each key in the
    template becomes a column in the output, and values are Jinja2 templates that
    can reference any column in the batch. The transformed dataset is written to
    a `processors-outputs/{processor_name}/` directory alongside the main dataset.

    Attributes:
        template (required): Dictionary defining the output schema. Keys are new column names,
            values are Jinja2 templates (strings, lists, or nested structures).
            Must be JSON-serializable.

    Inherited Attributes:
        name (required): Name of the processor.
    """

    template: dict[str, Any] = Field(
        ...,
        description="""
        Dictionary specifying columns and templates to use in the new dataset with transformed schema.

        Each key is a new column name, and each value is an object containing Jinja2 templates - for instance, a string or a list of strings.
        Values must be JSON-serializable.

        Example:

        ```python
        template = {
            "list_of_strings": ["{{ col1 }}", "{{ col2 }}"],
            "uppercase_string": "{{ col1 | upper }}",
            "lowercase_string": "{{ col2 | lower }}",
        }
        ```

        The above templates will create an new dataset with three columns: "list_of_strings", "uppercase_string", and "lowercase_string".
        References to columns "col1" and "col2" in the templates will be replaced with the actual values of the columns in the dataset.
        """,
    )
    processor_type: Literal[ProcessorType.SCHEMA_TRANSFORM] = ProcessorType.SCHEMA_TRANSFORM

    @field_validator("template")
    def validate_template(cls, v: dict[str, Any]) -> dict[str, Any]:
        try:
            json.dumps(v)
        except TypeError as e:
            if "not JSON serializable" in str(e):
                raise InvalidConfigError("Template must be JSON serializable")
        return v
