#  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.
"""Abstract base class for entrypoint configurations."""

import argparse
import os
import sys
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, Any, Dict, List, NoReturn, Optional
from uuid import UUID

from zenml.client import Client
from zenml.code_repositories import BaseCodeRepository
from zenml.enums import StackComponentType
from zenml.exceptions import CustomFlavorImportError
from zenml.logger import get_logger
from zenml.utils import (
    code_repository_utils,
    code_utils,
    source_utils,
    uuid_utils,
)

if TYPE_CHECKING:
    from zenml.artifact_stores import BaseArtifactStore
    from zenml.config import DockerSettings
    from zenml.models import CodeReferenceResponse, PipelineSnapshotResponse

logger = get_logger(__name__)
DEFAULT_ENTRYPOINT_COMMAND = [
    "python",
    "-m",
    "zenml.entrypoints.entrypoint",
]

ENTRYPOINT_CONFIG_SOURCE_OPTION = "entrypoint_config_source"
SNAPSHOT_ID_OPTION = "snapshot_id"


class BaseEntrypointConfiguration(ABC):
    """Abstract base class for entrypoint configurations.

    An entrypoint configuration specifies the arguments that should be passed
    to the entrypoint and what is running inside the entrypoint.

    Attributes:
        entrypoint_args: The parsed arguments passed to the entrypoint.
    """

    def __init__(self, arguments: List[str]):
        """Initializes the entrypoint configuration.

        Args:
            arguments: Command line arguments to configure this object.
        """
        self.entrypoint_args = self._parse_arguments(arguments)
        self._snapshot: Optional["PipelineSnapshotResponse"] = None

    @classmethod
    def get_entrypoint_command(cls) -> List[str]:
        """Returns a command that runs the entrypoint module.

        This entrypoint module is responsible for running the entrypoint
        configuration when called. Defaults to running the
        `zenml.entrypoints.entrypoint` module.

        **Note**: This command won't work on its own but needs to be called with
            the arguments returned by the `get_entrypoint_arguments(...)`
            method of this class.

        Returns:
            A list of strings with the command.
        """
        return DEFAULT_ENTRYPOINT_COMMAND

    @classmethod
    def get_entrypoint_options(cls) -> Dict[str, bool]:
        """Gets all options required for running with this configuration.

        Returns:
            A dictionary of options and whether they are required.
        """
        return {
            # Importable source pointing to the entrypoint configuration class
            # that should be used inside the entrypoint.
            ENTRYPOINT_CONFIG_SOURCE_OPTION: True,
            # ID of the pipeline snapshot to use in this entrypoint
            SNAPSHOT_ID_OPTION: True,
        }

    @classmethod
    def get_entrypoint_arguments(
        cls,
        **kwargs: Any,
    ) -> List[str]:
        """Gets all arguments that the entrypoint command should be called with.

        The argument list should be something that
        `argparse.ArgumentParser.parse_args(...)` can handle (e.g.
        `["--some_option", "some_value"]` or `["--some_option=some_value"]`).
        It needs to provide values for all options returned by the
        `get_entrypoint_options()` method of this class.

        Args:
            **kwargs: Keyword args.

        Returns:
            A list of strings with the arguments.

        Raises:
            ValueError: If no valid snapshot ID is passed.
        """
        arguments = [
            f"--{ENTRYPOINT_CONFIG_SOURCE_OPTION}",
            source_utils.resolve(cls).import_path,
        ]

        if SNAPSHOT_ID_OPTION in cls.get_entrypoint_options():
            snapshot_id = kwargs.get(SNAPSHOT_ID_OPTION)
            if not uuid_utils.is_valid_uuid(snapshot_id):
                raise ValueError(
                    f"Missing or invalid snapshot ID as argument for entrypoint "
                    f"configuration. Please make sure to pass a valid UUID to "
                    f"`{cls.__name__}.{cls.get_entrypoint_arguments.__name__}"
                    f"({SNAPSHOT_ID_OPTION}=<UUID>)`."
                )

            arguments.extend(
                [
                    f"--{SNAPSHOT_ID_OPTION}",
                    str(snapshot_id),
                ]
            )

        return arguments

    @classmethod
    def _parse_arguments(cls, arguments: List[str]) -> Dict[str, Any]:
        """Parses command line arguments.

        This method will create an `argparse.ArgumentParser` and add required
        arguments for all the options specified in the
        `get_entrypoint_options()` method of this class.

        Args:
            arguments: Arguments to parse. The format should be something that
                `argparse.ArgumentParser.parse_args(...)` can handle (e.g.
                `["--some_option", "some_value"]` or
                `["--some_option=some_value"]`).

        Returns:
            Dictionary of the parsed arguments.

        # noqa: DAR402
        Raises:
            ValueError: If the arguments are not valid.
        """

        class _CustomParser(argparse.ArgumentParser):
            """Argument parser subclass that suppresses some argparse logs.

            Also raises an exception instead of the `sys.exit()` call.
            """

            def error(self, message: str) -> NoReturn:
                raise ValueError(
                    f"Failed to parse entrypoint arguments: {message}"
                )

        parser = _CustomParser()

        for option_name, required in cls.get_entrypoint_options().items():
            if option_name == ENTRYPOINT_CONFIG_SOURCE_OPTION:
                # This option is already used by
                # `zenml.entrypoints.entrypoint` to read which config
                # class to use
                continue
            parser.add_argument(f"--{option_name}", required=required)

        result, _ = parser.parse_known_args(arguments)
        return vars(result)

    @property
    def snapshot(self) -> "PipelineSnapshotResponse":
        """The snapshot configured for this entrypoint configuration.

        Returns:
            The snapshot.
        """
        if self._snapshot is None:
            self._snapshot = self._load_snapshot()
        return self._snapshot

    @property
    def docker_settings(self) -> "DockerSettings":
        """The Docker settings configured for this entrypoint configuration.

        Returns:
            The Docker settings.
        """
        return self.snapshot.pipeline_configuration.docker_settings

    @property
    def should_download_code(self) -> bool:
        """Whether code should be downloaded.

        Returns:
            Whether code should be downloaded.
        """
        if (
            self.snapshot.code_reference
            and self.docker_settings.allow_download_from_code_repository
        ):
            return True

        if (
            self.snapshot.code_path
            and self.docker_settings.allow_download_from_artifact_store
        ):
            return True

        return False

    def _load_snapshot(self) -> "PipelineSnapshotResponse":
        """Loads the snapshot.

        Returns:
            The snapshot.
        """
        snapshot_id = UUID(self.entrypoint_args[SNAPSHOT_ID_OPTION])
        return Client().zen_store.get_snapshot(snapshot_id=snapshot_id)

    def prepare_code_environment(self) -> None:
        """Prepares the code environment.

        Raises:
            CustomFlavorImportError: If the artifact store flavor can't be
                imported.
            RuntimeError: If the current environment requires code download
                but the snapshot does not have a reference to any code.
        """
        # Change the working directory to make sure we're in the correct
        # directory where the files in the Docker image should be included.
        # This is necessary as some services overwrite the working directory
        # configured in the Docker image itself.
        os.makedirs("/app", exist_ok=True)
        os.chdir("/app")

        if self.should_download_code:
            if code_path := self.snapshot.code_path:
                # Load the artifact store not from the active stack but separately.
                # This is required in case the stack has custom flavor components
                # (other than the artifact store) for which the flavor
                # implementations will only be available once the download finishes.
                try:
                    artifact_store = self._load_active_artifact_store()
                except CustomFlavorImportError as e:
                    raise CustomFlavorImportError(
                        "Failed to import custom artifact store flavor. The "
                        "artifact store flavor is needed to download your code, "
                        "but it looks like it might be part of the files "
                        "that we're trying to download. If this is the case, you "
                        "should disable downloading code from the artifact store "
                        "using `DockerSettings(allow_download_from_artifact_store=False)` "
                        "or make sure the artifact flavor files are included in "
                        "Docker image by using a custom parent image or installing "
                        "them as part of a pip dependency."
                    ) from e
                code_utils.download_code_from_artifact_store(
                    code_path=code_path, artifact_store=artifact_store
                )
            elif code_reference := self.snapshot.code_reference:
                # TODO: This might fail if the code repository had unpushed changes
                # at the time the pipeline run was started.
                self.download_code_from_code_repository(
                    code_reference=code_reference
                )
            else:
                raise RuntimeError(
                    "Code download required but no code reference or path provided."
                )

            logger.info("Code download finished.")

        # If the working directory is not in the sys.path, we include it to make
        # sure user code gets correctly imported.
        cwd = os.getcwd()
        if cwd not in sys.path:
            sys.path.insert(0, cwd)

    def download_code_from_code_repository(
        self, code_reference: "CodeReferenceResponse"
    ) -> None:
        """Download code from a code repository.

        Args:
            code_reference: The reference to the code.
        """
        logger.info(
            "Downloading code from code repository `%s` (commit `%s`).",
            code_reference.code_repository.name,
            code_reference.commit,
        )

        model = Client().get_code_repository(code_reference.code_repository.id)
        repo = BaseCodeRepository.from_model(model)
        code_repo_root = os.path.abspath("code")
        download_dir = os.path.join(
            code_repo_root, code_reference.subdirectory
        )
        os.makedirs(download_dir, exist_ok=True)
        repo.download_files(
            commit=code_reference.commit,
            directory=download_dir,
            repo_sub_directory=code_reference.subdirectory,
        )
        source_utils.set_custom_source_root(download_dir)
        code_repository_utils.set_custom_local_repository(
            root=code_repo_root, commit=code_reference.commit, repo=repo
        )

        sys.path.insert(0, download_dir)
        os.chdir(download_dir)

    def _load_active_artifact_store(self) -> "BaseArtifactStore":
        """Load the active artifact store.

        Returns:
            The active artifact store.
        """
        from zenml.artifact_stores import BaseArtifactStore

        artifact_store_model = Client().active_stack_model.components[
            StackComponentType.ARTIFACT_STORE
        ][0]
        artifact_store = BaseArtifactStore.from_model(artifact_store_model)
        assert isinstance(artifact_store, BaseArtifactStore)

        return artifact_store

    @abstractmethod
    def run(self) -> None:
        """Runs the entrypoint configuration."""
