#  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.
"""Models representing pipeline deployments."""

from typing import Dict, Optional, Union
from uuid import UUID

from pydantic import Field

from zenml.config.docker_settings import SourceFileMode
from zenml.config.pipeline_configurations import PipelineConfiguration
from zenml.config.step_configurations import Step
from zenml.models.v2.base.base import BaseZenModel
from zenml.models.v2.base.scoped import (
    WorkspaceScopedFilter,
    WorkspaceScopedRequest,
    WorkspaceScopedResponse,
    WorkspaceScopedResponseBody,
    WorkspaceScopedResponseMetadata,
)
from zenml.models.v2.core.code_reference import (
    CodeReferenceRequest,
    CodeReferenceResponse,
)
from zenml.models.v2.core.pipeline import PipelineResponse
from zenml.models.v2.core.pipeline_build import (
    PipelineBuildResponse,
)
from zenml.models.v2.core.schedule import ScheduleResponse
from zenml.models.v2.core.stack import StackResponse

# ------------------ Request Model ------------------


class PipelineDeploymentBase(BaseZenModel):
    """Base model for pipeline deployments."""

    run_name_template: str = Field(
        title="The run name template for runs created using this deployment.",
    )
    pipeline_configuration: PipelineConfiguration = Field(
        title="The pipeline configuration for this deployment."
    )
    step_configurations: Dict[str, Step] = Field(
        default={}, title="The step configurations for this deployment."
    )
    client_environment: Dict[str, str] = Field(
        default={}, title="The client environment for this deployment."
    )
    client_version: Optional[str] = Field(
        default=None,
        title="The version of the ZenML installation on the client side.",
    )
    server_version: Optional[str] = Field(
        default=None,
        title="The version of the ZenML installation on the server side.",
    )

    @property
    def requires_included_files(self) -> bool:
        """Whether the deployment requires included files.

        Returns:
            Whether the deployment requires included files.
        """
        return any(
            step.config.docker_settings.source_files == SourceFileMode.INCLUDE
            for step in self.step_configurations.values()
        )

    @property
    def requires_code_download(self) -> bool:
        """Whether the deployment requires downloading some code files.

        Returns:
            Whether the deployment requires downloading some code files.
        """
        return any(
            step.config.docker_settings.source_files == SourceFileMode.DOWNLOAD
            for step in self.step_configurations.values()
        )


class PipelineDeploymentRequest(
    PipelineDeploymentBase, WorkspaceScopedRequest
):
    """Request model for pipeline deployments."""

    stack: UUID = Field(title="The stack associated with the deployment.")
    pipeline: Optional[UUID] = Field(
        default=None, title="The pipeline associated with the deployment."
    )
    build: Optional[UUID] = Field(
        default=None, title="The build associated with the deployment."
    )
    schedule: Optional[UUID] = Field(
        default=None, title="The schedule associated with the deployment."
    )
    code_reference: Optional["CodeReferenceRequest"] = Field(
        default=None,
        title="The code reference associated with the deployment.",
    )


# ------------------ Update Model ------------------

# There is no update model for pipeline deployments.

# ------------------ Response Model ------------------


class PipelineDeploymentResponseBody(WorkspaceScopedResponseBody):
    """Response body for pipeline deployments."""


class PipelineDeploymentResponseMetadata(WorkspaceScopedResponseMetadata):
    """Response metadata for pipeline deployments."""

    run_name_template: str = Field(
        title="The run name template for runs created using this deployment.",
    )
    pipeline_configuration: PipelineConfiguration = Field(
        title="The pipeline configuration for this deployment."
    )
    step_configurations: Dict[str, Step] = Field(
        default={}, title="The step configurations for this deployment."
    )
    client_environment: Dict[str, str] = Field(
        default={}, title="The client environment for this deployment."
    )
    client_version: Optional[str] = Field(
        title="The version of the ZenML installation on the client side."
    )
    server_version: Optional[str] = Field(
        title="The version of the ZenML installation on the server side."
    )
    pipeline: Optional[PipelineResponse] = Field(
        default=None, title="The pipeline associated with the deployment."
    )
    stack: Optional[StackResponse] = Field(
        default=None, title="The stack associated with the deployment."
    )
    build: Optional[PipelineBuildResponse] = Field(
        default=None,
        title="The pipeline build associated with the deployment.",
    )
    schedule: Optional[ScheduleResponse] = Field(
        default=None, title="The schedule associated with the deployment."
    )
    code_reference: Optional[CodeReferenceResponse] = Field(
        default=None,
        title="The code reference associated with the deployment.",
    )


class PipelineDeploymentResponse(
    WorkspaceScopedResponse[
        PipelineDeploymentResponseBody, PipelineDeploymentResponseMetadata
    ]
):
    """Response model for pipeline deployments."""

    def get_hydrated_version(self) -> "PipelineDeploymentResponse":
        """Return the hydrated version of this pipeline deployment.

        Returns:
            an instance of the same entity with the metadata field attached.
        """
        from zenml.client import Client

        return Client().zen_store.get_deployment(self.id)

    # Body and metadata properties
    @property
    def run_name_template(self) -> str:
        """The `run_name_template` property.

        Returns:
            the value of the property.
        """
        return self.get_metadata().run_name_template

    @property
    def pipeline_configuration(self) -> PipelineConfiguration:
        """The `pipeline_configuration` property.

        Returns:
            the value of the property.
        """
        return self.get_metadata().pipeline_configuration

    @property
    def step_configurations(self) -> Dict[str, Step]:
        """The `step_configurations` property.

        Returns:
            the value of the property.
        """
        return self.get_metadata().step_configurations

    @property
    def client_environment(self) -> Dict[str, str]:
        """The `client_environment` property.

        Returns:
            the value of the property.
        """
        return self.get_metadata().client_environment

    @property
    def client_version(self) -> Optional[str]:
        """The `client_version` property.

        Returns:
            the value of the property.
        """
        return self.get_metadata().client_version

    @property
    def server_version(self) -> Optional[str]:
        """The `server_version` property.

        Returns:
            the value of the property.
        """
        return self.get_metadata().server_version

    @property
    def pipeline(self) -> Optional[PipelineResponse]:
        """The `pipeline` property.

        Returns:
            the value of the property.
        """
        return self.get_metadata().pipeline

    @property
    def stack(self) -> Optional[StackResponse]:
        """The `stack` property.

        Returns:
            the value of the property.
        """
        return self.get_metadata().stack

    @property
    def build(self) -> Optional[PipelineBuildResponse]:
        """The `build` property.

        Returns:
            the value of the property.
        """
        return self.get_metadata().build

    @property
    def schedule(self) -> Optional[ScheduleResponse]:
        """The `schedule` property.

        Returns:
            the value of the property.
        """
        return self.get_metadata().schedule

    @property
    def code_reference(self) -> Optional[CodeReferenceResponse]:
        """The `code_reference` property.

        Returns:
            the value of the property.
        """
        return self.get_metadata().code_reference

    @property
    def requires_code_download(self) -> bool:
        """Whether the deployment requires downloading some code files.

        Returns:
            Whether the deployment requires downloading some code files.
        """
        return any(
            step.config.docker_settings.source_files == SourceFileMode.DOWNLOAD
            for step in self.step_configurations.values()
        )


# ------------------ Filter Model ------------------


class PipelineDeploymentFilter(WorkspaceScopedFilter):
    """Model to enable advanced filtering of all pipeline deployments."""

    workspace_id: Optional[Union[UUID, str]] = Field(
        default=None, description="Workspace for this deployment."
    )
    user_id: Optional[Union[UUID, str]] = Field(
        default=None, description="User that created this deployment."
    )
    pipeline_id: Optional[Union[UUID, str]] = Field(
        default=None, description="Pipeline associated with the deployment."
    )
    stack_id: Optional[Union[UUID, str]] = Field(
        default=None, description="Stack associated with the deployment."
    )
    build_id: Optional[Union[UUID, str]] = Field(
        default=None, description="Build associated with the deployment."
    )
    schedule_id: Optional[Union[UUID, str]] = Field(
        default=None, description="Schedule associated with the deployment."
    )
