#  Copyright (c) ZenML GmbH 2024. 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.
"""Implementation of the ZenML HyperAI orchestrator."""

import os
import re
import tempfile
from shlex import quote
from typing import IO, TYPE_CHECKING, Any, Dict, Optional, Type, cast

import paramiko
import yaml

from zenml.config.base_settings import BaseSettings
from zenml.entrypoints import StepEntrypointConfiguration
from zenml.enums import StackComponentType
from zenml.integrations.hyperai.flavors.hyperai_orchestrator_flavor import (
    HyperAIOrchestratorConfig,
    HyperAIOrchestratorSettings,
)
from zenml.logger import get_logger
from zenml.orchestrators import ContainerizedOrchestrator, SubmissionResult
from zenml.stack import Stack, StackValidator

if TYPE_CHECKING:
    from zenml.models import PipelineRunResponse, PipelineSnapshotResponse

logger = get_logger(__name__)

ENV_ZENML_HYPERAI_RUN_ID = "ZENML_HYPERAI_ORCHESTRATOR_RUN_ID"


class HyperAIOrchestrator(ContainerizedOrchestrator):
    """Orchestrator responsible for running pipelines on HyperAI instances."""

    @property
    def config(self) -> HyperAIOrchestratorConfig:
        """Returns the `HyperAIOrchestratorConfig` config.

        Returns:
            The configuration.
        """
        return cast(HyperAIOrchestratorConfig, self._config)

    @property
    def settings_class(self) -> Optional[Type["BaseSettings"]]:
        """Settings class for the HyperAI orchestrator.

        Returns:
            The settings class.
        """
        return HyperAIOrchestratorSettings

    @property
    def validator(self) -> Optional[StackValidator]:
        """Ensures there is an image builder in the stack.

        Returns:
            A `StackValidator` instance.
        """
        return StackValidator(
            required_components={
                StackComponentType.CONTAINER_REGISTRY,
                StackComponentType.IMAGE_BUILDER,
            }
        )

    def get_orchestrator_run_id(self) -> str:
        """Returns the active orchestrator run id.

        Raises:
            RuntimeError: If the environment variable specifying the run id
                is not set.

        Returns:
            The orchestrator run id.
        """
        try:
            return os.environ[ENV_ZENML_HYPERAI_RUN_ID]
        except KeyError:
            raise RuntimeError(
                "Unable to read run id from environment variable "
                f"{ENV_ZENML_HYPERAI_RUN_ID}."
            )

    def _validate_mount_path(self, path: str) -> str:
        """Validates if a given string is in a valid path format.

        Args:
            path: The path to be validated.

        Returns:
            The path in a valid format.

        Raises:
            RuntimeError: If the path is not in a valid format.
        """
        # Define a regular expression pattern to match a valid path format
        pattern = r'^(?:[a-zA-Z]:\\(\\[^\\/:*?"<>|]*)*$|^/([^\0]*)*$)'

        if bool(re.match(pattern, path)):
            return path
        else:
            raise RuntimeError(
                f"Path '{path}' is not in a valid format, so a mount cannot be established."
            )

    def _escape_shell_command(self, command: str) -> str:
        """Escapes a shell command.

        Args:
            command: The command to escape.

        Returns:
            The escaped command.
        """
        return quote(command)

    def _scp_to_hyperai_instance(
        self,
        paramiko_client: paramiko.SSHClient,
        f: IO[str],
        directory_name: str,
        file_name: str,
        description: str,
    ) -> None:
        """Copies a file to a HyperAI instance using SCP.

        Args:
            paramiko_client: The SSH client to use for the SCP transfer.
            f: The file to transfer.
            directory_name: The directory on the HyperAI instance to transfer
                the file to.
            file_name: The name of the file being transferred.
            description: A description of the file being transferred.

        Raises:
            RuntimeError: If the file cannot be written to the HyperAI instance.
        """
        try:
            scp_client = paramiko_client.open_sftp()
            scp_client.put(f.name, f"{directory_name}/{file_name}")
            scp_client.close()
        except FileNotFoundError:
            raise RuntimeError(
                f"Failed to write {description} to HyperAI instance. Does the user have permissions to write?"
            )

    def submit_pipeline(
        self,
        snapshot: "PipelineSnapshotResponse",
        stack: "Stack",
        base_environment: Dict[str, str],
        step_environments: Dict[str, Dict[str, str]],
        placeholder_run: Optional["PipelineRunResponse"] = None,
    ) -> Optional[SubmissionResult]:
        """Submits a pipeline to the orchestrator.

        This method should only submit the pipeline and not wait for it to
        complete. If the orchestrator is configured to wait for the pipeline run
        to complete, a function that waits for the pipeline run to complete can
        be passed as part of the submission result.

        Assumes that:
        - A HyperAI (hyperai.ai) instance is running on the configured IP address.
        - The HyperAI instance has been configured to allow SSH connections from the
            machine running the pipeline.
        - Docker and Docker Compose are installed on the HyperAI instance.
        - A key pair has been generated and the public key has been added to the
            HyperAI instance's `authorized_keys` file.
        - The private key is available in a HyperAI service connector linked to this
            orchestrator.

        Args:
            snapshot: The pipeline snapshot to submit.
            stack: The stack the pipeline will run on.
            base_environment: Base environment shared by all steps. This should
                be set if your orchestrator for example runs one container that
                is responsible for starting all the steps.
            step_environments: Environment variables to set when executing
                specific steps.
            placeholder_run: An optional placeholder run for the snapshot.

        Raises:
            RuntimeError: If running the pipeline fails.

        Returns:
            Optional submission result.
        """
        from zenml.integrations.hyperai.service_connectors.hyperai_service_connector import (
            HyperAIServiceConnector,
        )

        compose_definition: Dict[str, Any] = {"version": "3", "services": {}}
        snapshot_id = snapshot.id

        os.environ[ENV_ZENML_HYPERAI_RUN_ID] = str(snapshot_id)

        # Add each step as a service to the Docker Compose definition
        logger.info("Preparing pipeline steps for deployment.")
        for step_name, step in snapshot.step_configurations.items():
            # Get image
            image = self.get_image(snapshot=snapshot, step_name=step_name)

            step_settings = cast(
                HyperAIOrchestratorSettings, self.get_settings(step)
            )
            container_name = f"{snapshot_id}-{step_name}"

            # Make Compose service definition for step
            compose_definition["services"][container_name] = {
                "image": image,
                "container_name": container_name,
                "network_mode": "host",
                "entrypoint": StepEntrypointConfiguration.get_entrypoint_command(),
                "command": StepEntrypointConfiguration.get_entrypoint_arguments(
                    step_name=step_name, snapshot_id=snapshot.id
                ),
                "volumes": [
                    "{}:{}".format(
                        self._validate_mount_path(mount_from),
                        self._validate_mount_path(mount_to),
                    )
                    for mount_from, mount_to in step_settings.mounts_from_to.items()
                ],
            }

            # Depending on GPU setting, add GPU support to service definition
            if self.config.gpu_enabled_in_container:
                compose_definition["services"][container_name]["deploy"] = {
                    "resources": {
                        "reservations": {
                            "devices": [
                                {"driver": "nvidia", "capabilities": ["gpu"]}
                            ]
                        }
                    }
                }

            # Add run ID to step environment variables
            step_env = step_environments[step_name].copy()

            # Depending on whether it is a scheduled or a realtime pipeline, add
            # potential .env file to service definition for deployment ID override.
            if snapshot.schedule:
                compose_definition["services"][container_name]["env_file"] = [
                    ".env"
                ]
            else:
                step_env[ENV_ZENML_HYPERAI_RUN_ID] = str(snapshot_id)

            compose_definition["services"][container_name]["environment"] = (
                step_env
            )

            # Add dependency on upstream steps if applicable
            upstream_steps = step.spec.upstream_steps

            if len(upstream_steps) > 0:
                compose_definition["services"][container_name][
                    "depends_on"
                ] = {}

                for upstream_step_name in upstream_steps:
                    upstream_container_name = (
                        f"{snapshot_id}-{upstream_step_name}"
                    )
                    compose_definition["services"][container_name][
                        "depends_on"
                    ].update(
                        {
                            upstream_container_name: {
                                "condition": "service_completed_successfully"
                            }
                        }
                    )

        compose_definition_yaml: str = yaml.dump(compose_definition)

        # Connect to configured HyperAI instance
        logger.info(
            "Connecting to HyperAI instance and placing Docker Compose file."
        )
        if connector := self.get_connector():
            paramiko_client = connector.connect()
            if paramiko_client is None:
                raise RuntimeError(
                    "Expected to receive a `paramiko.SSHClient` object from the "
                    "linked connector, but got `None`. This likely originates from "
                    "a misconfigured service connector, typically caused by a wrong "
                    "SSH key type being selected. Please check your "
                    "`hyperai_orchestrator` configuration and make sure that the "
                    "`ssh_key_type` of its connected service connector is set to the "
                    "correct value."
                )
            elif not isinstance(paramiko_client, paramiko.SSHClient):
                raise RuntimeError(
                    f"Expected to receive a `paramiko.SSHClient` object from the "
                    f"linked connector, but got type `{type(paramiko_client)}`."
                )
        else:
            raise RuntimeError(
                "You must link a HyperAI service connector to the orchestrator."
            )

        # Get container registry autologin setting
        if self.config.container_registry_autologin:
            logger.info(
                "Attempting to automatically log in to container registry used by stack."
            )

            # Select stack container registry
            container_registry = stack.container_registry

            # Raise error if no container registry is found
            if not container_registry:
                raise RuntimeError(
                    "Unable to find container registry in stack."
                )

            # Get container registry credentials from its config
            credentials = container_registry.credentials
            if credentials is None:
                raise RuntimeError(
                    "The container registry in the active stack has no "
                    "credentials or service connector configured, but the "
                    "HyperAI orchestrator is set to autologin to the container "
                    "registry. Please configure the container registry with "
                    "credentials or turn off the `container_registry_autologin` "
                    "setting in the HyperAI orchestrator configuration."
                )

            container_registry_url = container_registry.config.uri
            (
                container_registry_username,
                container_registry_password,
            ) = credentials

            # Escape inputs
            container_registry_username = self._escape_shell_command(
                container_registry_username
            )
            container_registry_url = self._escape_shell_command(
                container_registry_url
            )

            # Log in to container registry using --password-stdin
            stdin, stdout, stderr = paramiko_client.exec_command(  # nosec
                f"docker login -u {container_registry_username} "
                f"--password-stdin {container_registry_url}"
            )
            # Send the password to stdin
            stdin.channel.send(
                f"{container_registry_password}\n".encode("utf-8")
            )
            stdin.channel.shutdown_write()

            # Log stdout
            for line in stdout.readlines():
                logger.info(line)

        # Get username from connector
        assert isinstance(connector, HyperAIServiceConnector)
        username = connector.config.username

        # Set up pipeline-runs directory if it doesn't exist
        nonscheduled_directory_name = self._escape_shell_command(
            f"/home/{username}/pipeline-runs"
        )
        directory_name = (
            nonscheduled_directory_name
            if not snapshot.schedule
            else self._escape_shell_command(
                f"/home/{username}/scheduled-pipeline-runs"
            )
        )
        stdin, stdout, stderr = paramiko_client.exec_command(  # nosec
            f"mkdir -p {directory_name}"
        )

        # Get pipeline run id and create directory for it
        orchestrator_run_id = self.get_orchestrator_run_id()
        directory_name = self._escape_shell_command(
            f"{directory_name}/{orchestrator_run_id}"
        )
        stdin, stdout, stderr = paramiko_client.exec_command(  # nosec
            f"mkdir -p {directory_name}"
        )

        # Remove all folders from nonscheduled pipelines if they are 7 days old or older
        if self.config.automatic_cleanup_pipeline_files:
            logger.info(
                "Cleaning up old pipeline files on HyperAI instance. This may take a while."
            )
            stdin, stdout, stderr = paramiko_client.exec_command(  # nosec
                f"find {nonscheduled_directory_name} -type d -ctime +7 -exec rm -rf {{}} +"
            )

        # Create temporary file and write Docker Compose file to it
        with tempfile.NamedTemporaryFile(mode="w", delete=True) as f:
            # Write Docker Compose file to temporary file
            with f.file as f_:
                f_.write(compose_definition_yaml)

            # Scp Docker Compose file to HyperAI instance
            self._scp_to_hyperai_instance(
                paramiko_client,
                f,
                directory_name,
                file_name="docker-compose.yml",
                description="Docker Compose file",
            )

        # Create temporary file and write script to it
        with tempfile.NamedTemporaryFile(mode="w", delete=True) as f:
            # Define bash line and command line
            bash_line = "#!/bin/bash\n"
            command_line = rf'cd {directory_name} && echo {ENV_ZENML_HYPERAI_RUN_ID}="{snapshot_id}_$(date +\%s)" > .env && docker compose up -d'

            # Write script to temporary file
            with f.file as f_:
                f_.write(bash_line)
                f_.write(command_line)

            # Scp script to HyperAI instance
            self._scp_to_hyperai_instance(
                paramiko_client,
                f,
                directory_name,
                file_name="run_pipeline.sh",
                description="startup script",
            )

        # Run or schedule Docker Compose file depending on settings
        if not snapshot.schedule:
            logger.info(
                "Starting ZenML pipeline on HyperAI instance. Depending on the size of your container image, this may take a while..."
            )
            stdin, stdout, stderr = paramiko_client.exec_command(  # nosec
                f"cd {directory_name} && docker compose up -d"
            )

            # Log errors in case of failure
            for line in stderr.readlines():
                logger.info(line)
        elif snapshot.schedule and snapshot.schedule.cron_expression:
            # Get cron expression for scheduled pipeline
            cron_expression = snapshot.schedule.cron_expression
            if not cron_expression:
                raise RuntimeError(
                    "A cron expression is required for scheduled pipelines."
                )
            expected_cron_pattern = r"^(?:(?:[0-9]|[1-5][0-9]|60)(?:,(?:[0-9]|[1-5][0-9]|60))*|[*](?:\/[1-9][0-9]*)?)(?:[ \t]+(?:(?:[0-9]|[0-5][0-9]|60)(?:,(?:[0-9]|[0-5][0-9]|60))*|[*](?:\/[1-9][0-9]*)?)){4}$"
            if not re.match(expected_cron_pattern, cron_expression):
                raise RuntimeError(
                    f"The cron expression '{cron_expression}' is not in a valid format."
                )

            # Log about scheduling
            logger.info(f"Requested cron expression: {cron_expression}")
            logger.info("Scheduling ZenML pipeline on HyperAI instance...")

            # Create cron job for scheduled pipeline on HyperAI instance
            stdin, stdout, stderr = paramiko_client.exec_command(  # nosec
                f"(crontab -l ; echo '{cron_expression} bash {directory_name}/run_pipeline.sh') | crontab -"
            )

            logger.info(
                f"Pipeline scheduled successfully in crontab with cron expression: {cron_expression}"
            )
        elif snapshot.schedule and snapshot.schedule.run_once_start_time:
            # Get start time for scheduled pipeline
            start_time = snapshot.schedule.run_once_start_time

            # Log about scheduling
            logger.info(f"Requested start time: {start_time}")
            logger.info("Scheduling ZenML pipeline on HyperAI instance...")

            # Check if `at` is installed on HyperAI instance
            stdin, stdout, stderr = paramiko_client.exec_command(  # nosec
                "which at"
            )
            if not stdout.readlines():
                raise RuntimeError(
                    "The `at` command is not installed on the HyperAI instance. Please install it to use start times for scheduled pipelines."
                )

            # Convert start time into YYYYMMDDHHMM.SS format
            start_time_str = start_time.strftime("%Y%m%d%H%M.%S")

            # Create cron job for scheduled pipeline on HyperAI instance
            stdin, stdout, stderr = paramiko_client.exec_command(  # nosec
                f"echo 'bash {directory_name}/run_pipeline.sh' | at -t {start_time_str}"
            )

            logger.info(
                f"Pipeline scheduled successfully to run once at: {start_time}"
            )
        else:
            raise RuntimeError(
                "A cron expression or start time is required for scheduled pipelines."
            )

        return None
