#!/usr/bin/env python
#
# Copyright 2015 Codethink Ltd
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; version 2 of the License.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.

'''tester

Test system images.

'''

import os
import sys
import pipes
import socket
import time
import uuid
import yaml
import shutil
import urllib2

import cliapp

from novaclient import client

class Host(object):
    def __init__(self):
        self.type = None
        self.systems = []

    def get_images_for_host(self):
        images_for_host = []
        for system in self.systems:
            images_for_host.append(system.image)
        return list(set(images_for_host))

    def upload_images(self, files):
        print('Dummy image upload')

    def delete_images(self, files):
        print('Dummy image remove')

    def deploy_systems(self, files):
        print('Dummy deploy systems')

    def cleanup_systems(self, files):
        print('Dummy cleanup systems')

class OpenstackHost(Host):
    def __init__(self, name, host_type):
        self.name = name
        self.type = host_type
        self.systems = []
        self.env_vars = {}
        print(str(self) + 'New host: ' + name)
        self.load_host_config()

    def __str__(self):
     return '[' + self.type + '|' + self.name + '] '

    def load_host_config(self):
        with open('openstack.setup', 'r') as stream:
            data = yaml.load(stream)
            for host in data:
                if host['host'] != self.name:
                    continue
                self.env_vars = { k: v for d in host['variables']
                                       for k, v in d.items() }
                print(str(self) + 'Loaded config:')
                # TODO: Is it only password we want to hide?
                for var in self.env_vars:
                    if var in ['OS_PASSWORD', 'OS_USERNAME', 'OS_TENANT_NAME']:
                        print('  ' + var + ': <hidden>')
                    else:
                        print('  ' + var + ': ' + self.env_vars[var])
                break
            else:
                # Not a fatal error, since we might be used for single openstack
                # host, with environment already set up.
                print(str(self) + 'WARNING: No config found')

    def set_env_vars(self):
        for var in self.env_vars:
            os.environ[var] = self.env_vars[var]

    def remove_env_vars(self):
        for var in self.env_vars:
            del os.environ[var]

    def add_system(self, name, image, flavour, net_id):
        system = OpenstackSystem(name, image, flavour, net_id)
        self.systems.append(system)
        return system

    def upload_images(self, files):
        images = self.get_images_for_host()
        self.set_env_vars()
        for image in images:
            print(str(self) + 'Uploading image: ' + files[image])
            args = ['glance', 'image-create', '--progress',
                    '--name', files[image],
                    '--disk-format', 'raw',
                    '--container-format', 'bare',
                    '--file', files[image]]
            output = cliapp.runcmd(args)
            print(output)
        self.remove_env_vars()

    def delete_images(self, files):
        images = self.get_images_for_host()
        self.set_env_vars()
        for image in images:
            print(str(self) + 'Removing image: ' + files[image])
            cliapp.runcmd(['nova', 'image-delete', files[image]])
        self.remove_env_vars()

    def deploy_systems(self, files):
        self.set_env_vars()
        for system in self.systems:
            print(str(self) + 'Deploying system: ' + system.name)
            system.deploy(files)
        self.remove_env_vars()

    def cleanup_systems(self):
        self.set_env_vars()
        print(str(self) + 'Terminating systems:')
        for system in self.systems:
            system.terminate()
        self.remove_env_vars()

class System(object):
    def __init__(self):
        self.name = None
        self.id = None
        self.instance = None
        self.ip_addr = None

    def __str__(self):
     return '[' + self.name + '] '

    def ssh_host(self, user):
        # TODO: Stop assuming we ssh into test instances as root
        return '{user}@{host}'.format(host=self.ip_addr, user=user)

    def runcmd(self, argv, user='root', chdir='.', **kwargs):
        print(str(self) + 'Running command as ' + user + ': ' + argv[-1])
        ssh_cmd = ['ssh', '-o', 'StrictHostKeyChecking=no',
                   '-o', 'BatchMode=yes',
                   '-o', 'UserKnownHostsFile=/dev/null', self.ssh_host(user)]
        cmd = ['sh', '-c', 'cd "$1" && shift && exec "$@"', '-', chdir]
        cmd += argv
        ssh_cmd.append(' '.join(map(pipes.quote, cmd)))
        return cliapp.runcmd(ssh_cmd, **kwargs)

    def wait_for_ssh(self, timeout=180):
        """Wait until the deployed VM is responding via SSH"""
        start_time = time.time()
        while True:
            try:
                self.runcmd(['true'], stdin=None, stdout=None, stderr=None)
                print(str(self) + 'SSH connection established.')
                return
            except cliapp.AppException:
                # TODO: Stop assuming the ssh part of the command is what failed
                if time.time() > start_time + timeout:
                    raise TimeoutError("%s sshd did not start after %i seconds"
                                       % (self.ip_address, timeout))
                time.sleep(0.5)

class OpenstackSystem(System):
    def __init__(self, name, image, flavour, net_id):
        self.name = name
        self.image = image
        self.flavour = flavour
        self.net_id = net_id
        self.id = str(uuid.uuid4()) + '-' + name
        self.instance = None
        self.ip_addr = None

    def deploy(self, files):
        print(str(self) + 'Launching instance with image ' + files[self.image])

        # Get a novaclient object
        nc = client.Client(2,
                           os.environ['OS_USERNAME'],
                           os.environ['OS_PASSWORD'],
                           os.environ['OS_TENANT_NAME'],
                           os.environ['OS_AUTH_URL'])

        args = ['nova', 'boot',
                '--flavor', self.flavour,
                '--image', files[self.image],
                '--nic', 'net-id=' + self.net_id,
                self.id]
        output = cliapp.runcmd(args)

        # Print nova boot output, with adminPass line removed
        output_lines = output.split('\n')
        for line in output_lines:
            if line.find('adminPass') != -1:
                password_line = line
        output_lines.remove(password_line)
        output = '\n'.join(output_lines)
        print output

        # Sleep for a bit, or nova explodes when trying to assign IP address
        time.sleep(20)

        print(str(self) + 'Requesting a floating IP address')

        for _ in xrange(5):
            ip_list = nc.floating_ips.list()
            free_ip = None
            for ip in ip_list:
                if ip.instance_id == None:
                    free_ip = ip
                    break
            else:
                free_ip = nc.floating_ips.create(
                              nc.floating_ip_pools.list()[0].name)
            if free_ip != None:
                instance = nc.servers.find(name=self.id)
                # TODO: switch back to cli tool, as python
                #       approach gave error.
                instance.add_floating_ip(free_ip)
                self.ip_addr = free_ip.ip
                break
        else:
            self.terminate()
            raise cliapp.AppException('Could not get a floating IP')

        print(str(self) + 'Obtained IP address: ' + self.ip_addr)

    def terminate(self):
        print(str(self) + 'Terminating instance')
        cliapp.runcmd(['nova', 'delete', self.id])

class OpenstackYAML(yaml.YAMLObject):
    yaml_tag = u'!openstack'
    def __init__(self, flavour, host, net_id):
        self.flavour = flavour
        self.net_id = net_id
        self.host = host
    def __repr__(self):
        return "%s(flavour=%r, net-id=%r, host=%r)" % (
            self.__class__.__name__, self.flavour, self.net_id, self.host)

def download_image(url):
    # TODO: Improve progress indication.
    path = str(uuid.uuid4()) + '-' + url.rpartition('/')[2]
    print('[Local] Fetching: ' + url)

    file_to_fetch = urllib2.urlopen(url)
    total_file_size = int(file_to_fetch.info()['Content-Length'])

    copied_size = 0
    last_percentaje_shown = -1
    CHUNK = 16 * 1024
    with open(path, 'wb') as f:
        print('[Local] Progress:')
        while True:
            chunk = file_to_fetch.read(CHUNK)
            if not chunk: break
            f.write(chunk)
            copied_size += len(chunk)
            percentaje = (copied_size * 100) / total_file_size
            if last_percentaje_shown < percentaje:
                last_percentaje_shown = percentaje
                print ("  Copied %s %% - %s of %s bytes" % (percentaje,
                                                            copied_size,
                                                            total_file_size))

    print('\n[Local] Saved as: ' + path)
    return path

class ReleaseApp(cliapp.Application):

    """Cliapp application which handles automatic builds and tests"""

    def add_settings(self):
        """Add the command line options needed"""
        group_main = 'Program Options'
        self.settings.string(['setup'],
                             'Test setup description file',
                             default=None,
                             group=group_main)

    def add_openstack_system(self, image_url, system_info):

        system_name = system_info['name']
        system_platform = system_info['platform']
        if system_name in self.systems.keys():
            print('WARNING: Ignoring duplicate system: ' + system_name)
            return

        # Create host, if needed
        host_type = 'openstack'
        host_name = system_platform.host
        host_id = 'openstack-' + host_name

        if host_id not in self.hosts.keys():
            host = OpenstackHost(host_name, host_type)
            self.hosts[host_id] = host

        # Create system
        host = self.hosts[host_id]
        system = host.add_system(system_name, image_url,
                                 system_platform.flavour,
                                 system_platform.net_id)
        self.systems[system_name] = system

    def load_setup(self, path):
        self.systems = {}
        self.images = {}
        self.hosts = {}
        with open(path, 'r') as stream:
            data = yaml.load(stream)
            for d in data:
                image_url = d['image']
                self.images[image_url] = download_image(image_url)
                for s in d['systems']:
                    if (type(s['platform']).__name__ == 'OpenstackYAML'):
                        self.add_openstack_system(image_url, s)

    def load_tests(self, tests):
        self.tests = {}

        for test in tests:
            print('[Local] Loading test set: ' + test)
            with open(test, 'r') as stream:
                data = yaml.load(stream)

            if data == None:
                continue
            if not isinstance(data, list):
                print('Bad test: no array of steps')
            for step in data:
                if 'systems' not in step:
                    print('Bad test step: no systems')
                    continue
                if not isinstance(step['systems'], list):
                    print('Bad test step: systems should be an array')
                    continue
                if 'commands' not in step:
                    print('Bad test step: no commands')
                    continue
                if not isinstance(step['commands'], list):
                    print('Bad test step: commands should be an array')
                    continue
                self.tests[test] = data

    def start(self):
        print('[Local] Upload images to test hosts:')
        for host in self.hosts:
            self.hosts[host].upload_images(self.images)

        print('[Local] Deploy systems on test hosts:')
        for host in self.hosts:
            self.hosts[host].deploy_systems(self.images)

        print('[Local] Check SSH connection to systems:')
        for system in self.systems:
            self.systems[system].wait_for_ssh()

        print('[Local] Building IP address prefix:')
        ip_addrs = {}
        for system in self.systems:
            sanitised = 'IP__' + system.replace('-','_').upper()
            ip_addrs[sanitised] = self.systems[system].ip_addr
        prefix = ''
        for entry in ip_addrs:
            prefix = prefix + entry + '=' + ip_addrs[entry] + ' '
        print('    ' + prefix)

        for test in self.tests:
            print('[Local] Running test: ' + test)
            for step in self.tests[test]:
                for cmd in step['commands']:
                    for system in step['systems']:
                        # TODO: Pass prefix, so systems can talk to each other

                        try:
                            user = step['user']
                        except KeyError:
                            user = 'root'

                        try:
                            s = self.systems[system]
                        except KeyError:
                            print "WARNING: System " + system + " not defined, ignoring test"
                        else:
                            s.runcmd(['sh', '-c', cmd], user, stdout=sys.stdout, stderr=sys.stderr)

    def clean_up(self):
        print('[Local] Clean up systems on remote hosts:')
        for host in self.hosts:
            self.hosts[host].cleanup_systems()

        # Sleep for a bit, or nova silently fails to delete the image
        # TODO: Test whether the image has been deleted in a retry loop.
        time.sleep(20)

        print('[Local] Remove images from remote hosts:')
        for host in self.hosts:
            self.hosts[host].delete_images(self.images)

        print('[Local] Delete local images:')
        for url in self.images:
             path = self.images[url]
             print('[Local] Deleting: ' + path)
             os.remove(path)

    def process_args(self, args):
        """Process the command line args and kick off the tests"""
        self.settings.require('setup')

        if len(args) == 0:
            print('WARNING: No tests specified.')

        self.load_tests(args)
        try:
            self.load_setup(self.settings['setup'])
            self.start()
        finally:
            self.clean_up()

if __name__ == '__main__':
    ReleaseApp().run()