#  Copyright (c) ZenML GmbH 2025. 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.
"""Utility functions for running pipelines."""

import time
from typing import (
    TYPE_CHECKING,
    Any,
    Dict,
    List,
    Optional,
    Set,
    Union,
)
from uuid import UUID

from pydantic import BaseModel

from zenml.client import Client
from zenml.config.pipeline_run_configuration import PipelineRunConfiguration
from zenml.config.source import Source, SourceType
from zenml.config.step_configurations import StepConfigurationUpdate
from zenml.enums import ExecutionStatus
from zenml.logger import get_logger
from zenml.models import (
    FlavorFilter,
    LogsRequest,
    PipelineRunRequest,
    PipelineRunResponse,
    PipelineRunTriggerInfo,
    PipelineSnapshotBase,
    PipelineSnapshotResponse,
    StackResponse,
)
from zenml.stack import Flavor, Stack
from zenml.utils import (
    code_utils,
    notebook_utils,
    source_utils,
    string_utils,
)
from zenml.utils.time_utils import utc_now
from zenml.zen_stores.base_zen_store import BaseZenStore

if TYPE_CHECKING:
    StepConfigurationUpdateOrDict = Union[
        Dict[str, Any], StepConfigurationUpdate
    ]

logger = get_logger(__name__)


def get_default_run_name(pipeline_name: str) -> str:
    """Gets the default name for a pipeline run.

    Args:
        pipeline_name: Name of the pipeline which will be run.

    Returns:
        Run name.
    """
    return f"{pipeline_name}-{{date}}-{{time}}"


def create_placeholder_run(
    snapshot: "PipelineSnapshotResponse",
    orchestrator_run_id: Optional[str] = None,
    logs: Optional["LogsRequest"] = None,
    trigger_info: Optional[PipelineRunTriggerInfo] = None,
    original_run_id: Optional[UUID] = None,
) -> "PipelineRunResponse":
    """Create a placeholder run for the snapshot.

    Args:
        snapshot: The snapshot for which to create the placeholder run.
        orchestrator_run_id: The orchestrator run ID for the run.
        logs: The logs for the run.
        trigger_info: The trigger information for the run.
        original_run_id: The original run ID for a replayed run.

    Returns:
        The placeholder run.
    """
    start_time = utc_now()
    run_request = PipelineRunRequest(
        name=string_utils.format_name_template(
            name_template=snapshot.run_name_template,
            substitutions=snapshot.pipeline_configuration.finalize_substitutions(
                start_time=start_time,
            ),
        ),
        # We set the start time on the placeholder run already to
        # make it consistent with the {time} placeholder in the
        # run name. This means the placeholder run will usually
        # have longer durations than scheduled runs, as for them
        # the start_time is only set once the first step starts
        # running.
        start_time=start_time,
        orchestrator_run_id=orchestrator_run_id,
        project=snapshot.project_id,
        snapshot=snapshot.id,
        status=ExecutionStatus.INITIALIZING,
        tags=snapshot.pipeline_configuration.tags,
        logs=logs,
        trigger_info=trigger_info,
        original_run_id=original_run_id,
    )
    run, _ = Client().zen_store.get_or_create_run(run_request)
    return run


def wait_for_pipeline_run_to_finish(run_id: UUID) -> "PipelineRunResponse":
    """Waits until a pipeline run is finished.

    Args:
        run_id: ID of the run for which to wait.

    Returns:
        Model of the finished run.
    """
    sleep_interval = 1
    max_sleep_interval = 64

    while True:
        run = Client().get_pipeline_run(run_id)

        if run.status.is_finished:
            return run

        logger.info(
            "Waiting for pipeline run with ID %s to finish (current status: %s)",
            run_id,
            run.status,
        )
        time.sleep(sleep_interval)
        if sleep_interval < max_sleep_interval:
            sleep_interval *= 2


def validate_stack_is_runnable_from_server(
    zen_store: BaseZenStore, stack: StackResponse
) -> None:
    """Validate if a stack model is runnable from the server.

    Args:
        zen_store: ZenStore to use for listing flavors.
        stack: The stack to validate.

    Raises:
        ValueError: If the stack has components of a custom flavor or local
            components.
    """
    for component_list in stack.components.values():
        assert len(component_list) == 1
        component = component_list[0]
        flavors = zen_store.list_flavors(
            FlavorFilter(name=component.flavor_name, type=component.type)
        )
        assert len(flavors) == 1
        flavor_model = flavors[0]

        if flavor_model.is_custom:
            raise ValueError(
                "Unable to run pipeline from the server on a stack that "
                "includes stack components with a custom flavor."
            )

        flavor = Flavor.from_model(flavor_model)
        component_config = flavor.config_class(**component.configuration)

        if component_config.is_local:
            raise ValueError(
                "Unable to run pipeline from the server on a stack that "
                "includes local stack components."
            )


def validate_run_config_is_runnable_from_server(
    run_configuration: "PipelineRunConfiguration",
    is_dynamic: bool,
) -> None:
    """Validates that the run configuration can be used to run from the server.

    Args:
        run_configuration: The run configuration to validate.
        is_dynamic: Whether the snapshot to run is dynamic.

    Raises:
        ValueError: If there are values in the run configuration that are not
            allowed when running a pipeline from the server.
    """
    if run_configuration.parameters and not is_dynamic:
        raise ValueError(
            "Can't set pipeline parameters when running a static pipeline via "
            "the REST API. You can either use dynamic pipelines for which you "
            "can pass pipeline parameters, or refactore your pipeline code to "
            "use step parameters instead of pipeline parameters. For example, "
            "instead of: "
            "```yaml "
            "parameters: "
            "  param1: 1 "
            "  param2: 2 "
            "``` "
            "You'll need to modify your pipeline code to pass parameters "
            "directly to steps: "
            "```yaml "
            "steps: "
            "  step1: "
            "    parameters: "
            "      param1: 1 "
            "      param2: 2 "
            "``` "
        )

    if run_configuration.build:
        raise ValueError("Can't set build when running pipeline via Rest API.")

    if run_configuration.schedule:
        raise ValueError(
            "Can't set schedule when running pipeline via Rest API."
        )

    if run_configuration.settings and run_configuration.settings.get("docker"):
        raise ValueError(
            "Can't set DockerSettings when running pipeline via Rest API."
        )

    if run_configuration.steps:
        for step_update in run_configuration.steps.values():
            if step_update.settings and step_update.settings.get("docker"):
                raise ValueError(
                    "Can't set DockerSettings when running pipeline via "
                    "Rest API."
                )


def upload_notebook_cell_code_if_necessary(
    snapshot: "PipelineSnapshotBase", stack: "Stack"
) -> None:
    """Upload notebook cell code if necessary.

    This function checks if any of the steps of the pipeline that will be
    executed in a different process are defined in a notebook. If that is the
    case, it will extract that notebook cell code into python files and upload
    an archive of all the necessary files to the artifact store.

    Args:
        snapshot: The snapshot.
        stack: The stack on which the snapshot will happen.

    Raises:
        RuntimeError: If the code for one of the steps that will run out of
            process cannot be extracted into a python file.
    """
    should_upload = False
    resolved_notebook_sources = source_utils.get_resolved_notebook_sources()

    for step in snapshot.step_configurations.values():
        source = step.spec.source

        if source.type == SourceType.NOTEBOOK:
            if (
                stack.orchestrator.flavor != "local"
                or step.config.step_operator
            ):
                should_upload = True
                cell_code = resolved_notebook_sources.get(
                    source.import_path, None
                )

                # Code does not run in-process, which means we need to
                # extract the step code into a python file
                if not cell_code:
                    raise RuntimeError(
                        f"Unable to run step {step.config.name}. This step is "
                        "defined in a notebook and you're trying to run it "
                        "in a remote environment, but ZenML was not able to "
                        "detect the step code in the notebook. To fix "
                        "this error, define your step in a python file instead "
                        "of a notebook."
                    )

    if should_upload:
        logger.info("Uploading notebook code...")

        for _, cell_code in resolved_notebook_sources.items():
            notebook_utils.warn_about_notebook_cell_magic_commands(
                cell_code=cell_code
            )
            module_name = notebook_utils.compute_cell_replacement_module_name(
                cell_code=cell_code
            )
            file_name = f"{module_name}.py"

            code_utils.upload_notebook_code(
                artifact_store=stack.artifact_store,
                cell_code=cell_code,
                file_name=file_name,
            )

        all_snapshot_sources = get_all_sources_from_value(snapshot)

        for source in all_snapshot_sources:
            if source.type == SourceType.NOTEBOOK:
                setattr(source, "artifact_store_id", stack.artifact_store.id)

        logger.info("Upload finished.")


def get_all_sources_from_value(value: Any) -> List[Source]:
    """Get all source objects from a value.

    Args:
        value: The value from which to get all the source objects.

    Returns:
        List of source objects for the given value.
    """
    sources = []
    if isinstance(value, Source):
        sources.append(value)
    elif isinstance(value, BaseModel):
        for v in value.__dict__.values():
            sources.extend(get_all_sources_from_value(v))
    elif isinstance(value, Dict):
        for v in value.values():
            sources.extend(get_all_sources_from_value(v))
    elif isinstance(value, (List, Set, tuple)):
        for v in value:
            sources.extend(get_all_sources_from_value(v))

    return sources
