#  Copyright (c) ZenML GmbH 2024. 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.
"""Model Version Data Lazy Loader definition."""

from typing import TYPE_CHECKING, Any, Dict, Optional

from pydantic import BaseModel, ConfigDict, model_validator

from zenml.pipelines.pipeline_context import get_pipeline_context
from zenml.utils.pydantic_utils import before_validator_handler

if TYPE_CHECKING:
    from zenml.models import ModelVersionResponse, PipelineRunResponse


class ModelVersionDataLazyLoader(BaseModel):
    """Model Version Data Lazy Loader helper class.

    It helps the inner codes to fetch proper artifact,
    model version metadata or artifact metadata from the
    model version during runtime time of the step.
    """

    model_name: str
    model_version: Optional[str] = None
    artifact_name: Optional[str] = None
    artifact_version: Optional[str] = None
    metadata_name: Optional[str] = None

    # TODO: In Pydantic v2, the `model_` is a protected namespaces for all
    #  fields defined under base models. If not handled, this raises a warning.
    #  It is possible to suppress this warning message with the following
    #  configuration, however the ultimate solution is to rename these fields.
    #  Even though they do not cause any problems right now, if we are not
    #  careful we might overwrite some fields protected by pydantic.
    model_config = ConfigDict(protected_namespaces=())

    @model_validator(mode="before")
    @classmethod
    @before_validator_handler
    def _root_validator(cls, data: Dict[str, Any]) -> Dict[str, Any]:
        """Validate all in one.

        Args:
            data: Dict of values.

        Returns:
            Dict of validated values.

        Raises:
            ValueError: If the model version id, but call is not internal.
        """
        if data.get("model_version", None) is None:
            try:
                context = get_pipeline_context()
                if (
                    not context.model
                    or context.model.name != data["model_name"]
                ):
                    raise ValueError(
                        "`version` must be set if you use the `Model` class "
                        "directly in the pipeline body, otherwise, you can use "
                        "`get_pipeline_context().model` to lazy load the current "
                        "Model Version from the pipeline context."
                    )
            except RuntimeError:
                pass
        data["suppress_class_validation_warnings"] = True
        return data

    def _get_model_response(
        self, pipeline_run: "PipelineRunResponse"
    ) -> "ModelVersionResponse":
        # if the version/number is None -> return the model in context
        if self.model_version is None:
            if mv := pipeline_run.model_version:
                if mv.model.name != self.model_name:
                    raise RuntimeError(
                        "Lazy loading of the model failed, since given name "
                        f"`{self.model_name}` does not match the model name "
                        f"in the pipeline context: `{mv.model.name}`."
                    )
                return mv
            else:
                raise RuntimeError(
                    "Lazy loading of the model failed, since the model version "
                    "is not set in the pipeline context."
                )

        # else return the model version by version
        else:
            from zenml.client import Client

            try:
                return Client().get_model_version(
                    model_name_or_id=self.model_name,
                    model_version_name_or_number_or_id=self.model_version,
                )
            except KeyError as e:
                raise RuntimeError(
                    "Lazy loading of the model version failed: "
                    f"no model `{self.model_name}` with version "
                    f"`{self.model_version}` could be found."
                ) from e
