#  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.
"""Implementation of the VertexAI experiment tracker for ZenML."""

import re
from typing import TYPE_CHECKING, Dict, Optional, Type, cast

from google.api_core import exceptions
from google.cloud import aiplatform
from google.cloud.aiplatform.compat.types import execution

from zenml.constants import METADATA_EXPERIMENT_TRACKER_URL
from zenml.experiment_trackers.base_experiment_tracker import (
    BaseExperimentTracker,
)
from zenml.integrations.gcp.flavors.vertex_experiment_tracker_flavor import (
    VertexExperimentTrackerConfig,
    VertexExperimentTrackerSettings,
)
from zenml.integrations.gcp.google_credentials_mixin import (
    GoogleCredentialsMixin,
)
from zenml.logger import get_logger
from zenml.metadata.metadata_types import Uri

if TYPE_CHECKING:
    from zenml.config.step_run_info import StepRunInfo
    from zenml.metadata.metadata_types import MetadataType

logger = get_logger(__name__)


class VertexExperimentTracker(BaseExperimentTracker, GoogleCredentialsMixin):
    """Track experiments using VertexAI."""

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

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

    @property
    def settings_class(self) -> Type[VertexExperimentTrackerSettings]:
        """Returns the `BaseSettings` settings class.

        Returns:
            The settings class.
        """
        return VertexExperimentTrackerSettings

    def prepare_step_run(self, info: "StepRunInfo") -> None:
        """Configures a VertexAI run.

        Args:
            info: Info about the step that will be executed.
        """
        self._initialize_vertex(info=info)
        self.experiment_name = self._get_experiment_name(info=info)
        self.run_name = self._get_run_name(info=info)

    def get_step_run_metadata(
        self, info: "StepRunInfo"
    ) -> Dict[str, "MetadataType"]:
        """Get component- and step-specific metadata after a step ran.

        Args:
            info: Info about the step that was executed.

        Returns:
            A dictionary of metadata.
        """
        experiment_name = self._get_experiment_name(info=info)
        run_name = self._get_run_name(info=info)
        tensorboard_resource_name = self._get_tensorboard_resource_name(
            experiment=experiment_name
        )
        dashboard_url = self._get_dashboard_url(experiment=experiment_name)
        return {
            METADATA_EXPERIMENT_TRACKER_URL: Uri(dashboard_url),
            "tensorboard_resource_name": tensorboard_resource_name or "",
            "vertex_run_name": run_name,
        }

    def _format_name(self, name: str) -> str:
        return re.sub(r"[^a-z0-9-]", "-", name.strip().lower())[:128].rstrip(
            "-"
        )

    def _get_experiment_name(self, info: "StepRunInfo") -> str:
        """Gets the experiment name.

        Args:
            info: Info about the step.

        Returns:
            The experiment name.
        """
        settings = cast(
            VertexExperimentTrackerSettings, self.get_settings(info)
        )
        name = settings.experiment or info.pipeline.name
        return self._format_name(name)

    def _get_run_name(self, info: "StepRunInfo") -> str:
        """Gets the run name.

        Args:
            info: Info about the step that will be executed.

        Returns:
            The run name.
        """
        return self._format_name(info.run_name)

    def _get_dashboard_url(self, experiment: str) -> str:
        """Gets the run URL.

        Args:
            experiment: The name of the experiment.

        Returns:
            The run URL.
        """
        resource = aiplatform.Experiment(experiment_name=experiment)
        return cast(str, resource.dashboard_url)

    def _get_tensorboard_resource_name(self, experiment: str) -> Optional[str]:
        resource = aiplatform.Experiment(
            experiment_name=experiment
        ).get_backing_tensorboard_resource()
        resource_name = (
            str(resource.resource_name) if resource is not None else None
        )
        return resource_name

    def _initialize_vertex(self, info: "StepRunInfo") -> None:
        """Initializes a VertexAI run.

        Args:
            info: Info about the step that will be executed.
        """
        settings = cast(
            VertexExperimentTrackerSettings, self.get_settings(info)
        )
        experiment = self._get_experiment_name(info=info)
        run_name = self._get_run_name(info=info)
        credentials, project = self._get_authentication()
        logger.info(
            f"Initializing VertexAI with experiment name {experiment} "
            f"and run name {run_name}."
        )

        aiplatform.init(
            project=project,
            location=self.config.location,
            experiment=experiment,
            experiment_tensorboard=settings.experiment_tensorboard,
            staging_bucket=self.config.staging_bucket,
            credentials=credentials,
            encryption_spec_key_name=self.config.encryption_spec_key_name,
            network=self.config.network,
            api_endpoint=self.config.api_endpoint,
            api_key=self.config.api_key,
            api_transport=self.config.api_transport,
            request_metadata=self.config.request_metadata,
        )

        try:
            aiplatform.start_run(
                run=run_name,
                tensorboard=settings.experiment_tensorboard,
                resume=True,
            )
        except exceptions.NotFound:
            aiplatform.start_run(
                run=run_name,
                tensorboard=settings.experiment_tensorboard,
                resume=False,
            )

        logger.info(
            f"VertexAI experiment dashboard: {self._get_dashboard_url(experiment=experiment)}"
        )
        logger.info(
            f"Tensorboard resource name: {self._get_tensorboard_resource_name(experiment=experiment)}"
        )

    def cleanup_step_run(self, info: "StepRunInfo", step_failed: bool) -> None:
        """Stops the VertexAI run.

        Args:
            info: Info about the step that was executed.
            step_failed: Whether the step failed or not.
        """
        state = (
            execution.Execution.State.FAILED
            if step_failed
            else execution.Execution.State.COMPLETE
        )
        aiplatform.end_run(state=state)
