#!/usr/bin/python
# Copyright (c) 2014 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Orchestrate virtual machines to setup a toy instance of the lab for testing.

This module is meant to help create a closed loop development flow for members
of the lab team which looks something like this:
                    ______________
                   |              |
                   |gs vm resistry|<+
                   |______________| |
                          |         |
                          v         |
        New change -> puppylab -> New core_cluster box
                          |
         Vagrantfile specifies cluster settings
         _________________|____________________
        |                                      |
        |  puppet provisions core_cluster box  |
        |______________________________________|
                |          | ........... |
                v          v             v
              master     shard1       shardn
             |     |     |     |      |     |
            mysql  afe  tko heartbt   tko heartbt
             |     |     |     |      |     |
host ports  8001  8002  8001  8002    8001  8002
        [host ports liable to autocorrect as needed]

This module can work with any vm hosting service/provider as long as they
adhere to the vagrant interface. VirtualBox is the only implementation so
far, though GCE will be an ideal candidate.

Class spec:
* VagrantProvisioner: Provision boxes per a VagrantFile.
    * VirtualBoxProvisioner: Generate a Virtualbox VagrantFile.
* CoreVM: Manage individual core_cluster vms.
* ClusterManager: Spin up cluster.

Usage: clusterctl --admin-repo /usr/local/autotest/chromeos-internal
"""

import argparse
import logging
import os
import sys

import common
from autotest_lib.puppylab import lab_manifest
from autotest_lib.puppylab import vm_manager
from autotest_lib.site_utils.lib import infra


# TODO: Enable multiple shards via command line args.
NUM_SHARDS = 1
SHADOW_PATH = '/usr/local/autotest/shadow_config.ini'


class ConfigError(Exception):
    """Raised if one of the vms in the cluster is misconfigured."""


class CoreVM(object):
    """Interface to create and manage a core_cluster vm image.

    A core_cluster vm image has base packages shared by all server roles.
    """
    _core_vm_name = 'chromeos_lab_core_cluster'
    _core_image_source = 'gs://vms/%s.box' % _core_vm_name
    _core_image_name = '%s.box' % _core_vm_name
    _core_image_destination = os.path.join(
            vm_manager.VAGRANT_DIR, _core_image_name)

    # TODO: Preperation is currently by hand. Use the provisioner to
    # create a box of name '_core_image_name', with the CoreClusterTemplate
    # in the VAGRANT_DIR if you wish to prepare a new vm. You can achieve
    # this by:
    # * Copying the CoreClusterTemplate to a Vagrantfile and replacing the
    #   modulepath with the path to you chromeos-admin/puppet directory.
    # * Calling `vagrant up` in the directory with this vagrant file.
    # * When it's done, calling vagrant package.
    # This should produce a package.box in the same dir.

    def __init__(self, provisioner):
        self.provisioner = provisioner


    def setup_core_box(self):
        """Setup a core cluster vm.

        Download a core_cluster image if one isn't present on disk and
        register it with vagrant.
        """
        if not os.path.exists(self._core_image_destination):
            infra.execute_command(
                'localhost', 'gsutil cp %s %s' %
                (self._core_image_source, self._core_image_destination))
        self.provisioner.register_box(
                self._core_image_destination, self._core_vm_name)


    def teardown_core_box(self):
        """Teardown a core cluster vm."""
        # TODO: delete the box file.
        self.provisioner.unregister_box(self._core_vm_name)


class ClusterManager(object):
    """Interface to spin up a cluster of CoreVMs.

    This class manages all the details between creating a core_cluster image
    and running tests on a full fledged cluster.
    """

    def _register_shards(self, num_shards):
        """Register num_shards worth of shard info.

        This includes the name, port address and board of the new shard. This
        information is piped through to each vm, so the cluster manager is
        actually in control of all the shards in the cluster and can address
        them by name.

        Consider a shard, shard1, assigned to board stumpy:
            * You will be able to ssh into it with 'vagrant ssh stumpyshard'.
            * The afe for the shard will be running on a incrementally
              designated port starting from shards_base_port.
            * The afe port of the shard is piped through to the shadow_config.
              This is required for 2 reasons:
                # `cli/atest shard add` should use this name, because it is
                  the name the shard-client will use to request jobs.
                # the master afe should show links to the shard using this name.

        @param num_shards: The number of shards we wish to add to the cluster.
        """
        self.vagrantfile_shard_args = {}
        self.shard_board_map = {}
        self.vagrant_shard_names = []

        for num in range(1, num_shards+1):
            # The name to use for vagrant ssh
            shard_name = 'shard%s' % num
            # The port for the shard's afe
            shard_port = lab_manifest.shards_base_port + num
            # The hostname to use in the shadow_config of the shard
            shard_hostname = '%s:%s' % (lab_manifest.vm_host_name, shard_port)

            self.vagrantfile_shard_args.update({
                shard_name: shard_name,
                '%s_shadow_config_hostname' % shard_name: shard_hostname,
                '%s_port' % shard_name: shard_port,
            })
            if lab_manifest.shards:
                board = lab_manifest.shards.pop()
                # Assign a board to a shard. Use the shard_hostname as this
                # settings is not meant to be human understandable.
                self.shard_board_map[shard_hostname] = board
                vagrant_shard_name = '%sshard' % board.rsplit(':')[-1]
                # Replace the shard<int>-type-name with board_shard
                self.vagrantfile_shard_args[shard_name] = vagrant_shard_name
                self.vagrant_shard_names.append(vagrant_shard_name)


    def __init__(self, vm_provisioner, vagrant_master_name='master',
                 num_shards=1):
        """Initialize parameters for the cluster.

        @param vm_provisioner: A provisioner object, currently the only one
            supported is VirtualBox.
        @param master_name: The name to give the cluster master.
        @param num_shards: The number of shards in the cluster. Each shard
            gets a name allocated based on its number (eg: shard1).
        """
        self.provisioner = vm_provisioner
        self.vm_manager = CoreVM(provisioner=self.provisioner)
        self._register_shards(num_shards)
        self.vagrant_master_name = vagrant_master_name


    def start_cluster(self):
        """Start a cluster."""
        self.vm_manager.setup_core_box()

        # TODO: Add a --rebuild-cluster option.
        needs_destroy = self.provisioner.initialize_vagrant(
                master=self.vagrant_master_name,
                master_port=lab_manifest.master_afe_port,
                **self.vagrantfile_shard_args)
        self.provisioner.provision(needs_destroy)


    def shutdown_cluster(self):
        """Shutdown the current cluster."""
        # TODO: Actually destroy. Halt is useful for debugging.
        self.provisioner.vagrant_cmd('halt')


    def execute_against_vm(self, vm_name, cmd):
        """Execute cmd against vm_name.

        @param cmd: The command to execute.
        @param vm_name: The name of the vm, eg: stumpyshard.
        """
        return self.provisioner.vagrant_cmd(
                "ssh %s -- '%s'" % (vm_name, cmd)).rstrip('\n')


    def _get_shadow_config_value(self, vm_name, key):
        cmd = 'grep "^%s:" %s' % (key, SHADOW_PATH)
        shadow_value = self.execute_against_vm(vm_name, cmd)
        return shadow_value.rsplit(':')[-1].lstrip(' ')


    def _check_shadow_config(self, vm, key, expected_value):
        """Sanity check the shadow_configs of all vms in the cluster.

        @raises ConfigError: If a shadow_config is misconfigured.
        """
        value = self._get_shadow_config_value(vm, key)
        if value != expected_value:
            raise ConfigError(
                    '%s vm has misconfigued config %s = %s, expected %s' %
                    (vm, key, value, expected_value))
        logging.info('%s has %s = %s', vm, key, value)


    def _upstart_cmd(self, vm, job_name, cmd='status'):
        """Execute an upstart command.

        @param vm: The name of the vm to execute it against.
        @param job_name: The name of the upstart job.
        @param cmd: The upstart command.

        @return: The output of the upstart command.
        """
        status_cmd = 'sudo %s %s' % (cmd, job_name)
        try:
            return self.execute_against_vm(vm, status_cmd)
        except vm_manager.VagrantCmdError as e:
            return '%s service not found on %s' % (job_name, vm)


    def check_services(self, action='start'):
        """Get the status of all core services on the vms.

        This method is designed to start srevices on the master/all
        shards if their shadow configs are as expected. If the shadow
        config option on a vm has an unexpected setting, services
        are not started on it.

        @param action: The action to perform on servcies. Start will
            start all of them, stop will stop them all.

        @raises ConfigError: If a shadow_config option is unexpected.
        """
        core_services = set(
                ['scheduler', 'host-scheduler',
                 'gs_offloader', 'gs_offloader_s', 'shard-client'])
        gateway = self.execute_against_vm(
                self.vagrant_master_name,
                "netstat -rn | grep \"^0.0.0.0 \" | cut -d \" \" -f10 | head -1"
                ).rstrip('\n')

        for vm in self.vagrant_shard_names + [self.vagrant_master_name]:
            vm_manager.format_msg('Checking services on %s' % vm)
            self._check_shadow_config(vm, 'host', 'localhost')
            global_db = ('localhost' if vm == self.vagrant_master_name
                         else gateway)
            self._check_shadow_config(vm, 'global_db_host', global_db)

            for service in core_services:
                logging.info('Checking %s on %s', service, vm)
                status = self._upstart_cmd(vm, service, action)
                logging.info(status)


def bringup_cluster(admin_repo, num_shards=NUM_SHARDS, start_safe=False):
    """Start a cluster.

    @param admin_repo: Path to the chromeos-admin repo.
    @param num_shards: Number of shards. You cannot change
        the number of shards on a running cluster, you need
        to destroy the cluster, remove the vagrant file,
        modify the ClusterTemplate to include a new section
        for the additional shard, and rerun clusterctl.
    @param start_safe: Start the cluster in safe mode. This means
        all core services will be stopped.
    """
    puppet_path = os.path.join(admin_repo, 'puppet')
    if not os.path.exists(puppet_path):
        raise ValueError('Admin repo %s does not contain puppet module' %
                         admin_repo)
    cluster_manager = ClusterManager(
            vm_provisioner=vm_manager.VirtualBox(puppet_path=puppet_path),
            vagrant_master_name='master', num_shards=num_shards)
    cluster_manager.start_cluster()
    try:
        cluster_manager.check_services(action='stop' if start_safe else 'start')
    except ConfigError as e:
        logging.error(
                'Shutting down cluster: %s', e)
        cluster_manager.shutdown_cluster()
        return 1


def sync():
    """Sync autotest from the host to all vms in the cluster."""
    vm_manager.format_msg('Syncing Cluster')
    vm_manager.VagrantProvisioner.vagrant_cmd('rsync', stream_output=True)
    vm_manager.VagrantProvisioner.vagrant_cmd(
            'provision --provision-with shell', stream_output=True)
    vm_manager.format_msg('Please restart services as required')


def _parse_args(args):
    """Parse command line arguments.

    @param args: A list of command line arguments, eg sys.argv[1:]

    @return: A tuple with the parsed args, as returned by parser.parse_args.
    """
    if not args:
        print ('Too few arguments, try clusterctl --help')
        sys.exit(1)

    description = ('A script to orchestrate a toy test lab. Provided '
                   'with a path to the internal repo it will download a '
                   'vm image and spin up a cluster against which you can '
                   'test core autotest changes without DUTs.')
    parser = argparse.ArgumentParser(description=description)
    subparsers = parser.add_subparsers()
    provision_subparser = subparsers.add_parser(
            'provision', help='provision a cluster')
    provision_subparser.required = False
    provision_subparser.set_defaults(which='provision')
    provision_subparser.add_argument(
            '--admin-repo', dest='admin_repo', type=str,
            help=('Path to the admin repo that has puppet scripts used for '
                  'provisioning the cluster. If you do not already have it you '
                  'can git clone the chromeos/chromeos-admin repo.'))
    provision_subparser.add_argument(
            '--safe', dest='start_safe', action='store_true',
            help='If sepcified services will not be started automatically.')

    # TODO: Automate restart of services via a --restart option.
    update_subparser = subparsers.add_parser('update', help='Update a cluster')
    update_subparser.required = False
    update_subparser.set_defaults(which='update')
    update_subparser.add_argument(
            '--sync', dest='sync', action='store_true',
            help='Sync autotest from host to all vms in cluster.')
    return parser.parse_args(args)


def main(args):
    """Main function.

    @param args: command line arguments for the script.
    """
    args = _parse_args(args)
    if args.which == 'update' and args.sync:
        sync()
    else:
        bringup_cluster(
                admin_repo=args.admin_repo, start_safe=args.start_safe)


if __name__ == '__main__':
    sys.exit(main(sys.argv[1:]))