#  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.
"""Step run info."""

from typing import Any, Callable
from uuid import UUID

from zenml.config.frozen_base_model import FrozenBaseModel
from zenml.config.pipeline_configurations import PipelineConfiguration
from zenml.config.step_configurations import StepConfiguration, StepSpec
from zenml.logger import get_logger
from zenml.models import PipelineSnapshotResponse, StepRunResponse

logger = get_logger(__name__)


class StepRunInfo(FrozenBaseModel):
    """All information necessary to run a step."""

    step_run_id: UUID
    run_id: UUID
    run_name: str
    pipeline_step_name: str

    config: StepConfiguration
    spec: StepSpec
    pipeline: PipelineConfiguration
    snapshot: PipelineSnapshotResponse
    step_run: StepRunResponse

    force_write_logs: Callable[..., Any]

    def get_image(self, key: str) -> str:
        """Gets the Docker image for the given key.

        Args:
            key: The key for which to get the image.

        Raises:
            RuntimeError: If the run does not have an associated build.

        Returns:
            The image name or digest.
        """
        if not self.snapshot.build:
            raise RuntimeError(
                f"Missing build for snapshot {self.snapshot.id}. This is "
                "probably because the build was manually deleted."
            )

        if self.snapshot.is_dynamic:
            if self.config.docker_settings.skip_build:
                # If there are custom Docker settings for the step and no
                # build is required, we can use the parent image directly.
                assert self.config.docker_settings.parent_image
                return self.config.docker_settings.parent_image

            step_key = self.config.template
            if not step_key:
                logger.warning(
                    "Unable to find config template for step %s. Falling "
                    "back to the pipeline image.",
                    self.pipeline_step_name,
                )
                step_key = None
        else:
            step_key = self.pipeline_step_name

        return self.snapshot.build.get_image(component_key=key, step=step_key)
