"""Abstractions around GCP resources and nodes.

The logic has been abstracted away here to allow for different GCP resources
(API endpoints), which can differ widely, making it impossible to use
the same logic for everything.

Classes inheriting from ``GCPResource`` represent different GCP resources -
API endpoints that allow for nodes to be created, removed, listed and
otherwise managed. Those classes contain methods abstracting GCP REST API
calls.
Each resource has a corresponding node type, represented by a
class inheriting from ``GCPNode``. Those classes are essentially dicts
with some extra methods. The instances of those classes will be created
from API responses.

The ``GCPNodeType`` enum is a lightweight way to classify nodes.

Currently, Compute and TPU resources & nodes are supported.

In order to add support for new resources, create classes inheriting from
``GCPResource`` and ``GCPNode``, update the ``GCPNodeType`` enum,
update the ``_generate_node_name`` method and finally update the
node provider.
"""

import abc
import logging
import re
import time
from collections import UserDict
from copy import deepcopy
from enum import Enum
from functools import wraps
from typing import Any, Dict, List, Optional, Tuple, Union
from uuid import uuid4

import httplib2
from google_auth_httplib2 import AuthorizedHttp
from googleapiclient.discovery import Resource
from googleapiclient.errors import HttpError

from ray.autoscaler.tags import TAG_RAY_CLUSTER_NAME, TAG_RAY_NODE_NAME

logger = logging.getLogger(__name__)

INSTANCE_NAME_MAX_LEN = 64
INSTANCE_NAME_UUID_LEN = 8
MAX_POLLS = 12
# TPUs take a long while to respond, so we increase the MAX_POLLS
# considerably - this probably could be smaller
# TPU deletion uses MAX_POLLS
MAX_POLLS_TPU = MAX_POLLS * 8
POLL_INTERVAL = 5


def _retry_on_exception(
    exception: Union[Exception, Tuple[Exception]],
    regex: Optional[str] = None,
    max_retries: int = MAX_POLLS,
    retry_interval_s: int = POLL_INTERVAL,
):
    """Retry a function call n-times for as long as it throws an exception."""

    def dec(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            def try_catch_exc():
                try:
                    value = func(*args, **kwargs)
                    return value
                except Exception as e:
                    if not isinstance(e, exception) or (
                        regex and not re.search(regex, str(e))
                    ):
                        raise e
                    return e

            for _ in range(max_retries):
                ret = try_catch_exc()
                if not isinstance(ret, Exception):
                    break
                time.sleep(retry_interval_s)
            if isinstance(ret, Exception):
                raise ret
            return ret

        return wrapper

    return dec


def _generate_node_name(labels: dict, node_suffix: str) -> str:
    """Generate node name from labels and suffix.

    This is required so that the correct resource can be selected
    when the only information autoscaler has is the name of the node.

    The suffix is expected to be one of 'compute' or 'tpu'
    (as in ``GCPNodeType``).
    """
    name_label = labels[TAG_RAY_NODE_NAME]
    assert len(name_label) <= (INSTANCE_NAME_MAX_LEN - INSTANCE_NAME_UUID_LEN - 1), (
        name_label,
        len(name_label),
    )
    return f"{name_label}-{uuid4().hex[:INSTANCE_NAME_UUID_LEN]}-{node_suffix}"


class GCPNodeType(Enum):
    """Enum for GCP node types (compute & tpu)"""

    COMPUTE = "compute"
    TPU = "tpu"

    @staticmethod
    def from_gcp_node(node: "GCPNode"):
        """Return GCPNodeType based on ``node``'s class"""
        if isinstance(node, GCPTPUNode):
            return GCPNodeType.TPU
        if isinstance(node, GCPComputeNode):
            return GCPNodeType.COMPUTE
        raise TypeError(f"Wrong GCPNode type {type(node)}.")

    @staticmethod
    def name_to_type(name: str):
        """Provided a node name, determine the type.

        This expects the name to be in format '[NAME]-[UUID]-[TYPE]',
        where [TYPE] is either 'compute' or 'tpu'.
        """
        return GCPNodeType(name.split("-")[-1])


class GCPNode(UserDict, metaclass=abc.ABCMeta):
    """Abstraction around compute and tpu nodes"""

    NON_TERMINATED_STATUSES = None
    RUNNING_STATUSES = None
    STATUS_FIELD = None

    def __init__(self, base_dict: dict, resource: "GCPResource", **kwargs) -> None:
        super().__init__(base_dict, **kwargs)
        self.resource = resource
        assert isinstance(self.resource, GCPResource)

    def is_running(self) -> bool:
        return self.get(self.STATUS_FIELD) in self.RUNNING_STATUSES

    def is_terminated(self) -> bool:
        return self.get(self.STATUS_FIELD) not in self.NON_TERMINATED_STATUSES

    @abc.abstractmethod
    def get_labels(self) -> dict:
        return

    @abc.abstractmethod
    def get_external_ip(self) -> str:
        return

    @abc.abstractmethod
    def get_internal_ip(self) -> str:
        return

    def __repr__(self) -> str:
        return f"<{self.__class__.__name__}: {self.get('name')}>"


class GCPComputeNode(GCPNode):
    """Abstraction around compute nodes"""

    # https://cloud.google.com/compute/docs/instances/instance-life-cycle
    NON_TERMINATED_STATUSES = {"PROVISIONING", "STAGING", "RUNNING"}
    TERMINATED_STATUSES = {"TERMINATED", "SUSPENDED"}
    RUNNING_STATUSES = {"RUNNING"}
    STATUS_FIELD = "status"

    def get_labels(self) -> dict:
        return self.get("labels", {})

    def get_external_ip(self) -> str:
        return (
            self.get("networkInterfaces", [{}])[0]
            .get("accessConfigs", [{}])[0]
            .get("natIP", None)
        )

    def get_internal_ip(self) -> str:
        return self.get("networkInterfaces", [{}])[0].get("networkIP")


class GCPTPUNode(GCPNode):
    """Abstraction around tpu nodes"""

    # https://cloud.google.com/tpu/docs/reference/rest/v2alpha1/projects.locations.nodes#State

    NON_TERMINATED_STATUSES = {"CREATING", "STARTING", "RESTARTING", "READY"}
    RUNNING_STATUSES = {"READY"}
    STATUS_FIELD = "state"

    def get_labels(self) -> dict:
        return self.get("labels", {})

    @property
    def num_workers(self) -> int:
        return len(self.get("networkEndpoints", [{}]))

    def get_external_ips(self) -> List[str]:
        return self.get("networkEndpoints", [{}])

    def get_external_ip(self, worker_index: int = 0) -> str:
        return (
            self.get_external_ips()[worker_index]
            .get("accessConfig", {})
            .get("externalIp", None)
        )

    def get_internal_ips(self) -> List[str]:
        return self.get("networkEndpoints", [{}])

    def get_internal_ip(self, worker_index: int = 0) -> str:
        return self.get_internal_ips()[worker_index].get("ipAddress", None)


class GCPResource(metaclass=abc.ABCMeta):
    """Abstraction around compute and TPU resources"""

    def __init__(
        self,
        resource: Resource,
        project_id: str,
        availability_zone: str,
        cluster_name: str,
    ) -> None:
        self.resource = resource
        self.project_id = project_id
        self.availability_zone = availability_zone
        self.cluster_name = cluster_name

    @abc.abstractmethod
    def get_new_authorized_http(self, http: AuthorizedHttp) -> AuthorizedHttp:
        """Generate a new AuthorizedHttp object with the given credentials."""
        return

    @abc.abstractmethod
    def wait_for_operation(
        self,
        operation: dict,
        max_polls: int = MAX_POLLS,
        poll_interval: int = POLL_INTERVAL,
    ) -> dict:
        """Waits a preset amount of time for operation to complete."""
        return None

    @abc.abstractmethod
    def list_instances(
        self,
        label_filters: Optional[dict] = None,
        is_terminated: bool = False,
    ) -> List["GCPNode"]:
        """Returns a filtered list of all instances.

        The filter removes all terminated instances and, if ``label_filters``
        are provided, all instances which labels are not matching the
        ones provided.
        """
        return

    @abc.abstractmethod
    def get_instance(self, node_id: str) -> "GCPNode":
        """Returns a single instance."""
        return

    @abc.abstractmethod
    def set_labels(
        self, node: GCPNode, labels: dict, wait_for_operation: bool = True
    ) -> dict:
        """Sets labels on an instance and returns result.

        Completely replaces the labels dictionary."""
        return

    @abc.abstractmethod
    def create_instance(
        self, base_config: dict, labels: dict, wait_for_operation: bool = True
    ) -> Tuple[dict, str]:
        """Creates a single instance and returns result.

        Returns a tuple of (result, node_name).
        """
        return

    def create_instances(
        self,
        base_config: dict,
        labels: dict,
        count: int,
        wait_for_operation: bool = True,
    ) -> List[Tuple[dict, str]]:
        """Creates multiple instances and returns result.

        Returns a list of tuples of (result, node_name).
        """
        operations = [
            self.create_instance(base_config, labels, wait_for_operation=False)
            for i in range(count)
        ]

        if wait_for_operation:
            results = [
                (self.wait_for_operation(operation), node_name)
                for operation, node_name in operations
            ]
        else:
            results = operations

        return results

    @abc.abstractmethod
    def delete_instance(self, node_id: str, wait_for_operation: bool = True) -> dict:
        """Deletes an instance and returns result."""
        return

    @abc.abstractmethod
    def stop_instance(self, node_id: str, wait_for_operation: bool = True) -> dict:
        """Deletes an instance and returns result."""
        return

    @abc.abstractmethod
    def start_instance(self, node_id: str, wait_for_operation: bool = True) -> dict:
        """Starts a single instance and returns result."""
        return


class GCPCompute(GCPResource):
    """Abstraction around GCP compute resource"""

    def get_new_authorized_http(self, http: AuthorizedHttp) -> AuthorizedHttp:
        """Generate a new AuthorizedHttp object with the given credentials."""
        new_http = AuthorizedHttp(http.credentials, http=httplib2.Http())
        return new_http

    def wait_for_operation(
        self,
        operation: dict,
        max_polls: int = MAX_POLLS,
        poll_interval: int = POLL_INTERVAL,
    ) -> dict:
        """Poll for compute zone operation until finished."""
        logger.info(
            "wait_for_compute_zone_operation: "
            f"Waiting for operation {operation['name']} to finish..."
        )

        for _ in range(max_polls):
            result = (
                self.resource.zoneOperations()
                .get(
                    project=self.project_id,
                    operation=operation["name"],
                    zone=self.availability_zone,
                )
                .execute(http=self.get_new_authorized_http(self.resource._http))
            )
            if "error" in result:
                raise Exception(result["error"])

            if result["status"] == "DONE":
                logger.info(
                    "wait_for_compute_zone_operation: "
                    f"Operation {operation['name']} finished."
                )
                break

            time.sleep(poll_interval)

        return result

    def list_instances(
        self,
        label_filters: Optional[dict] = None,
        is_terminated: bool = False,
    ) -> List[GCPComputeNode]:
        label_filters = label_filters or {}

        if label_filters:
            label_filter_expr = (
                "("
                + " AND ".join(
                    [
                        "(labels.{key} = {value})".format(key=key, value=value)
                        for key, value in label_filters.items()
                    ]
                )
                + ")"
            )
        else:
            label_filter_expr = ""

        statuses = (
            GCPComputeNode.TERMINATED_STATUSES
            if is_terminated
            else GCPComputeNode.NON_TERMINATED_STATUSES
        )

        instance_state_filter_expr = (
            "("
            + " OR ".join(
                ["(status = {status})".format(status=status) for status in statuses]
            )
            + ")"
        )

        cluster_name_filter_expr = "(labels.{key} = {value})".format(
            key=TAG_RAY_CLUSTER_NAME, value=self.cluster_name
        )

        # TPU VMs spawn accompanying Compute Instances that must be filtered out,
        # else this results in duplicated nodes.
        tpu_negation_filter_expr = "(NOT labels.{label}:*)".format(label="tpu_cores")

        not_empty_filters = [
            f
            for f in [
                label_filter_expr,
                instance_state_filter_expr,
                cluster_name_filter_expr,
                tpu_negation_filter_expr,
            ]
            if f
        ]

        filter_expr = " AND ".join(not_empty_filters)

        response = (
            self.resource.instances()
            .list(
                project=self.project_id,
                zone=self.availability_zone,
                filter=filter_expr,
            )
            .execute(http=self.get_new_authorized_http(self.resource._http))
        )

        instances = response.get("items", [])
        return [GCPComputeNode(i, self) for i in instances]

    def get_instance(self, node_id: str) -> GCPComputeNode:
        instance = (
            self.resource.instances()
            .get(
                project=self.project_id,
                zone=self.availability_zone,
                instance=node_id,
            )
            .execute()
        )

        return GCPComputeNode(instance, self)

    def set_labels(
        self, node: GCPComputeNode, labels: dict, wait_for_operation: bool = True
    ) -> dict:
        body = {
            "labels": dict(node["labels"], **labels),
            "labelFingerprint": node["labelFingerprint"],
        }
        node_id = node["name"]
        operation = (
            self.resource.instances()
            .setLabels(
                project=self.project_id,
                zone=self.availability_zone,
                instance=node_id,
                body=body,
            )
            .execute(http=self.get_new_authorized_http(self.resource._http))
        )

        if wait_for_operation:
            result = self.wait_for_operation(operation)
        else:
            result = operation

        return result

    def _convert_resources_to_urls(
        self, configuration_dict: Dict[str, Any]
    ) -> Dict[str, Any]:
        """Ensures that resources are in their full URL form.

        GCP expects machineType and acceleratorType to be a full URL (e.g.
        `zones/us-west1/machineTypes/n1-standard-2`) instead of just the
        type (`n1-standard-2`)

        Args:
            configuration_dict: Dict of options that will be passed to GCP
        Returns:
            Input dictionary, but with possibly expanding `machineType` and
                `acceleratorType`.
        """
        configuration_dict = deepcopy(configuration_dict)
        existing_machine_type = configuration_dict["machineType"]
        if not re.search(".*/machineTypes/.*", existing_machine_type):
            configuration_dict[
                "machineType"
            ] = "zones/{zone}/machineTypes/{machine_type}".format(
                zone=self.availability_zone,
                machine_type=configuration_dict["machineType"],
            )

        for accelerator in configuration_dict.get("guestAccelerators", []):
            gpu_type = accelerator["acceleratorType"]
            if not re.search(".*/acceleratorTypes/.*", gpu_type):
                accelerator[
                    "acceleratorType"
                ] = "projects/{project}/zones/{zone}/acceleratorTypes/{accelerator}".format(  # noqa: E501
                    project=self.project_id,
                    zone=self.availability_zone,
                    accelerator=gpu_type,
                )

        return configuration_dict

    def create_instance(
        self, base_config: dict, labels: dict, wait_for_operation: bool = True
    ) -> Tuple[dict, str]:

        config = self._convert_resources_to_urls(base_config)
        # removing TPU-specific default key set in config.py
        config.pop("networkConfig", None)
        name = _generate_node_name(labels, GCPNodeType.COMPUTE.value)

        labels = dict(config.get("labels", {}), **labels)

        config.update(
            {
                "labels": dict(labels, **{TAG_RAY_CLUSTER_NAME: self.cluster_name}),
                "name": name,
            }
        )

        # Allow Google Compute Engine instance templates.
        #
        # Config example:
        #
        #     ...
        #     node_config:
        #         sourceInstanceTemplate: global/instanceTemplates/worker-16
        #         machineType: e2-standard-16
        #     ...
        #
        # node_config parameters override matching template parameters, if any.
        #
        # https://cloud.google.com/compute/docs/instance-templates
        # https://cloud.google.com/compute/docs/reference/rest/v1/instances/insert
        source_instance_template = config.pop("sourceInstanceTemplate", None)

        operation = (
            self.resource.instances()
            .insert(
                project=self.project_id,
                zone=self.availability_zone,
                sourceInstanceTemplate=source_instance_template,
                body=config,
            )
            .execute(http=self.get_new_authorized_http(self.resource._http))
        )

        if wait_for_operation:
            result = self.wait_for_operation(operation)
        else:
            result = operation

        return result, name

    def delete_instance(self, node_id: str, wait_for_operation: bool = True) -> dict:
        operation = (
            self.resource.instances()
            .delete(
                project=self.project_id,
                zone=self.availability_zone,
                instance=node_id,
            )
            .execute()
        )

        if wait_for_operation:
            result = self.wait_for_operation(operation)
        else:
            result = operation

        return result

    def stop_instance(self, node_id: str, wait_for_operation: bool = True) -> dict:
        operation = (
            self.resource.instances()
            .stop(
                project=self.project_id,
                zone=self.availability_zone,
                instance=node_id,
            )
            .execute()
        )

        if wait_for_operation:
            result = self.wait_for_operation(operation)
        else:
            result = operation
        return result

    def start_instance(self, node_id: str, wait_for_operation: bool = True) -> dict:

        operation = (
            self.resource.instances()
            .start(
                project=self.project_id,
                zone=self.availability_zone,
                instance=node_id,
            )
            .execute(http=self.get_new_authorized_http(self.resource._http))
        )

        if wait_for_operation:
            result = self.wait_for_operation(operation)
        else:
            result = operation
        return result


class GCPTPU(GCPResource):
    """Abstraction around GCP TPU resource"""

    # node names already contain the path, but this is required for `parent`
    # arguments
    @property
    def path(self):
        return f"projects/{self.project_id}/locations/{self.availability_zone}"

    def get_new_authorized_http(self, http: AuthorizedHttp) -> AuthorizedHttp:
        """Generate a new AuthorizedHttp object with the given credentials."""
        new_http = AuthorizedHttp(http.credentials, http=httplib2.Http())
        return new_http

    def wait_for_operation(
        self,
        operation: dict,
        max_polls: int = MAX_POLLS_TPU,
        poll_interval: int = POLL_INTERVAL,
    ) -> dict:
        """Poll for TPU operation until finished."""
        logger.info(
            "wait_for_tpu_operation: "
            f"Waiting for operation {operation['name']} to finish..."
        )

        for _ in range(max_polls):
            result = (
                self.resource.projects()
                .locations()
                .operations()
                .get(name=f"{operation['name']}")
                .execute(http=self.get_new_authorized_http(self.resource._http))
            )
            if "error" in result:
                raise Exception(result["error"])

            if "response" in result:
                logger.info(
                    "wait_for_tpu_operation: "
                    f"Operation {operation['name']} finished."
                )
                break

            time.sleep(poll_interval)

        return result

    def list_instances(
        self,
        label_filters: Optional[dict] = None,
        is_terminated: bool = False,
    ) -> List[GCPTPUNode]:
        response = (
            self.resource.projects()
            .locations()
            .nodes()
            .list(parent=self.path)
            .execute(http=self.get_new_authorized_http(self.resource._http))
        )

        instances = response.get("nodes", [])
        instances = [GCPTPUNode(i, self) for i in instances]

        # filter_expr cannot be passed directly to API
        # so we need to filter the results ourselves

        # same logic as in GCPCompute.list_instances
        label_filters = label_filters or {}
        label_filters[TAG_RAY_CLUSTER_NAME] = self.cluster_name

        def filter_instance(instance: GCPTPUNode) -> bool:
            if instance.is_terminated():
                return False

            labels = instance.get_labels()
            if label_filters:
                for key, value in label_filters.items():
                    if key not in labels:
                        return False
                    if value != labels[key]:
                        return False

            return True

        instances = list(filter(filter_instance, instances))

        return instances

    def get_instance(self, node_id: str) -> GCPTPUNode:
        instance = (
            self.resource.projects()
            .locations()
            .nodes()
            .get(name=node_id)
            .execute(http=self.get_new_authorized_http(self.resource._http))
        )

        return GCPTPUNode(instance, self)

    # this sometimes fails without a clear reason, so we retry it
    # MAX_POLLS times
    @_retry_on_exception(HttpError, "unable to queue the operation")
    def set_labels(
        self, node: GCPTPUNode, labels: dict, wait_for_operation: bool = True
    ) -> dict:
        body = {
            "labels": dict(node["labels"], **labels),
        }
        update_mask = "labels"

        operation = (
            self.resource.projects()
            .locations()
            .nodes()
            .patch(
                name=node["name"],
                updateMask=update_mask,
                body=body,
            )
            .execute(http=self.get_new_authorized_http(self.resource._http))
        )

        if wait_for_operation:
            result = self.wait_for_operation(operation)
        else:
            result = operation

        return result

    def create_instance(
        self, base_config: dict, labels: dict, wait_for_operation: bool = True
    ) -> Tuple[dict, str]:
        config = base_config.copy()
        # removing Compute-specific default key set in config.py
        config.pop("networkInterfaces", None)
        name = _generate_node_name(labels, GCPNodeType.TPU.value)

        labels = dict(config.get("labels", {}), **labels)

        config.update(
            {
                "labels": dict(labels, **{TAG_RAY_CLUSTER_NAME: self.cluster_name}),
            }
        )

        if "networkConfig" not in config:
            config["networkConfig"] = {}
        if "enableExternalIps" not in config["networkConfig"]:
            # this is required for SSH to work, per google documentation
            # https://cloud.google.com/tpu/docs/users-guide-tpu-vm#create-curl
            config["networkConfig"]["enableExternalIps"] = True

        # replace serviceAccounts with serviceAccount, and scopes with scope
        # this is necessary for the head node to work
        # see here: https://tpu.googleapis.com/$discovery/rest?version=v2alpha1
        if "serviceAccounts" in config:
            config["serviceAccount"] = config.pop("serviceAccounts")[0]
            config["serviceAccount"]["scope"] = config["serviceAccount"].pop("scopes")

        operation = (
            self.resource.projects()
            .locations()
            .nodes()
            .create(
                parent=self.path,
                body=config,
                nodeId=name,
            )
            .execute(http=self.get_new_authorized_http(self.resource._http))
        )

        if wait_for_operation:
            result = self.wait_for_operation(operation)
        else:
            result = operation

        return result, name

    def delete_instance(self, node_id: str, wait_for_operation: bool = True) -> dict:
        operation = (
            self.resource.projects()
            .locations()
            .nodes()
            .delete(name=node_id)
            .execute(http=self.get_new_authorized_http(self.resource._http))
        )

        # No need to increase MAX_POLLS for deletion
        if wait_for_operation:
            result = self.wait_for_operation(operation, max_polls=MAX_POLLS)
        else:
            result = operation

        return result

    def stop_instance(self, node_id: str, wait_for_operation: bool = True) -> dict:
        operation = (
            self.resource.projects()
            .locations()
            .nodes()
            .stop(name=node_id)
            .execute(http=self.get_new_authorized_http(self.resource._http))
        )

        if wait_for_operation:
            result = self.wait_for_operation(operation, max_polls=MAX_POLLS)
        else:
            result = operation

        return result

    def start_instance(self, node_id: str, wait_for_operation: bool = True) -> dict:
        operation = (
            self.resource.projects()
            .locations()
            .nodes()
            .start(name=node_id)
            .execute(http=self.get_new_authorized_http(self.resource._http))
        )

        if wait_for_operation:
            result = self.wait_for_operation(operation, max_polls=MAX_POLLS)
        else:
            result = operation

        return result