from google.cloud import storage
import googleapiclient.discovery
import shutil
import os
import time
from .function import port_open, post_slack
from .libraries import extra_libraries, important_libraries
from herpetologist import check_type
import subprocess
import cloudpickle
from typing import Callable, List


additional_command = [
    'gsutil cp gs://general-bucket/dask.zip dask.zip',
    'unzip dask.zip',
    'worker_size=1 name=a project=a zone=a expired=99999 docker-compose -f docker-compose.yaml up --build',
]
dask_network = {
    'allowed': [{'IPProtocol': 'tcp', 'ports': ['8787', '8786']}],
    'description': '',
    'direction': 'INGRESS',
    'kind': 'compute#firewall',
    'name': 'dask-network',
    'priority': 1000.0,
    'sourceRanges': ['0.0.0.0/0'],
    'targetTags': ['dask'],
}


@check_type
def build_image(
    project: str,
    zone: str,
    bucket_name: str,
    instance_name: str,
    image_name: str,
    project_vm: str = 'ubuntu-os-cloud',
    family_vm: str = 'ubuntu-1804-lts',
    storage_image: str = 'asia-southeast1',
    webhook_function: Callable = post_slack,
    validate_webhook: bool = True,
    additional_libraries: List[str] = extra_libraries,
    install_bash: str = None,
    **kwargs
):
    """
    Parameters
    ----------

    project: str
        project id
    zone: str
    bucket_name: str
        bucket name to upload dask code, can be private.
    image_name: str
        image name for dask bootloader.
    project_vm: str, (default='ubuntu-os-cloud')
        project name for vm. 
    family_vm: str, (default='ubuntu-1804-lts')
        family name for vm.
    storage_image: str, (default='asia-southeast1')
        storage location for dask image.
    webhook_function: Callable, (default=post_slack)
        Callable function to send alert during gracefully delete, default is post_slack.
    validate_webhook: bool, (default=True)
        if True, will validate `webhook_function`. 
        Not suggest to set it as False because this webhook_function will use during gracefully delete.
    additional_libraries: List[str], (default=extra_libraries). 
        add more libraries from PYPI. This is necessary if want dask cluster able to necessary libraries.
    **kwargs:
        Keyword arguments to pass to webhook_function.
    """

    if webhook_function.__name__ == 'post_slack':

        def nested_post(msg):
            return webhook_function(msg, **kwargs)

    else:

        def nested_post(msg):
            return webhook_function(msg)

    if validate_webhook:
        if nested_post('Testing from ondemand-dask') != 200:
            raise Exception('`webhook_function` must returned 200.')

    compute = googleapiclient.discovery.build('compute', 'v1')
    storage_client = storage.Client()
    bucket = storage_client.bucket(bucket_name)

    this_dir = os.path.dirname(__file__)
    pkl = os.path.join(this_dir, 'image', 'dask', 'post.pkl')
    with open(pkl, 'wb') as fopen:
        cloudpickle.dump(nested_post, fopen)

    reqs = important_libraries + additional_libraries
    reqs = list(set(reqs))

    req = os.path.join(this_dir, 'image', 'dask', 'requirements.txt')
    with open(req, 'w') as fopen:
        fopen.write('\n'.join(reqs))

    image = os.path.join(this_dir, 'image')
    shutil.make_archive('dask', 'zip', image)
    blob = bucket.blob('dask.zip')
    blob.upload_from_filename('dask.zip')
    os.remove('dask.zip')
    image_response = (
        compute.images()
        .getFromFamily(project = project_vm, family = family_vm)
        .execute()
    )
    source_disk_image = image_response['selfLink']

    try:
        print('Creating `dask-network` firewall rule.')
        compute.firewalls().insert(
            project = project, body = dask_network
        ).execute()
        print('Done.')
    except:
        print('`dask-network` exists.')

    machine_type = 'zones/%s/machineTypes/n1-standard-1' % zone

    if install_bash is None:
        install_bash = 'install.sh'
        install_bash = os.path.join(this_dir, install_bash)

    startup_script = open(install_bash).read()
    startup_script = '\n'.join(
        startup_script.split('\n') + additional_command
    ).replace('general-bucket', bucket_name)

    config = {
        'name': instance_name,
        'tags': {'items': ['dask']},
        'machineType': machine_type,
        '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',
                    'https://www.googleapis.com/auth/compute',
                ],
            }
        ],
        'metadata': {
            'items': [
                {'key': 'startup-script', 'value': startup_script},
                {'key': 'bucket', 'value': bucket_name},
            ]
        },
    }

    operation = (
        compute.instances()
        .insert(project = project, zone = zone, body = config)
        .execute()
    )

    print('Waiting instance `%s` to run.' % (instance_name))

    while True:
        result = (
            compute.zoneOperations()
            .get(project = project, zone = zone, operation = operation['name'])
            .execute()
        )

        if result['status'] == 'DONE':
            if 'error' in result:
                raise Exception(result['error'])
            else:
                print('Done.')
            break

        time.sleep(1)

    print('Waiting IP Address to check health.')

    while True:
        result = (
            compute.instances().list(project = project, zone = zone).execute()
        )
        results = result['items'] if 'items' in result else None
        dask = [r for r in results if r['name'] == instance_name]
        if len(dask) > 0:
            dask = dask[0]
            ip_address = dask['networkInterfaces'][0]['accessConfigs'][0][
                'natIP'
            ]
            print('Got it, Public IP: %s' % (ip_address))
            break

        time.sleep(2)

    print('Waiting Dask cluster to run.')
    while True:
        if port_open(ip_address, 8786) and port_open(ip_address, 8787):
            print('Done.')
            break
        time.sleep(5)

    compute = googleapiclient.discovery.build('compute', 'v1')
    print('Deleting image `%s` if exists.' % (image_name))
    try:
        compute.images().delete(project = project, image = image_name).execute()
        print('Done.')
    except:
        pass

    # give a rest to gcp API before build the image.
    time.sleep(20)

    print('Building image `%s`.' % (image_name))
    try:
        o = subprocess.check_output(
            [
                'gcloud',
                'compute',
                'images',
                'create',
                image_name,
                '--source-disk',
                instance_name,
                '--source-disk-zone',
                zone,
                '--family',
                family_vm,
                '--storage-location',
                storage_image,
                '--force',
            ],
            stderr = subprocess.STDOUT,
        )
        print('Done.')
    except subprocess.CalledProcessError as e:
        print(e.output.decode('utf-8'))
        raise

    print('Deleting instance `%s`.' % (instance_name))
    compute = googleapiclient.discovery.build('compute', 'v1')
    compute.instances().delete(
        project = project, zone = zone, instance = instance_name
    ).execute()
    print('Done.')
    return True
