#!/usr/bin/env python3
#
# Copyright 2018 The Bazel Authors. 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
#
#    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.

from datetime import datetime
import json
import os
import queue
import re
import subprocess
import sys
import tempfile
import threading
import time

import gcloud
import gcloud_utils

DEBUG = False

LOCATION = 'europe-west1-d'

IMAGE_CREATION_VMS = {
    # Find the newest FreeBSD 11 image via:
    # gcloud compute images list --project freebsd-org-cloud-dev \
    #     --no-standard-images
    # ('buildkite-freebsd11',): {
    #     'source_image': 'https://www.googleapis.com/compute/v1/projects/freebsd-org-cloud-dev/global/images/freebsd-11-1-stable-amd64-2017-12-28',
    #     'scripts': [
    #         'setup-freebsd.sh',
    #         'install-buildkite-agent.sh'
    #     ]
    # },
    ('buildkite-ubuntu1404',): {
        'source_image_project': 'ubuntu-os-cloud',
        'source_image_family': 'ubuntu-1404-lts',
        'scripts': [
            'shell-utils.sh',
            'setup-ubuntu.sh',
            'install-azul-zulu.sh',
            'install-bazel.sh',
            'install-buildkite-agent.sh',
            'install-docker.sh',
            'install-nodejs.sh',
            'install-python36.sh',
            'install-android-sdk.sh',
            'shutdown.sh'
        ],
        'licenses': [
            'https://www.googleapis.com/compute/v1/projects/vm-options/global/licenses/enable-vmx'
        ]
    },
    ('buildkite-ubuntu1604', 'buildkite-trusted-ubuntu1604', 'buildkite-pipeline-ubuntu1604'): {
        'source_image_project': 'ubuntu-os-cloud',
        'source_image_family': 'ubuntu-1604-lts',
        'scripts': [
            'shell-utils.sh',
            'setup-ubuntu.sh',
            'install-azul-zulu.sh',
            'install-bazel.sh',
            'install-buildkite-agent.sh',
            'install-docker.sh',
            'install-nodejs.sh',
            'install-python36.sh',
            'install-android-sdk.sh',
            'shutdown.sh'
        ],
        'licenses': [
            'https://www.googleapis.com/compute/v1/projects/vm-options/global/licenses/enable-vmx'
        ]
    },
    ('buildkite-windows',): {
        'source_image_project': 'windows-cloud',
        'source_image_family': 'windows-1709-core',
        'scripts': [
            'setup-windows-manual.ps1'
        ]
    }
}

WORK_QUEUE = queue.Queue()


def run(args, **kwargs):
    return subprocess.run(args, **kwargs)


def merge_setup_scripts(instance_name, scripts):
    newline = '\r\n' if 'windows' in instance_name else None
    # Merge all setup scripts into one.
    merged_script_path = tempfile.mkstemp()[1]
    with open(merged_script_path, 'w', newline=newline) as merged_script_file:
        for script in scripts:
            with open(script, 'r') as script_file:
                merged_script_file.write(script_file.read() + '\n')
    return merged_script_path


def create_instance(instance_name, params):
    merged_script_path = merge_setup_scripts(instance_name, params['scripts'])
    try:
        if 'windows' in instance_name:
            startup_script = 'windows-startup-script-ps1=' + merged_script_path
        else:
            startup_script = 'startup-script=' + merged_script_path

        if 'source_image' in params:
            image = {
                'image': params['source_image']
            }
        else:
            image = {
                'image-project': params['source_image_project'],
                'image-family': params['source_image_family']
            }

        gcloud.create_instance(
            instance_name,
            zone=LOCATION,
            machine_type='n1-standard-8',
            network='buildkite',
            metadata_from_file=startup_script,
            min_cpu_platform='Intel Skylake',
            boot_disk_type='pd-ssd',
            boot_disk_size='50GB',
            **image)
    finally:
        os.remove(merged_script_path)


# https://stackoverflow.com/a/25802742
def write_to_clipboard(output):
    process = subprocess.Popen(
        'pbcopy', env={'LANG': 'en_US.UTF-8'}, stdin=subprocess.PIPE)
    process.communicate(output.encode('utf-8'))


def print_windows_instructions(instance_name):
    tail_start = gcloud_utils.tail_serial_console(instance_name, zone=LOCATION, until='Finished running startup scripts')

    pw = json.loads(gcloud.reset_windows_password(instance_name, format='json', zone=LOCATION).stdout)
    rdp_file = tempfile.mkstemp(suffix='.rdp')[1]
    with open(rdp_file, 'w') as f:
        f.write('full address:s:' + pw['ip_address'] + '\n')
        f.write('username:s:' + pw['username'] + '\n')
    subprocess.run(['open', rdp_file])
    write_to_clipboard(pw['password'])
    with gcloud.PRINT_LOCK:
        print('Use this password to connect to the Windows VM: ' + pw['password'])
        print('Please run the setup script C:\\setup.ps1 once you\'re logged in.')

    # Wait until the VM reboots once, then open RDP again.
    tail_start = gcloud_utils.tail_serial_console(
        instance_name, zone=LOCATION, start=tail_start, until='Finished running startup scripts')
    print('Connecting via RDP a second time to finish the setup...')
    write_to_clipboard(pw['password'])
    run(['open', rdp_file])
    return tail_start


def workflow(name, params):
    instance_name = "%s-image-%s" % (name, int(datetime.now().timestamp()))
    try:
        # Create the VM.
        create_instance(instance_name, params)

        # Wait for the VM to become ready.
        gcloud_utils.wait_for_instance(instance_name, zone=LOCATION, status='RUNNING')

        if 'windows' in instance_name:
            # Wait for VM to be ready, then print setup instructions.
            tail_start = print_windows_instructions(instance_name)
            # Continue printing the serial console until the VM shuts down.
            gcloud_utils.tail_serial_console(instance_name, zone=LOCATION, start=tail_start)
        else:
            # Continuously print the serial console.
            gcloud_utils.tail_serial_console(instance_name, zone=LOCATION)

        # Wait for the VM to completely shutdown.
        gcloud_utils.wait_for_instance(instance_name, zone=LOCATION, status='TERMINATED')

        # Create a new image from our VM.
        gcloud.create_image(
            instance_name,
            family=name,
            source_disk=instance_name,
            source_disk_zone=LOCATION,
            licenses=params.get('licenses', []))
    finally:
        gcloud.delete_instance(instance_name)


def worker():
    while True:
        item = WORK_QUEUE.get()
        if not item:
            break
        try:
            workflow(**item)
        finally:
            WORK_QUEUE.task_done()


def main(argv=None):
    if argv is None:
        argv = sys.argv[1:]

    # Put VM creation instructions into the work queue.
    for names, params in IMAGE_CREATION_VMS.items():
        for name in names:
            if argv and name not in argv:
                continue
            WORK_QUEUE.put({
                'name': name,
                'params': params
            })

    # Spawn worker threads that will create the VMs.
    threads = []
    for _ in range(WORK_QUEUE.qsize()):
        t = threading.Thread(target=worker)
        t.start()
        threads.append(t)

    # Wait for all VMs to be created.
    WORK_QUEUE.join()

    # Signal worker threads to exit.
    for _ in range(len(threads)):
        WORK_QUEUE.put(None)

    # Wait for worker threads to exit.
    for t in threads:
        t.join()

    return 0


if __name__ == '__main__':
    sys.exit(main())