import atexit
import logging
import os

from parsl.jobs.states import JobState, JobStatus
from parsl.launchers import SingleNodeLauncher

logger = logging.getLogger(__name__)

try:
    import googleapiclient.discovery

except ImportError:
    _google_enabled = False
else:
    _google_enabled = True

translate_table = {
    'PENDING': JobState.PENDING,
    'PROVISIONING': JobState.PENDING,
    "STAGING": JobState.PENDING,
    'RUNNING': JobState.RUNNING,
    'DONE': JobState.COMPLETED,
    'STOPPING': JobState.COMPLETED,
    'STOPPED': JobState.COMPLETED,
    'TERMINATED': JobState.COMPLETED,
    'SUSPENDING': JobState.COMPLETED,
    'SUSPENDED': JobState.COMPLETED,
}


class GoogleCloudProvider():
    """A provider for using resources from the Google Compute Engine.

    Parameters
    ----------
    project_id : str
        Project ID from Google compute engine.
    key_file : str
        Path to authorization private key json file. This is required for auth.
        A new one can be generated here: https://console.cloud.google.com/apis/credentials
    region : str
        Region in which to start instances
    os_project : str
        OS project code for Google compute engine.
    os_family : str
        OS family to request.
    google_version : str
        Google compute engine version to use. Possibilies include 'v1' (default) or 'beta'.
    instance_type: str
        'n1-standard-1',
    init_blocks : int
        Number of blocks to provision immediately. Default is 1.
    min_blocks : int
        Minimum number of blocks to maintain. Default is 0.
    max_blocks : int
        Maximum number of blocks to maintain. Default is 10.
    parallelism : float
        Ratio of provisioned task slots to active tasks. A parallelism value of 1 represents aggressive
        scaling where as many resources as possible are used; parallelism close to 0 represents
        the opposite situation in which as few resources as possible (i.e., min_blocks) are used.

     """

    def __init__(self,
                 project_id,
                 key_file,
                 region,
                 os_project,
                 os_family,
                 google_version='v1',
                 instance_type='n1-standard-1',
                 init_blocks=1,
                 min_blocks=0,
                 max_blocks=10,
                 launcher=SingleNodeLauncher(),
                 parallelism=1):
        self.project_id = project_id
        os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = key_file
        self.zone = self.get_zone(region)
        self.os_project = os_project
        self.os_family = os_family
        self.label = 'google_cloud'
        self.client = googleapiclient.discovery.build('compute', google_version)
        self.instance_type = instance_type
        self.init_blocks = init_blocks
        self.min_blocks = min_blocks
        self.max_blocks = max_blocks
        self.parallelism = parallelism
        self.num_instances = 0
        self.launcher = launcher
        self.script_dir = None

        # Dictionary that keeps track of jobs, keyed on job_id
        self.resources = {}
        atexit.register(self.bye)

    def submit(self, command, tasks_per_node, job_name="parsl.gcs"):
        ''' The submit method takes the command string to be executed upon
        instantiation of a resource most often to start a pilot.

        Args :
             - command (str) : The bash command string to be executed.
             - tasks_per_node (int) : command invocations to be launched per node

        KWargs:
             - job_name (str) : Human friendly name to be assigned to the job request

        Returns:
             - A job identifier, this could be an integer, string etc

        Raises:
             - ExecutionProviderException or its subclasses
        '''
        wrapped_cmd = self.launcher(command,
                                    tasks_per_node,
                                    1,
                                    self.script_dir)

        instance, name = self.create_instance(command=wrapped_cmd)
        self.resources[name] = {"job_id": name, "status": JobStatus(translate_table[instance['status']])}
        return name

    def status(self, job_ids):
        ''' Get the status of a list of jobs identified by the job identifiers
        returned from the submit request.

        Args:
             - job_ids (list) : A list of job identifiers

        Returns:
             - A list of JobStatus objects corresponding to each job_id in the job_ids list.

        Raises:
             - ExecutionProviderException or its subclasses

        '''
        statuses = []
        for job_id in job_ids:
            instance = self.client.instances().get(instance=job_id, project=self.project_id, zone=self.zone).execute()
            job_status = JobStatus(translate_table[instance['status']])
            self.resources[job_id]['status'] = job_status
            statuses.append(job_status)
        return statuses

    def cancel(self, job_ids):
        ''' Cancels the resources identified by the job_ids provided by the user.

        Args:
             - job_ids (list): A list of job identifiers

        Returns:
             - A list of status from cancelling the job which can be True, False

        Raises:
             - ExecutionProviderException or its subclasses
        '''
        statuses = []
        for job_id in job_ids:
            try:
                self.delete_instance(job_id)
                statuses.append(True)
            except Exception:
                statuses.append(False)
        return statuses

    def bye(self):
        self.cancel([i for i in list(self.resources)])

    def create_instance(self, command=""):
        name = "parslauto{}".format(self.num_instances)
        self.num_instances += 1
        compute = self.client
        project = self.project_id
        image_response = compute.images().getFromFamily(
            project=self.os_project, family=self.os_family).execute()
        source_disk_image = image_response['selfLink']

        # Configure the machine
        machine_type = "zones/{}/machineTypes/{}".format(self.zone, self.instance_type)
        startup_script = command

        config = {
            'name': name,
            'machineType': machine_type,

            # Specify the boot disk and the image to use as a source.
            'disks': [{
                'boot': True,
                'autoDelete': True,
                'initializeParams': {
                    'sourceImage': source_disk_image,
                }
            }],
            'networkInterfaces': [{
                'network': 'global/networks/default',
                'accessConfigs': [{
                    'type': 'ONE_TO_ONE_NAT',
                    'name': 'External NAT'
                }]
            }],
            'serviceAccounts': [{
                'email':
                'default',
                'scopes': [
                    'https://www.googleapis.com/auth/devstorage.read_write',
                    'https://www.googleapis.com/auth/logging.write'
                ]
            }],
            'metadata': {
                'items': [{
                    # Startup script is automatically executed by the
                    # instance upon startup.
                    'key': 'startup-script',
                    'value': startup_script
                }]
            }
        }

        return compute.instances().insert(project=project, zone=self.zone, body=config).execute(), name

    def get_zone(self, region):
        res = self.client.zones().list(project=self.project_id).execute()
        for zone in res['items']:
            if region in zone['name'] and zone['status'] == "UP":
                return zone["name"]

    def delete_instance(self, name):

        compute = self.client
        project = self.project_id
        zone = self.zone

        return compute.instances().delete(project=project, zone=zone, instance=name).execute()

    @property
    def status_polling_interval(self):
        return 60
