#  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.
"""Utilities for outputs."""

import os
from typing import TYPE_CHECKING, Dict, Sequence
from uuid import uuid4

from zenml.client import Client
from zenml.constants import IN_MEMORY_ARTIFACT_URI_PREFIX
from zenml.logger import get_logger
from zenml.utils import string_utils

if TYPE_CHECKING:
    from zenml.config.step_configurations import Step
    from zenml.models import StepRunResponse
    from zenml.stack import Stack


logger = get_logger(__name__)


def generate_artifact_uri(
    artifact_store_path: str,
    step_run: "StepRunResponse",
    output_name: str,
) -> str:
    """Generates a URI for an output artifact.

    Args:
        artifact_store_path: The path of the artifact store in which the
            artifact will be stored.
        step_run: The step run that created the artifact.
        output_name: The name of the output in the step run for this artifact.

    Returns:
        The URI of the output artifact.
    """
    for banned_character in ["<", ">", ":", '"', "/", "\\", "|", "?", "*"]:
        output_name = output_name.replace(banned_character, "_")
    return os.path.join(
        artifact_store_path,
        step_run.name,
        output_name,
        str(step_run.id),
        str(uuid4())[:8],  # add random subfolder to avoid collisions
    )


def prepare_output_artifact_uris(
    step_run: "StepRunResponse",
    stack: "Stack",
    step: "Step",
    *,
    skip_artifact_materialization: bool = False,
) -> Dict[str, str]:
    """Prepares the output artifact URIs to run the current step.

    Args:
        step_run: The step run for which to prepare the artifact URIs.
        stack: The stack on which the pipeline is running.
        step: The step configuration.
        skip_artifact_materialization: Whether to skip artifact materialization.

    Raises:
        RuntimeError: If an artifact URI already exists.

    Returns:
        A dictionary mapping output names to artifact URIs.
    """
    artifact_store = stack.artifact_store
    output_artifact_uris: Dict[str, str] = {}

    for output_name in step.config.outputs.keys():
        substituted_output_name = string_utils.format_name_template(
            output_name, substitutions=step_run.config.substitutions
        )
        artifact_store_path = (
            IN_MEMORY_ARTIFACT_URI_PREFIX
            if skip_artifact_materialization
            else artifact_store.path
        )
        artifact_uri = generate_artifact_uri(
            artifact_store_path=artifact_store_path,
            step_run=step_run,
            output_name=substituted_output_name,
        )

        if not skip_artifact_materialization:
            if artifact_store.exists(artifact_uri):
                raise RuntimeError("Artifact already exists")
            artifact_store.makedirs(artifact_uri)
        output_artifact_uris[output_name] = artifact_uri
    return output_artifact_uris


def remove_artifact_dirs(artifact_uris: Sequence[str]) -> None:
    """Removes the artifact directories.

    Args:
        artifact_uris: URIs of the artifacts to remove the directories for.
    """
    # TODO: maybe keep non-empty dirs here, if the step saved intermediate
    # checkpoints?
    artifact_store = Client().active_stack.artifact_store
    for artifact_uri in artifact_uris:
        if artifact_store.isdir(artifact_uri):
            artifact_store.rmtree(artifact_uri)
