# Copyright (C) 2015-2021 Regents of the University of California
#
# 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
#
#     http://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.
import json
import logging
import os
import threading
import time
import uuid
from typing import Optional

from libcloud.compute.drivers.gce import GCEFailedNode
from libcloud.compute.providers import get_driver
from libcloud.compute.types import Provider

from toil.jobStores.googleJobStore import GoogleJobStore
from toil.lib.compatibility import compat_bytes_recursive
from toil.lib.conversions import human2bytes
from toil.lib.web import web_session
from toil.provisioners import NoSuchClusterException
from toil.provisioners.abstractProvisioner import AbstractProvisioner, Shape
from toil.provisioners.node import Node

logger = logging.getLogger(__name__)
logging.getLogger("urllib3.connectionpool").setLevel(logging.WARNING)


class GCEProvisioner(AbstractProvisioner):
    """
    Implements a Google Compute Engine Provisioner using libcloud.
    """

    NODE_BOTO_PATH = "/root/.boto"  # boto file path on instances
    SOURCE_IMAGE = b"projects/kinvolk-public/global/images/family/flatcar-stable"

    def __init__(
        self,
        clusterName,
        clusterType,
        zone,
        nodeStorage,
        nodeStorageOverrides,
        sseKey,
        enable_fuse,
    ):
        self.cloud = "gce"
        self._sseKey = sseKey

        # Call base class constructor, which will call createClusterSettings()
        # or readClusterSettings()
        super().__init__(
            clusterName,
            clusterType,
            zone,
            nodeStorage,
            nodeStorageOverrides,
            enable_fuse,
        )

    def supportedClusterTypes(self):
        return {"mesos"}

    def createClusterSettings(self):
        # All we need to do is read the Google credentials we need to provision
        # things
        self._readCredentials()

    def readClusterSettings(self):
        """
        Read the cluster settings from the instance, which should be the leader.
        See https://cloud.google.com/compute/docs/storing-retrieving-metadata for details about
        reading the metadata.
        """
        metadata_server = "http://metadata/computeMetadata/v1/instance/"
        metadata_flavor = {"Metadata-Flavor": "Google"}
        zone = web_session.get(metadata_server + "zone", headers=metadata_flavor).text
        self._zone = zone.split("/")[-1]

        project_metadata_server = "http://metadata/computeMetadata/v1/project/"
        self._projectId = web_session.get(
            project_metadata_server + "project-id", headers=metadata_flavor
        ).text

        # From a GCE instance, these values can be blank. Only the projectId is needed
        self._googleJson = ""
        self._clientEmail = ""

        self._tags = web_session.get(
            metadata_server + "description", headers=metadata_flavor
        ).text
        tags = json.loads(self._tags)
        self.clusterName = tags["clusterName"]
        self._gceDriver = self._getDriver()
        self._instanceGroup = self._gceDriver.ex_get_instancegroup(
            self.clusterName, zone=self._zone
        )

        leader = self.getLeader()
        self._leaderPrivateIP = leader.privateIP

        # The location of the Google credentials file on instances.
        self._credentialsPath = GoogleJobStore.nodeServiceAccountJson
        self._keyName = "core"  # key name leader users to communicate with works
        self._botoPath = (
            self.NODE_BOTO_PATH
        )  # boto credentials (used if reading an AWS bucket)

        # Let the base provisioner work out how to deploy duly authorized
        # workers for this leader.
        self._setLeaderWorkerAuthentication()

    def _readCredentials(self):
        """
        Get the credentials from the file specified by GOOGLE_APPLICATION_CREDENTIALS.
        """
        self._googleJson = os.getenv("GOOGLE_APPLICATION_CREDENTIALS")
        if not self._googleJson:
            raise RuntimeError("GOOGLE_APPLICATION_CREDENTIALS not set.")
        try:
            with open(self._googleJson) as jsonFile:
                self.googleConnectionParams = json.loads(jsonFile.read())
        except:
            raise RuntimeError(
                "GCEProvisioner: Could not parse the Google service account json file %s"
                % self._googleJson
            )

        self._projectId = self.googleConnectionParams["project_id"]
        self._clientEmail = self.googleConnectionParams["client_email"]
        self._credentialsPath = self._googleJson
        self._clearLeaderWorkerAuthentication()  # TODO: Why are we doing this?
        self._gceDriver = self._getDriver()

    def _write_file_to_cloud(self, key: str, contents: bytes) -> str:
        raise NotImplementedError(
            "The gceProvisioner doesn't support _write_file_to_cloud()."
        )

    def _get_user_data_limit(self) -> int:
        # See: https://cloud.google.com/compute/docs/metadata/setting-custom-metadata#limitations
        return human2bytes("256KB")

    def launchCluster(self, leaderNodeType, leaderStorage, owner, **kwargs):
        """
        In addition to the parameters inherited from the abstractProvisioner,
        the Google launchCluster takes the following parameters:
        keyName: The key used to communicate with instances
        botoPath: Boto credentials for reading an AWS jobStore (optional).
        network: a network (optional)
        vpcSubnet: A subnet (optional).
        use_private_ip: even though a public ip exists, ignore it (optional)
        """
        if "keyName" not in kwargs:
            raise RuntimeError("A keyPairName is required for the GCE provisioner.")
        self._keyName = kwargs["keyName"]
        if "botoPath" in kwargs:
            self._botoPath = kwargs["botoPath"]
        self._vpcSubnet = kwargs.get("vpcSubnet", None)
        self._network = kwargs.get("network", None)
        self._use_private_ip = kwargs.get("use_private_ip", None)

        # Throws an error if cluster exists
        self._instanceGroup = self._gceDriver.ex_create_instancegroup(
            self.clusterName, self._zone
        )
        logger.debug("Launching leader")

        # GCE doesn't have a dictionary tags field. The tags field is just a string list.
        # Therefore, dumping tags into the description.
        tags = {"Owner": self._keyName, "clusterName": self.clusterName}
        if "userTags" in kwargs:
            tags.update(kwargs["userTags"])
        self._tags = json.dumps(tags)

        metadata = {
            "items": [
                {"key": "user-data", "value": self._getIgnitionUserData("leader")}
            ]
        }
        imageType = "flatcar-stable"
        sa_scopes = [{"scopes": ["compute", "storage-full"]}]
        disk = {}
        disk["initializeParams"] = {
            "sourceImage": self.SOURCE_IMAGE,
            "diskSizeGb": leaderStorage,
        }
        disk.update({"boot": True, "autoDelete": True})
        name = "l" + str(uuid.uuid4())

        leader = self._gceDriver.create_node(
            name,
            leaderNodeType,
            imageType,
            location=self._zone,
            ex_service_accounts=sa_scopes,
            ex_metadata=compat_bytes_recursive(metadata),
            ex_network=self._network,
            ex_subnetwork=self._vpcSubnet,
            ex_disks_gce_struct=[compat_bytes_recursive(disk)],
            description=self._tags,
            ex_preemptible=False,
        )

        self._instanceGroup.add_instances([leader])
        self._leaderPrivateIP = leader.private_ips[0]  # needed if adding workers
        # self.subnetID = leader.subnet_id  # TODO: get subnetID

        # Wait for the appliance to start and inject credentials.
        leaderNode = Node(
            publicIP=leader.public_ips[0],
            privateIP=leader.private_ips[0],
            name=leader.name,
            launchTime=leader.created_at,
            nodeType=leader.size,
            preemptible=False,
            tags=self._tags,
            use_private_ip=self._use_private_ip,
        )
        leaderNode.waitForNode("toil_leader", keyName=self._keyName)
        leaderNode.copySshKeys(self._keyName)
        leaderNode.injectFile(
            self._credentialsPath, GoogleJobStore.nodeServiceAccountJson, "toil_leader"
        )
        if self._botoPath:
            leaderNode.injectFile(self._botoPath, self.NODE_BOTO_PATH, "toil_leader")
        # Download credentials
        self._setLeaderWorkerAuthentication(leaderNode)

        logger.debug("Launched leader")

    def getNodeShape(self, instance_type: str, preemptible=False) -> Shape:
        # TODO: read this value only once
        sizes = self._gceDriver.list_sizes(location=self._zone)
        sizes = [x for x in sizes if x.name == instance_type]
        assert len(sizes) == 1
        instanceType = sizes[0]

        disk = 0  # instanceType.disks * instanceType.disk_capacity * 2 ** 30
        if disk == 0:
            # This is an EBS-backed instance. We will use the root
            # volume, so add the amount of EBS storage requested forhe root volume
            disk = (
                self._nodeStorageOverrides.get(instance_type, self._nodeStorage) * 2**30
            )

        # Ram is in M.
        # Underestimate memory by 100M to prevent autoscaler from disagreeing with
        # mesos about whether a job can run on a particular node type
        memory = (instanceType.ram / 1000 - 0.1) * 2**30
        return Shape(
            wallTime=60 * 60,
            memory=memory,
            cores=instanceType.extra["guestCpus"],
            disk=disk,
            preemptible=preemptible,
        )

    @staticmethod
    def retryPredicate(e):
        """Not used by GCE"""
        return False

    def destroyCluster(self) -> None:
        """
        Try a few times to terminate all of the instances in the group.
        """
        logger.debug("Destroying cluster %s" % self.clusterName)
        instancesToTerminate = self._getNodesInCluster()
        attempts = 0
        while instancesToTerminate and attempts < 3:
            self._terminateInstances(instances=instancesToTerminate)
            instancesToTerminate = self._getNodesInCluster()
            attempts += 1

        # remove group
        instanceGroup = self._gceDriver.ex_get_instancegroup(
            self.clusterName, zone=self._zone
        )
        instanceGroup.destroy()

    def terminateNodes(self, nodes):
        if nodes:
            nodeNames = [n.name for n in nodes]
            instances = self._getNodesInCluster()
            instancesToKill = [i for i in instances if i.name in nodeNames]
            self._terminateInstances(instancesToKill)

    def addNodes(self, nodeTypes: set[str], numNodes, preemptible, spotBid=None) -> int:
        assert self._leaderPrivateIP

        # We don't support any balancing here so just pick one of the
        # equivalent node types
        node_type = next(iter(nodeTypes))

        # If keys are rsynced, then the mesos-slave needs to be started after the keys have been
        # transferred. The waitForKey.sh script loops on the new VM until it finds the keyPath file, then it starts the
        # mesos-slave. If there are multiple keys to be transferred, then the last one to be transferred must be
        # set to keyPath.
        keyPath = None
        botoExists = False
        if self._botoPath is not None and os.path.exists(self._botoPath):
            keyPath = self.NODE_BOTO_PATH
            botoExists = True
        elif self._sseKey:
            keyPath = self._sseKey

        if not preemptible:
            logger.debug("Launching %s non-preemptible nodes", numNodes)
        else:
            logger.debug("Launching %s preemptible nodes", numNodes)

        # kwargs["subnet_id"] = self.subnetID if self.subnetID else self._getClusterInstance(self.instanceMetaData).subnet_id
        userData = self._getIgnitionUserData("worker", keyPath, preemptible)
        metadata = {"items": [{"key": "user-data", "value": userData}]}
        imageType = "flatcar-stable"
        sa_scopes = [{"scopes": ["compute", "storage-full"]}]
        disk = {}
        disk["initializeParams"] = {
            "sourceImage": self.SOURCE_IMAGE,
            "diskSizeGb": self._nodeStorageOverrides.get(node_type, self._nodeStorage),
        }
        disk.update({"boot": True, "autoDelete": True})

        # TODO:
        #  - bug in gce.py for ex_create_multiple_nodes (erroneously, doesn't allow image and disk to specified)
        #  - ex_create_multiple_nodes is limited to 1000 nodes
        #    - use a different function
        #    - or write a loop over the rest of this function, with 1K nodes max on each iteration
        retries = 0
        workersCreated = 0
        # Try a few times to create the requested number of workers
        while numNodes - workersCreated > 0 and retries < 3:
            instancesLaunched = self.ex_create_multiple_nodes(
                "",
                node_type,
                imageType,
                numNodes - workersCreated,
                location=self._zone,
                ex_service_accounts=sa_scopes,
                ex_metadata=metadata,
                ex_disks_gce_struct=[disk],
                description=self._tags,
                ex_preemptible=preemptible,
            )
            failedWorkers = []
            for instance in instancesLaunched:
                if isinstance(instance, GCEFailedNode):
                    logger.error(
                        "Worker failed to launch with code %s. Error message: %s"
                        % (instance.code, instance.error)
                    )
                    continue

                node = Node(
                    publicIP=instance.public_ips[0],
                    privateIP=instance.private_ips[0],
                    name=instance.name,
                    launchTime=instance.created_at,
                    nodeType=instance.size,
                    preemptible=False,
                    tags=self._tags,
                    use_private_ip=self._use_private_ip,
                )  # FIXME: what should tags be set to?

                try:
                    self._injectWorkerFiles(node, botoExists)
                    logger.debug("Created worker %s" % node.publicIP)
                    self._instanceGroup.add_instances([instance])
                    workersCreated += 1
                except Exception as e:
                    logger.error(
                        f"Failed to configure worker {node.name}. Error message: {e}"
                    )
                    failedWorkers.append(instance)
            if failedWorkers:
                logger.error("Terminating %d failed workers" % len(failedWorkers))
                self._terminateInstances(failedWorkers)
            retries += 1

        logger.debug("Launched %d new instance(s)", numNodes)
        if numNodes != workersCreated:
            logger.error("Failed to launch %d worker(s)", numNodes - workersCreated)
        return workersCreated

    def getProvisionedWorkers(
        self, instance_type: Optional[str] = None, preemptible: Optional[bool] = None
    ):
        assert self._leaderPrivateIP
        entireCluster = self._getNodesInCluster(instance_type=instance_type)
        logger.debug("All nodes in cluster: %s", entireCluster)
        workerInstances = []
        for instance in entireCluster:
            if preemptible is not None:
                scheduling = instance.extra.get("scheduling")
                # If this field is not found in the extra meta-data, assume the node is not preemptible.
                if scheduling and scheduling.get("preemptible", False) != preemptible:
                    continue
            isWorker = True
            for ip in instance.private_ips:
                if ip == self._leaderPrivateIP:
                    isWorker = False
                    break  # don't include the leader
            if isWorker and instance.state == "running":
                workerInstances.append(instance)

        logger.debug("All workers found in cluster: %s", workerInstances)
        return [
            Node(
                publicIP=i.public_ips[0],
                privateIP=i.private_ips[0],
                name=i.name,
                launchTime=i.created_at,
                nodeType=i.size,
                preemptible=i.extra.get("scheduling", {}).get("preemptible", False),
                tags=None,
                use_private_ip=self._use_private_ip,
            )
            for i in workerInstances
        ]

    def getLeader(self):
        instances = self._getNodesInCluster()
        instances.sort(key=lambda x: x.created_at)
        try:
            leader = instances[0]  # assume leader was launched first
        except IndexError:
            raise NoSuchClusterException(self.clusterName)
        return Node(
            publicIP=leader.public_ips[0],
            privateIP=leader.private_ips[0],
            name=leader.name,
            launchTime=leader.created_at,
            nodeType=leader.size,
            preemptible=False,
            tags=None,
            use_private_ip=self._use_private_ip,
        )

    def _injectWorkerFiles(self, node, botoExists):
        """
        Set up the credentials on the worker.
        """
        node.waitForNode("toil_worker", keyName=self._keyName)
        node.copySshKeys(self._keyName)
        node.injectFile(
            self._credentialsPath, GoogleJobStore.nodeServiceAccountJson, "toil_worker"
        )
        if self._sseKey:
            node.injectFile(self._sseKey, self._sseKey, "toil_worker")
        if botoExists:
            node.injectFile(self._botoPath, self.NODE_BOTO_PATH, "toil_worker")

    def _getNodesInCluster(self, instance_type: Optional[str] = None):
        instanceGroup = self._gceDriver.ex_get_instancegroup(
            self.clusterName, zone=self._zone
        )
        instances = instanceGroup.list_instances()
        if instance_type:
            instances = [
                instance for instance in instances if instance.size == instance_type
            ]
        return instances

    def _getDriver(self):
        """Connect to GCE"""
        driverCls = get_driver(Provider.GCE)
        return driverCls(
            self._clientEmail,
            self._googleJson,
            project=self._projectId,
            datacenter=self._zone,
        )

    def _terminateInstances(self, instances):
        def worker(driver, instance):
            logger.debug("Terminating instance: %s", instance.name)
            driver.destroy_node(instance)

        threads = []
        for instance in instances:
            t = threading.Thread(target=worker, args=(self._gceDriver, instance))
            threads.append(t)
            t.start()

        logger.debug("... Waiting for instance(s) to shut down...")
        for t in threads:
            t.join()

    # MONKEY PATCH - This function was copied form libcloud to fix a bug.
    DEFAULT_TASK_COMPLETION_TIMEOUT = 180

    def ex_create_multiple_nodes(
        self,
        base_name,
        size,
        image,
        number,
        location=None,
        ex_network="default",
        ex_subnetwork=None,
        ex_tags=None,
        ex_metadata=None,
        ignore_errors=True,
        use_existing_disk=True,
        poll_interval=2,
        external_ip="ephemeral",
        ex_disk_type="pd-standard",
        ex_disk_auto_delete=True,
        ex_service_accounts=None,
        timeout=DEFAULT_TASK_COMPLETION_TIMEOUT,
        description=None,
        ex_can_ip_forward=None,
        ex_disks_gce_struct=None,
        ex_nic_gce_struct=None,
        ex_on_host_maintenance=None,
        ex_automatic_restart=None,
        ex_image_family=None,
        ex_preemptible=None,
    ):
        """
        Monkey patch to gce.py in libcloud to allow disk and images to be specified.
        Also changed name to a uuid below.
        The prefix 'wp' identifies preemptible nodes and 'wn' non-preemptible nodes.
        """
        # if image and ex_disks_gce_struct:
        #    raise ValueError("Cannot specify both 'image' and "
        #                     "'ex_disks_gce_struct'.")

        driver = self._getDriver()
        if image and ex_image_family:
            raise ValueError("Cannot specify both 'image' and " "'ex_image_family'")

        location = location or driver.zone
        if not hasattr(location, "name"):
            location = driver.ex_get_zone(location)
        if not hasattr(size, "name"):
            size = driver.ex_get_size(size, location)
        if not hasattr(ex_network, "name"):
            ex_network = driver.ex_get_network(ex_network)
        if ex_subnetwork and not hasattr(ex_subnetwork, "name"):
            ex_subnetwork = driver.ex_get_subnetwork(
                ex_subnetwork, region=driver._get_region_from_zone(location)
            )
        if ex_image_family:
            image = driver.ex_get_image_from_family(ex_image_family)
        if image and not hasattr(image, "name"):
            image = driver.ex_get_image(image)
        if not hasattr(ex_disk_type, "name"):
            ex_disk_type = driver.ex_get_disktype(ex_disk_type, zone=location)

        node_attrs = {
            "size": size,
            "image": image,
            "location": location,
            "network": ex_network,
            "subnetwork": ex_subnetwork,
            "tags": ex_tags,
            "metadata": ex_metadata,
            "ignore_errors": ignore_errors,
            "use_existing_disk": use_existing_disk,
            "external_ip": external_ip,
            "ex_disk_type": ex_disk_type,
            "ex_disk_auto_delete": ex_disk_auto_delete,
            "ex_service_accounts": ex_service_accounts,
            "description": description,
            "ex_can_ip_forward": ex_can_ip_forward,
            "ex_disks_gce_struct": ex_disks_gce_struct,
            "ex_nic_gce_struct": ex_nic_gce_struct,
            "ex_on_host_maintenance": ex_on_host_maintenance,
            "ex_automatic_restart": ex_automatic_restart,
            "ex_preemptible": ex_preemptible,
        }
        # List for holding the status information for disk/node creation.
        status_list = []

        for i in range(number):
            name = "wp" if ex_preemptible else "wn"
            name += str(uuid.uuid4())  # '%s-%03d' % (base_name, i)
            status = {"name": name, "node_response": None, "node": None}
            status_list.append(status)

        start_time = time.time()
        complete = False
        while not complete:
            if time.time() - start_time >= timeout:
                raise Exception(
                    "Timeout (%s sec) while waiting for multiple " "instances"
                )
            complete = True
            time.sleep(poll_interval)
            for status in status_list:
                # Create the node or check status if already in progress.
                if not status["node"]:
                    if not status["node_response"]:
                        driver._multi_create_node(status, node_attrs)
                    else:
                        driver._multi_check_node(status, node_attrs)
                # If any of the nodes have not been created (or failed) we are
                # not done yet.
                if not status["node"]:
                    complete = False

        # Return list of nodes
        node_list = []
        for status in status_list:
            node_list.append(status["node"])
        return node_list
