#  Copyright (c) ZenML GmbH 2022. 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 steps runs."""

from datetime import datetime, timedelta
from typing import (
    TYPE_CHECKING,
    ClassVar,
    Dict,
    List,
    Optional,
    Type,
    TypeVar,
    Union,
)
from uuid import UUID

from pydantic import BaseModel, ConfigDict, Field

from zenml.config.step_configurations import Step, StepConfiguration, StepSpec
from zenml.constants import STR_FIELD_MAX_LENGTH, TEXT_FIELD_MAX_LENGTH
from zenml.enums import (
    ArtifactSaveType,
    ExecutionStatus,
    StepRunInputArtifactType,
    StepType,
)
from zenml.metadata.metadata_types import MetadataType
from zenml.models.v2.base.base import BaseUpdate
from zenml.models.v2.base.scoped import (
    ProjectScopedFilter,
    ProjectScopedRequest,
    ProjectScopedResponse,
    ProjectScopedResponseBody,
    ProjectScopedResponseMetadata,
    ProjectScopedResponseResources,
    RunMetadataFilterMixin,
)
from zenml.models.v2.core.artifact_version import ArtifactVersionResponse
from zenml.models.v2.core.model_version import ModelVersionResponse
from zenml.models.v2.misc.exception_info import ExceptionInfo
from zenml.utils.time_utils import utc_now

if TYPE_CHECKING:
    from sqlalchemy.sql.elements import ColumnElement

    from zenml.models.v2.core.logs import (
        LogsRequest,
        LogsResponse,
    )
    from zenml.zen_stores.schemas import BaseSchema

    AnySchema = TypeVar("AnySchema", bound=BaseSchema)


class StepRunInputResponse(ArtifactVersionResponse):
    """Response model for step run inputs."""

    input_type: StepRunInputArtifactType
    index: Optional[int] = Field(
        title="The index of the input artifact in the step run.",
        default=None,
    )
    chunk_index: Optional[int] = Field(
        title="The index of the chunk in the input artifact.",
        default=None,
    )
    chunk_size: Optional[int] = Field(
        title="The size of the chunk in the input artifact.",
        default=None,
    )

    def get_hydrated_version(self) -> "StepRunInputResponse":
        """Get the hydrated version of this step run input.

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

        return StepRunInputResponse(
            input_type=self.input_type,
            **Client().zen_store.get_artifact_version(self.id).model_dump(),
        )


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


class StepRunRequest(ProjectScopedRequest):
    """Request model for step runs."""

    name: str = Field(
        title="The name of the pipeline run step.",
        max_length=STR_FIELD_MAX_LENGTH,
    )
    start_time: datetime = Field(
        title="The start time of the step run.",
    )
    end_time: Optional[datetime] = Field(
        title="The end time of the step run.",
        default=None,
    )
    status: ExecutionStatus = Field(title="The status of the step.")
    cache_key: Optional[str] = Field(
        title="The cache key of the step run.",
        default=None,
        max_length=STR_FIELD_MAX_LENGTH,
    )
    cache_expires_at: Optional[datetime] = Field(
        title="The time at which this step run should not be used for cached "
        "results anymore. If not set, the result will never expire.",
        default=None,
    )
    code_hash: Optional[str] = Field(
        title="The code hash of the step run.",
        default=None,
        max_length=STR_FIELD_MAX_LENGTH,
    )
    docstring: Optional[str] = Field(
        title="The docstring of the step function or class.",
        default=None,
        max_length=TEXT_FIELD_MAX_LENGTH,
    )
    source_code: Optional[str] = Field(
        title="The source code of the step function or class.",
        default=None,
        max_length=TEXT_FIELD_MAX_LENGTH,
    )
    pipeline_run_id: UUID = Field(
        title="The ID of the pipeline run that this step run belongs to.",
    )
    original_step_run_id: Optional[UUID] = Field(
        title="The ID of the original step run if this step was cached.",
        default=None,
    )
    parent_step_ids: List[UUID] = Field(
        title="The IDs of the parent steps of this step run.",
        default_factory=list,
        deprecated=True,
    )
    inputs: Dict[str, List[UUID]] = Field(
        title="The IDs of the input artifact versions of the step run.",
        default_factory=dict,
    )
    outputs: Dict[str, List[UUID]] = Field(
        title="The IDs of the output artifact versions of the step run.",
        default_factory=dict,
    )
    logs: Optional[Union[UUID, "LogsRequest"]] = Field(
        title="Logs associated with this step run.",
        default=None,
    )
    exception_info: Optional[ExceptionInfo] = Field(
        default=None,
        title="The exception information of the step run.",
    )
    dynamic_config: Optional["Step"] = Field(
        title="The dynamic configuration of the step run.",
        default=None,
    )

    model_config = ConfigDict(protected_namespaces=())


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


class StepRunUpdate(BaseUpdate):
    """Update model for step runs."""

    outputs: Dict[str, List[UUID]] = Field(
        title="The IDs of the output artifact versions of the step run.",
        default={},
    )
    loaded_artifact_versions: Dict[str, UUID] = Field(
        title="The IDs of artifact versions that were loaded by this step run.",
        default={},
    )
    status: Optional[ExecutionStatus] = Field(
        title="The status of the step.",
        default=None,
    )
    end_time: Optional[datetime] = Field(
        title="The end time of the step run.",
        default=None,
    )
    exception_info: Optional[ExceptionInfo] = Field(
        default=None,
        title="The exception information of the step run.",
    )
    cache_expires_at: Optional[datetime] = Field(
        title="The time at which this step run should not be used for cached "
        "results anymore.",
        default=None,
    )
    add_logs: Optional[List["LogsRequest"]] = Field(
        default=None, title="New logs to add to the step run."
    )
    model_config = ConfigDict(protected_namespaces=())


# ------------------ Response Model ------------------
class StepRunResponseBody(ProjectScopedResponseBody):
    """Response body for step runs."""

    type: Optional[StepType] = Field(
        title="The type of the step.",
        default=None,
    )
    status: ExecutionStatus = Field(title="The status of the step.")
    version: int = Field(
        title="The version of the step run.",
    )
    is_retriable: bool = Field(
        title="Whether the step run is retriable.",
    )
    start_time: Optional[datetime] = Field(
        title="The start time of the step run.",
        default=None,
    )
    end_time: Optional[datetime] = Field(
        title="The end time of the step run.",
        default=None,
    )
    latest_heartbeat: Optional[datetime] = Field(
        title="The latest heartbeat of the step run.",
        default=None,
    )
    model_version_id: Optional[UUID] = Field(
        title="The ID of the model version that was "
        "configured by this step run explicitly.",
        default=None,
    )
    substitutions: Dict[str, str] = Field(
        title="The substitutions of the step run.",
        default={},
    )
    heartbeat_threshold: Optional[int] = Field(
        title="The applied heartbeat healthiness threshold ",
        default=None,
    )
    model_config = ConfigDict(protected_namespaces=())


class StepRunResponseMetadata(ProjectScopedResponseMetadata):
    """Response metadata for step runs."""

    __zenml_skip_dehydration__: ClassVar[List[str]] = [
        "config",
        "spec",
        "metadata",
    ]

    # Configuration
    config: "StepConfiguration" = Field(title="The configuration of the step.")
    spec: "StepSpec" = Field(title="The spec of the step.")

    # Code related fields
    cache_key: Optional[str] = Field(
        title="The cache key of the step run.",
        default=None,
        max_length=STR_FIELD_MAX_LENGTH,
    )
    cache_expires_at: Optional[datetime] = Field(
        title="The time at which this step run should not be used for cached "
        "results anymore. If not set, the result will never expire.",
        default=None,
    )
    code_hash: Optional[str] = Field(
        title="The code hash of the step run.",
        default=None,
        max_length=STR_FIELD_MAX_LENGTH,
    )
    docstring: Optional[str] = Field(
        title="The docstring of the step function or class.",
        default=None,
        max_length=TEXT_FIELD_MAX_LENGTH,
    )
    source_code: Optional[str] = Field(
        title="The source code of the step function or class.",
        default=None,
        max_length=TEXT_FIELD_MAX_LENGTH,
    )
    exception_info: Optional[ExceptionInfo] = Field(
        default=None,
        title="The exception information of the step run.",
    )

    # References
    snapshot_id: UUID = Field(
        title="The snapshot associated with the step run."
    )
    pipeline_run_id: UUID = Field(
        title="The ID of the pipeline run that this step run belongs to.",
    )
    original_step_run_id: Optional[UUID] = Field(
        title="The ID of the original step run if this step was cached.",
        default=None,
    )
    parent_step_ids: List[UUID] = Field(
        title="The IDs of the parent steps of this step run.",
        default_factory=list,
    )
    run_metadata: Dict[str, MetadataType] = Field(
        title="Metadata associated with this step run.",
        default={},
    )


class StepRunResponseResources(ProjectScopedResponseResources):
    """Class for all resource models associated with the step run entity."""

    log_collection: Optional[List["LogsResponse"]] = Field(
        title="Logs associated with this step run.",
        default=None,
    )

    model_version: Optional[ModelVersionResponse] = None
    inputs: Dict[str, List[StepRunInputResponse]] = Field(
        title="The input artifact versions of the step run.",
        default_factory=dict,
    )
    outputs: Dict[str, List[ArtifactVersionResponse]] = Field(
        title="The output artifact versions of the step run.",
        default_factory=dict,
    )

    # 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=())


class StepRunResponse(
    ProjectScopedResponse[
        StepRunResponseBody, StepRunResponseMetadata, StepRunResponseResources
    ]
):
    """Response model for step runs."""

    name: str = Field(
        title="The name of the pipeline run step.",
        max_length=STR_FIELD_MAX_LENGTH,
    )

    def get_hydrated_version(self) -> "StepRunResponse":
        """Get the hydrated version of this step run.

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

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

    # Helper properties
    @property
    def input(self) -> StepRunInputResponse:
        """Returns the input artifact that was used to run this step.

        Returns:
            The input artifact.

        Raises:
            ValueError: If there were zero or multiple inputs to this step.
        """
        if not self.inputs:
            raise ValueError(f"Step {self.name} has no inputs.")
        if len(self.inputs) > 1 or (
            len(self.inputs) == 1 and len(next(iter(self.inputs.values()))) > 1
        ):
            raise ValueError(
                f"Step {self.name} has multiple inputs, so `Step.input` is "
                "ambiguous. Please use `Step.inputs` instead."
            )
        return next(iter(self.inputs.values()))[0]

    @property
    def output(self) -> ArtifactVersionResponse:
        """Returns the output artifact that was written by this step.

        Returns:
            The output artifact.

        Raises:
            ValueError: If there were zero or multiple step outputs.
        """
        if not self.outputs:
            raise ValueError(f"Step {self.name} has no outputs.")
        if len(self.outputs) > 1 or (
            len(self.outputs) == 1
            and len(next(iter(self.outputs.values()))) > 1
        ):
            raise ValueError(
                f"Step {self.name} has multiple outputs, so `Step.output` is "
                "ambiguous. Please use `Step.outputs` instead."
            )
        return next(iter(self.outputs.values()))[0]

    @property
    def regular_inputs(self) -> Dict[str, List[StepRunInputResponse]]:
        """Returns the regular step inputs of the step run.

        Regular step inputs are the inputs that are defined in the step function
        signature, and are not manually loaded during the step execution.

        Returns:
            The regular step inputs.
        """
        result = {}

        for input_name, input_artifacts in self.inputs.items():
            filtered = [
                input_artifact
                for input_artifact in input_artifacts
                if input_artifact.input_type != StepRunInputArtifactType.MANUAL
            ]
            if filtered:
                result[input_name] = filtered

        return result

    @property
    def regular_outputs(self) -> Dict[str, ArtifactVersionResponse]:
        """Returns the regular step outputs of the step run.

        Regular step outputs are the outputs that are defined in the step
        function signature, and are not manually saved during the step
        execution.

        Raises:
            ValueError: If there were multiple regular output artifacts for the
                same output name.

        Returns:
            The regular step outputs.
        """
        result = {}

        for output_name, output_artifacts in self.outputs.items():
            filtered = [
                output_artifact
                for output_artifact in output_artifacts
                if output_artifact.save_type == ArtifactSaveType.STEP_OUTPUT
            ]
            if len(filtered) > 1:
                raise ValueError(
                    f"Expected 1 regular output artifact for {output_name}, "
                    f"got {len(filtered)}."
                )
            if filtered:
                result[output_name] = filtered[0]

        return result

    # Body and metadata properties

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

        Returns:
            the value of the property.
        """
        return self.get_body().type

    @property
    def status(self) -> ExecutionStatus:
        """The `status` property.

        Returns:
            the value of the property.
        """
        return self.get_body().status

    @property
    def version(self) -> int:
        """The `version` property.

        Returns:
            the value of the property.
        """
        return self.get_body().version

    @property
    def is_retriable(self) -> bool:
        """The `is_retriable` property.

        Returns:
            the value of the property.
        """
        return self.get_body().is_retriable

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

        Returns:
            the value of the property.
        """
        return self.get_resources().inputs

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

        Returns:
            the value of the property.
        """
        return self.get_resources().outputs

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

        Returns:
            the value of the property.
        """
        return self.get_body().model_version_id

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

        Returns:
            the value of the property.
        """
        return self.get_body().substitutions

    @property
    def config(self) -> "StepConfiguration":
        """The `config` property.

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

    @property
    def spec(self) -> "StepSpec":
        """The `spec` property.

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

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

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

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

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

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

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

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

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

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

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

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

        Returns:
            the value of the property.
        """
        return self.get_body().start_time

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

        Returns:
            the value of the property.
        """
        return self.get_body().end_time

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

        Returns:
            The duration of the step run.
        """
        if self.end_time is None or self.start_time is None:
            return None
        return self.end_time - self.start_time

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

        Returns:
            the value of the property.
        """
        return self.get_body().latest_heartbeat

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

        Returns:
            the value of the property.
        """
        return self.get_body().heartbeat_threshold

    @property
    def snapshot_id(self) -> UUID:
        """The `snapshot_id` property.

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

    @property
    def pipeline_run_id(self) -> UUID:
        """The `pipeline_run_id` property.

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

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

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

    @property
    def parent_step_ids(self) -> List[UUID]:
        """The `parent_step_ids` property.

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

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

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

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

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

    @property
    def log_collection(self) -> Optional[List["LogsResponse"]]:
        """The `log_collection` property.

        Returns:
            the value of the property.
        """
        return self.get_resources().log_collection

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

        Returns:
            the value of the property.
        """
        return self.get_resources().model_version


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


class StepRunFilter(ProjectScopedFilter, RunMetadataFilterMixin):
    """Model to enable advanced filtering of step runs."""

    FILTER_EXCLUDE_FIELDS: ClassVar[List[str]] = [
        *ProjectScopedFilter.FILTER_EXCLUDE_FIELDS,
        *RunMetadataFilterMixin.FILTER_EXCLUDE_FIELDS,
        "model",
        "exclude_retried",
        "cache_expired",
    ]
    CLI_EXCLUDE_FIELDS: ClassVar[List[str]] = [
        *ProjectScopedFilter.CLI_EXCLUDE_FIELDS,
        *RunMetadataFilterMixin.CLI_EXCLUDE_FIELDS,
    ]
    CUSTOM_SORTING_OPTIONS: ClassVar[List[str]] = [
        *ProjectScopedFilter.CUSTOM_SORTING_OPTIONS,
        *RunMetadataFilterMixin.CUSTOM_SORTING_OPTIONS,
    ]
    API_MULTI_INPUT_PARAMS: ClassVar[List[str]] = [
        *ProjectScopedFilter.API_MULTI_INPUT_PARAMS,
        *RunMetadataFilterMixin.API_MULTI_INPUT_PARAMS,
    ]

    name: Optional[str] = Field(
        default=None,
        description="Name of the step run",
    )
    code_hash: Optional[str] = Field(
        default=None,
        description="Code hash for this step run",
    )
    cache_key: Optional[str] = Field(
        default=None,
        description="Cache key for this step run",
    )
    status: Optional[str] = Field(
        default=None,
        description="Status of the Step Run",
    )
    start_time: Optional[Union[datetime, str]] = Field(
        default=None,
        description="Start time for this run",
        union_mode="left_to_right",
    )
    end_time: Optional[Union[datetime, str]] = Field(
        default=None,
        description="End time for this run",
        union_mode="left_to_right",
    )
    pipeline_run_id: Optional[Union[UUID, str]] = Field(
        default=None,
        description="Pipeline run of this step run",
        union_mode="left_to_right",
    )
    snapshot_id: Optional[Union[UUID, str]] = Field(
        default=None,
        description="Snapshot of this step run",
        union_mode="left_to_right",
    )
    original_step_run_id: Optional[Union[UUID, str]] = Field(
        default=None,
        description="Original id for this step run",
        union_mode="left_to_right",
    )
    model_version_id: Optional[Union[UUID, str]] = Field(
        default=None,
        description="Model version associated with the step run.",
        union_mode="left_to_right",
    )
    model: Optional[Union[UUID, str]] = Field(
        default=None,
        description="Name/ID of the model associated with the step run.",
    )
    version: Union[int, str, None] = Field(
        default=None,
        description="Version of the step run.",
        union_mode="left_to_right",
    )
    exclude_retried: Optional[bool] = Field(
        default=None,
        description="Whether to exclude retried step runs.",
    )
    cache_expires_at: Optional[Union[datetime, str]] = Field(
        default=None,
        description="Cache expiration time of the step run.",
        union_mode="left_to_right",
    )
    cache_expired: Optional[bool] = Field(
        default=None,
        description="Whether the cache expiration time of the step run has "
        "passed.",
    )
    model_config = ConfigDict(protected_namespaces=())

    def get_custom_filters(
        self, table: Type["AnySchema"]
    ) -> List["ColumnElement[bool]"]:
        """Get custom filters.

        Args:
            table: The query table.

        Returns:
            A list of custom filters.
        """
        custom_filters = super().get_custom_filters(table)

        from sqlmodel import and_, col, or_

        from zenml.zen_stores.schemas import (
            ModelSchema,
            ModelVersionSchema,
            StepRunSchema,
        )

        if self.model:
            model_filter = and_(
                StepRunSchema.model_version_id == ModelVersionSchema.id,
                ModelVersionSchema.model_id == ModelSchema.id,
                self.generate_name_or_id_query_conditions(
                    value=self.model, table=ModelSchema
                ),
            )
            custom_filters.append(model_filter)

        if self.exclude_retried:
            custom_filters.append(
                col(StepRunSchema.status) != ExecutionStatus.RETRIED.value
            )

        if self.cache_expired is True:
            cache_expiration_filter = and_(
                col(StepRunSchema.cache_expires_at).is_not(None),
                col(StepRunSchema.cache_expires_at) < utc_now(),
            )
            custom_filters.append(cache_expiration_filter)
        elif self.cache_expired is False:
            cache_expiration_filter = or_(
                col(StepRunSchema.cache_expires_at) > utc_now(),
                col(StepRunSchema.cache_expires_at).is_(None),
            )
            custom_filters.append(cache_expiration_filter)

        return custom_filters


# ------------------ Heartbeat Model ---------------


class StepHeartbeatResponse(BaseModel, use_enum_values=True):
    """Light-weight model for Step Heartbeat responses."""

    id: UUID
    status: ExecutionStatus
    latest_heartbeat: datetime
    pipeline_run_status: ExecutionStatus | None = None
    heartbeat_enabled: bool
