#  Copyright (c) ZenML GmbH 2023. All Rights Reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at:
#
#       https://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
#  or implied. See the License for the specific language governing
#  permissions and limitations under the License.
"""Util functions for step and pipeline entrypoint functions."""

import inspect
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    Dict,
    NamedTuple,
    NoReturn,
    Optional,
    Sequence,
    Type,
    Union,
)

from pydantic import ConfigDict, ValidationError, create_model

from zenml.constants import ENFORCE_TYPE_ANNOTATIONS
from zenml.exceptions import StepInterfaceError
from zenml.logger import get_logger
from zenml.materializers.base_materializer import BaseMaterializer
from zenml.metadata.lazy_load import LazyRunMetadataResponse
from zenml.steps.utils import (
    OutputSignature,
    parse_return_type_annotations,
    resolve_type_annotation,
)
from zenml.utils import yaml_utils

if TYPE_CHECKING:
    from zenml.config.source import Source
    from zenml.pipelines.pipeline_definition import Pipeline

    MaterializerClassOrSource = Union[str, "Source", Type["BaseMaterializer"]]

logger = get_logger(__name__)


class StepArtifact:
    """Class to represent step output artifacts."""

    def __init__(
        self,
        invocation_id: str,
        output_name: str,
        annotation: OutputSignature,
        pipeline: "Pipeline",
        chunk_index: Optional[int] = None,
        chunk_size: Optional[int] = None,
    ) -> None:
        """Initialize a step artifact.

        Args:
            invocation_id: The ID of the invocation that produces this artifact.
            output_name: The name of the output that produces this artifact.
            annotation: The output type annotation.
            pipeline: The pipeline which the invocation is part of.
            chunk_index: The index of the chunk to load.
            chunk_size: The size of the chunk to load.
        """
        self.invocation_id = invocation_id
        self.output_name = output_name
        self.annotation = annotation
        self.pipeline = pipeline
        self.chunk_index = chunk_index
        self.chunk_size = chunk_size

    def __iter__(self) -> NoReturn:
        """Raise a custom error if someone is trying to iterate this object.

        Raises:
            StepInterfaceError: If trying to iterate this object.
        """
        raise StepInterfaceError(
            "Unable to unpack step artifact. This error is probably because "
            "you're trying to unpack the return value of your step but the "
            "step only returns a single artifact. For more information on how "
            "to add type annotations to your step to indicate multiple "
            "artifacts visit https://docs.zenml.io/concepts/steps_and_pipelines#multiple-return-values."
        )


def validate_reserved_arguments(
    signature: inspect.Signature, reserved_arguments: Sequence[str]
) -> None:
    """Validates that the signature does not contain any reserved arguments.

    Args:
        signature: The signature to validate.
        reserved_arguments: The reserved arguments for the signature.

    Raises:
        RuntimeError: If the signature contains a reserved argument.
    """
    for arg in reserved_arguments:
        if arg in signature.parameters:
            raise RuntimeError(f"Reserved argument name '{arg}'.")


class EntrypointFunctionDefinition(NamedTuple):
    """Class representing a step entrypoint function.

    Attributes:
        inputs: The entrypoint function inputs.
        outputs: The entrypoint function outputs. This dictionary maps output
            names to output annotations.
    """

    inputs: Dict[str, inspect.Parameter]
    outputs: Dict[str, OutputSignature]

    def validate_input(self, key: str, value: Any) -> None:
        """Validates an input to the step entrypoint function.

        Args:
            key: The key for which the input was passed
            value: The input value.

        Raises:
            KeyError: If the function has no input for the given key.
            RuntimeError: If a parameter is passed for an input that is
                annotated as an `UnmaterializedArtifact`.
            RuntimeError: If the input value is not valid for the type
                annotation provided for the function parameter.
            StepInterfaceError: If the input is a parameter and not JSON
                serializable.
        """
        from zenml.artifacts.external_artifact import ExternalArtifact
        from zenml.artifacts.unmaterialized_artifact import (
            UnmaterializedArtifact,
        )
        from zenml.client_lazy_loader import ClientLazyLoader
        from zenml.models import ArtifactVersionResponse

        if key not in self.inputs:
            raise KeyError(
                f"Received step entrypoint input for invalid key {key}."
            )

        parameter = self.inputs[key]

        if isinstance(
            value,
            (
                StepArtifact,
                ExternalArtifact,
                ArtifactVersionResponse,
                ClientLazyLoader,
                LazyRunMetadataResponse,
            ),
        ):
            # If we were to do any type validation for artifacts here, we
            # would not be able to leverage pydantics type coercion (e.g.
            # providing an `int` artifact for a `float` input)
            return
        elif isinstance(value, list) and all(
            isinstance(item, StepArtifact) for item in value
        ):
            # Same as above, we don't do any type validation for a list of
            # artifacts here.
            return

        # Not an artifact -> This is a parameter
        if parameter.annotation is UnmaterializedArtifact:
            raise RuntimeError(
                "Passing parameter for input of type `UnmaterializedArtifact` "
                "is not allowed."
            )

        if not yaml_utils.is_json_serializable(value):
            raise StepInterfaceError(
                f"Argument type (`{type(value)}`) for argument "
                f"'{key}' is not JSON serializable and can not be passed as "
                "a parameter. This input can either be provided by the "
                "output of another step or as an external artifact: "
                "https://docs.zenml.io/user-guides/starter-guide/manage-artifacts#managing-artifacts-not-produced-by-zenml-pipelines"
            )

        try:
            self._validate_input_value(parameter=parameter, value=value)
        except ValidationError as e:
            raise RuntimeError(
                f"Input validation failed for input '{parameter.name}': "
                f"Expected type `{parameter.annotation}` but received type "
                f"`{type(value)}`."
            ) from e

    def _validate_input_value(
        self, parameter: inspect.Parameter, value: Any
    ) -> None:
        """Validates an input value to the step entrypoint function.

        Args:
            parameter: The function parameter for which the value was provided.
            value: The input value.
        """
        # We allow passing None for optional annotations that would otherwise
        # not be allowed as a parameter
        config_dict = ConfigDict(arbitrary_types_allowed=value is None)

        # Create a pydantic model with just a single required field with the
        # type annotation of the parameter to verify the input type including
        # pydantics type coercion
        validation_model_class = create_model(
            "input_validation_model",
            __config__=config_dict,
            value=(parameter.annotation, ...),
        )
        validation_model_class(value=value)


def validate_entrypoint_function(
    func: Callable[..., Any],
    reserved_arguments: Sequence[str] = (),
) -> EntrypointFunctionDefinition:
    """Validates a step entrypoint function.

    Args:
        func: The step entrypoint function to validate.
        reserved_arguments: The reserved arguments for the entrypoint function.

    Raises:
        StepInterfaceError: If the entrypoint function has variable arguments
            or keyword arguments.
        RuntimeError: If type annotations should be enforced and a type
            annotation is missing.

    Returns:
        A validated definition of the entrypoint function.
    """
    signature = inspect.signature(func, follow_wrapped=True)
    validate_reserved_arguments(
        signature=signature, reserved_arguments=reserved_arguments
    )

    inputs = {}

    signature_parameters = list(signature.parameters.items())
    for key, parameter in signature_parameters:
        if parameter.kind in {parameter.VAR_POSITIONAL, parameter.VAR_KEYWORD}:
            raise StepInterfaceError(
                f"Variable args or kwargs not allowed for function "
                f"{func.__name__}."
            )

        annotation = parameter.annotation
        if annotation is parameter.empty:
            if ENFORCE_TYPE_ANNOTATIONS:
                raise RuntimeError(
                    f"Missing type annotation for input '{key}' of step "
                    f"function '{func.__name__}'."
                )

            # If a type annotation is missing, use `Any` instead
            parameter = parameter.replace(annotation=Any)

        annotation = resolve_type_annotation(annotation)
        inputs[key] = parameter

    outputs = parse_return_type_annotations(
        func=func,
        enforce_type_annotations=ENFORCE_TYPE_ANNOTATIONS,
    )

    return EntrypointFunctionDefinition(
        inputs=inputs,
        outputs=outputs,
    )
