"""
Transformation Registry.

Central registry for all transformation operations.
"""

from typing import Any, Callable

import pandas as pd
import polars as pl

from app.core.exceptions import InvalidTransformationError, TransformationExecutionError
from app.core.logging import get_logger

logger = get_logger(__name__)


class TransformationRegistry:
    """Registry for transformation operations."""

    def __init__(self):
        """Initialize registry."""
        self.transformations: dict[str, dict[str, Callable]] = {}
        self._register_transformations()

    def _register_transformations(self) -> None:
        """Register all transformations."""
        from app.transformations import cleaning, encoding, normalization, feature_engineering

        # Cleaning operations
        self.register("drop_duplicates", cleaning.drop_duplicates_pandas, "pandas")
        self.register("drop_duplicates", cleaning.drop_duplicates_polars, "polars")
        self.register("drop_na", cleaning.drop_na_pandas, "pandas")
        self.register("drop_na", cleaning.drop_na_polars, "polars")
        self.register("fill_missing", cleaning.fill_missing_pandas, "pandas")
        self.register("fill_missing", cleaning.fill_missing_polars, "polars")

        # Encoding operations
        self.register("encode_onehot", encoding.onehot_encode_pandas, "pandas")
        self.register("encode_onehot", encoding.onehot_encode_polars, "polars")
        self.register("encode_label", encoding.label_encode_pandas, "pandas")
        self.register("encode_label", encoding.label_encode_polars, "polars")

        # Normalization operations
        self.register("normalize", normalization.normalize_pandas, "pandas")
        self.register("normalize", normalization.normalize_polars, "polars")
        self.register("standardize", normalization.standardize_pandas, "pandas")
        self.register("standardize", normalization.standardize_polars, "polars")

        # Feature engineering
        self.register("create_feature", feature_engineering.create_feature_pandas, "pandas")
        self.register("binning", feature_engineering.binning_pandas, "pandas")

        # Add more transformations...

    def register(self, operation: str, func: Callable, engine: str) -> None:
        """Register a transformation function."""
        if operation not in self.transformations:
            self.transformations[operation] = {}
        self.transformations[operation][engine] = func
        logger.debug(f"Registered transformation: {operation} for engine: {engine}")

    def execute(
        self,
        operation: str,
        df: pd.DataFrame | pl.DataFrame,
        params: dict[str, Any],
        engine: str,
    ) -> pd.DataFrame | pl.DataFrame:
        """
        Execute a transformation.

        Args:
            operation: Operation name
            df: DataFrame to transform
            params: Operation parameters
            engine: Engine to use ('pandas' or 'polars')

        Returns:
            Transformed DataFrame

        Raises:
            InvalidTransformationError: If operation not found
            TransformationExecutionError: If execution fails
        """
        if operation not in self.transformations:
            raise InvalidTransformationError(
                f"Unknown transformation: {operation}",
                {"available": list(self.transformations.keys())},
            )

        if engine not in self.transformations[operation]:
            # Try fallback to other engine
            available_engines = list(self.transformations[operation].keys())
            if available_engines:
                logger.warning(
                    f"Engine {engine} not available for {operation}, using {available_engines[0]}"
                )
                # Convert DataFrame if needed
                if available_engines[0] == "pandas" and isinstance(df, pl.DataFrame):
                    df = df.to_pandas()
                elif available_engines[0] == "polars" and isinstance(df, pd.DataFrame):
                    df = pl.from_pandas(df)
                engine = available_engines[0]
            else:
                raise InvalidTransformationError(
                    f"No implementation available for {operation}",
                )

        func = self.transformations[operation][engine]

        try:
            result = func(df, **params)
            logger.info(f"Executed transformation: {operation}", engine=engine)
            return result
        except Exception as e:
            logger.error(f"Transformation failed: {operation}", error=str(e), exc_info=True)
            raise TransformationExecutionError(
                f"Failed to execute {operation}: {e}",
                {"params": params},
            )
