"""This module defines the abstract base class for pipeline steps in a data processing pipeline.

The module provides a standard interface for implementing data processing steps as part of a larger pipeline,
via the `PipelineStep` abstract base class. Each `PipelineStep` subclass should have a specific purpose and
should be able to run independently or as part of the pipeline.
"""

from __future__ import annotations

from abc import ABC, abstractmethod
from typing import Any, Literal, TypeVar

from typing_extensions import TypedDict, get_type_hints

from mleko.pipeline.data_container import DataContainer
from mleko.utils.custom_logger import CustomLogger


logger = CustomLogger()
"""The logger for the module."""

T = TypeVar("T")
"""Type variable for generic type annotations."""

TypedDictType = TypeVar("T", bound=TypedDict)  # type: ignore
"""Type variable for TypedDict type annotations."""


def validate_typeddict(data: dict[str, Any], typeddict: type[TypedDictType]) -> None:
    """Validates that all keys defined in the TypedDict are present in the data dict.

    This function checks for the presence of all keys, including those marked as Optional.
    The presence of a key is mandatory, but the value can be None if the key is Optional.

    Args:
        data: The data dictionary to validate.
        typeddict: The TypedDict class against which to validate the data.

    Raises:
        ValueError: If the data is missing required keys or has extra keys not defined in the TypedDict.
    """
    type_hints = get_type_hints(typeddict)

    # Check for missing keys
    missing_keys = [key for key in type_hints if key not in data]
    if missing_keys:
        msg = f"Missing required key(s) in {typeddict.__name__}: {', '.join(missing_keys)}"
        logger.error(msg)
        raise ValueError(msg)

    # Check for extra keys
    extra_keys = [key for key in data if key not in type_hints]
    if extra_keys:
        msg = f"Extra key(s) in data not defined in {typeddict.__name__}: {', '.join(extra_keys)}"
        logger.error(msg)
        raise ValueError(msg)


class PipelineStep(ABC):
    """Base class for all pipeline steps, ensuring the standardized interface for performing data processing operations.

    Descendants of this class must implement the `execute` method, which carries out the data processing operation
    related to the step. Additionally, they must implement the `_get_input_model` and `_get_output_model` methods,
    which dynamically determine the TypedDict models to use for validation.
    """

    def __init__(
        self,
        inputs: dict[str, Any] | TypedDictType,
        outputs: dict[str, str] | TypedDictType,
        cache_group: str | None,
    ) -> None:
        """Initialize a new PipelineStep with the provided input and output keys.

        Args:
            action: The action to perform, such as "fit", "transform", or "fit_transform".
            inputs: Input data required for this step's processing operation.
            outputs: Output data generated by this step's processing operation.
            cache_group: The cache group to use.
        """
        self._inputs = inputs
        self._outputs = outputs
        self._cache_group = cache_group

        self._validate_inputs()
        self._validate_outputs()

    @abstractmethod
    def execute(self, data_container: DataContainer, force_recompute: bool) -> DataContainer:
        """Execute the data processing operation associated with this pipeline step.

        The `execute` method is the main entry point for the data processing operation associated with this step.
        It receives a `DataContainer` instance as input, containing the data to be processed by this step.
        The method should perform the processing operation and return the processed data as a `DataContainer` instance.

        Args:
            data_container: Input data for this step's processing operation.
            force_recompute: Whether to force the step to recompute its output, even if it already exists.

        Raises:
            NotImplementedError: Must be implemented by subclass.
        """
        raise NotImplementedError

    @abstractmethod
    def _get_input_model(self) -> type[TypedDictType]:  # type: ignore
        """Dynamically determine the TypedDict model to use for validation.

        Returns:
            The TypedDict model to use for validation.
        """
        pass

    @abstractmethod
    def _get_output_model(self) -> type[TypedDictType]:  # type: ignore
        """Dynamically determine the TypedDict model to use for validation.

        Returns:
            The TypedDict model to use for validation.
        """
        pass

    def _validate_and_get_input(
        self,
        input_object: str | T,
        expected_type: type[T],
        data_container: DataContainer,
        is_optional: bool = False,
    ) -> T:
        """Validate and get the input from the data container or as a direct value.

        If the input is a string, it is treated as a key to look up in the data container. If it is not a string,
        it is treated as the input value itself. The method validates the input type, and returns the input value
        if it is valid.

        Warning:
            This method does not handle subscripted types, such as `List[str]`. It only handles simple types. If you
            need to validate a subscripted type, you should do so manually after retrieving the input value from this
            method.

        Args:
            input_object: The input key or value to validate and retrieve.
            expected_type: The expected type of the input value.
            data_container: The data container containing the input data.
            is_optional: Whether the input is optional.

        Raises:
            ValueError: If the input is invalid or not found in the data container.

        Returns:
            The input value if it is valid, or None if it is optional and not found in the data container.
        """
        if isinstance(input_object, str):
            input_value: T = data_container.data.get(input_object, None)
        else:
            input_value: T = input_object

        if input_value is None and not is_optional and isinstance(input_object, str):
            msg = f"Input {input_object!r} is required but not found in data container."
            logger.error(msg)
            raise ValueError(msg)

        if input_value is None and not is_optional and not isinstance(input_object, str):
            msg = "Input is required but found None."
            logger.error(msg)
            raise ValueError(msg)

        if input_value is not None and not isinstance(input_value, expected_type):
            msg = f"Invalid data type: {type(input_value)}. Expected {expected_type}."
            logger.error(msg)
            raise ValueError(msg)

        return input_value

    def _validate_inputs(self) -> None:
        """Validates the step's inputs using TypeDict models.

        Raises:
            ValueError: If the inputs are not a dictionary.
        """
        if isinstance(self._inputs, dict):
            input_model = self._get_input_model()
            validate_typeddict(self._inputs, input_model)
        else:
            msg = f"Inputs must be a dictionary, got {type(self._inputs)}."
            logger.error(msg)
            raise ValueError(msg)

    def _validate_outputs(self) -> None:
        """Validates the step's outputs using TypeDict models.

        Raises:
            ValueError: If the outputs are not a dictionary.
        """
        if isinstance(self._outputs, dict):
            output_model = self._get_output_model()
            validate_typeddict(self._outputs, output_model)
        else:
            msg = f"Outputs must be a dictionary, got {type(self._outputs)}."
            logger.error(msg)
            raise ValueError(msg)


FitTransformAction = Literal["fit", "transform", "fit_transform"]
"""Type alias for the action to perform, either "fit", "transform", or "fit_transform"."""


class FitTransformPipelineStep(PipelineStep):
    """Base class for all pipeline steps, ensuring the standardized interface for performing data processing operations.

    Specialized for steps that perform actions "fit", "transform", or "fit_transform". Descendants of this class must
    implement the `execute` method, which carries out the data processing operation related to the step. Additionally,
    they must implement the `_get_input_model` and `_get_output_model` methods, which dynamically determine the
    TypedDict models to use for validation based on the action to be performed.
    """

    _action: FitTransformAction

    def __init__(
        self,
        action: FitTransformAction,
        inputs: dict[str, Any] | TypedDictType,
        outputs: dict[str, str] | TypedDictType,
        cache_group: str | None,
    ) -> None:
        """Initialize a new PipelineStep with the provided input and output keys.

        Args:
            action: The action to perform, either "fit", "transform", or "fit_transform".
            inputs: Input data required for this step's processing operation.
            outputs: Output data generated by this step's processing operation.
            cache_group: The cache group to use.

        Raises:
            ValueError: If the action is not one of "fit", "transform", or "fit_transform".
        """
        if action not in ("fit", "transform", "fit_transform"):
            raise ValueError(f"Invalid action: {action}. Expected one of 'fit', 'transform', or 'fit_transform'.")

        self._action = action
        super().__init__(inputs, outputs, cache_group)
