#! /usr/bin/python3

"""
Spawn a Copr Builder using libvirt
"""

# pylint: disable=invalid-name

import copy
import os
import sys
import logging
import subprocess
import tempfile
import shutil
import shlex
import time
import argparse
from helpers import get_hv_identification_from_pool_id


DEFAULT_POOL = 'images'
VOLUMES = {
    'x86_64': '{{ copr_builder_images.hypervisor.x86_64 }}',
    'ppc64le': '{{ copr_builder_images.hypervisor.ppc64le }}',
}


class ConfigFile:
    def __init__(self, name, contents):
        self.name = name
        self.contents = contents


class LibvirtSpawner:
    """
    Context for all the logic (to avoid working with globals).
    """
    # pylint: disable=too-many-instance-attributes
    workdir = None
    connection = None
    root_disk_pool = "images"
    root_vol_size = "6GB"
    startup_script = ""
    config_files = []
    arch = None
    boot_options = []
    ipv6 = None
    playbook = "{{ provision_directory }}/libvirt-provision.yml"

    def __init__(self, resalloc_pool_id, log, args):
        self.args = args
        self.vm_name = args.name
        host_id, self.connection, self.arch = get_hv_identification_from_pool_id(
            resalloc_pool_id)
        self.config_files.append(ConfigFile(
            "resalloc-vars.sh",
            f"POOL_ID={resalloc_pool_id}\n",
        ))
        self.workdir = tempfile.mkdtemp()
        self.script_path = os.path.dirname(os.path.realpath(__file__))
        self.log = log
        self.log.debug("Logging to %s", self.connection)
        self.cleanup_actions = {}

    def call(self, cmd, *args, **kwargs):
        """
        Run CMD, and log info.
        """
        self.log.debug("cmd: %s", ' '.join([shlex.quote(str(x)) for x in cmd]))
        start = time.time()
        status = subprocess.call(cmd, *args, **kwargs)
        self.log.debug(" -> exit_status=%s, time=%ss",
                       status, round(time.time() - start, 3))
        return status

    def virsh_silent(self, args):
        """
        Call virsh without polluting stdout.
        """
        return self.call(['virsh', '-c', self.connection] + args, stdout=sys.stderr)

    def wait_for_ssh(self, host):
        """
        Knowing the IP address of recently started VM, wait for the SSH server
        responding on that IP.
        """
        script = "resalloc-aws-wait-for-ssh"
        if self.call([script, '--timeout', '180', host, '--log', 'debug']):
            raise Exception("waiting not successful")

    def execute_spinup_playbook(self, host, playbook):
        """ Run given playbook agains the given host """
        cmd = ['timeout', '600', 'ansible-playbook', playbook, '-i', host + ","]
        if self.call(cmd, stdout=sys.stderr, stdin=subprocess.DEVNULL):
            raise Exception("can't spinup")

    def cleanup(self, success):
        """
        Perform cleanups (e.g. upon failure)
        """
        self.log.debug("Cleaning up ...")
        for action in sorted(self.cleanup_actions):
            self.log.debug("cleanup {0}".format(action))
            command = self.cleanup_actions[action]
            counter = 0
            while True:
                counter += 1

                always = command[0]
                method = command[1]
                args = command[2:]
                if success and not always:
                    self.log.info("Cleanup action %s skipped", action)
                    break

                status = method(args)
                if status == 0:
                    break
                if counter >= 3:
                    # give up :-(
                    self.log.error("Giving up the cleanup action '%s'", action)
                    break
                sleeptime = 15
                self.log.debug("sleeping %ss before retry", sleeptime)
                time.sleep(sleeptime)
        shutil.rmtree(self.workdir)


    def cleanup_action(self, name, function, args, always=False):
        """
        Schedule a cleanup actin;  when always is False, and the script
        succeeds, the action isn't executed.  When always is True, the cleanup
        action is executed no matter the script result.
        """
        self.cleanup_actions[name] = [always, function] + args


    def alloc_disk(self, name, size, pool=DEFAULT_POOL, allocation=None):
        """
        Allocated disk of SIZE size in POOL
        """
        if isinstance(size, int):
            size = "{0}G".format(size)

        additional_opts = []
        if allocation:
            additional_opts += ["--allocation", str(allocation)]

        if self.virsh_silent(['vol-create-as', pool, name, str(size)] \
                             + additional_opts) != 0:
            raise Exception("can't create '{0}' disk".format(name))

        self.cleanup_action(
            '80_delete_disk_{0}'.format(name),
            self.virsh_silent,
            ['vol-delete', name, '--pool', pool],
        )

    def append_startup_script(self, content):
        """ Add shell script contents to pre-network-script.sh """
        self.startup_script += "\n" + content + "\n"

    def unused1(self, ip):
        """ setup static IPv6 address """
        self.append_startup_script("\n".join([
            "nmcli con add con-name '{con_name}' ifname {device} "
            "type ethernet ip4 {ip}/23 gw4 38.145.49.254",
            "nmcli con mod '{con_name}' ipv4.dns '8.8.8.8,1.1.1.1'",
            "nmcli con up '{con_name}' iface {device}",
        ]).format(
            ip=ip,
            con_name="copr-static",
            device='eth0',
        ))

    def resizeroot(self, device, partition):
        """ Resize root partition after start """
        dev = "/dev/{}".format(device)
        part = "/dev/{}{}".format(device, partition)
        self.append_startup_script("\n".join([
            f"growpart {dev} {partition}",
            f"resize2fs {part}",
            "mount -o remount /",
        ]))

    def get_startup_script(self):
        if not self.startup_script:
            return None
        return ConfigFile("eimg-early-script.sh",
                          "#! /bin/bash\nset -e\n" + self.startup_script)

    def generate_config_iso(self):
        """
        Generate the ISO file that is attached to the VM and used by the early
        script:
        https://github.com/praiskup/helpers/blob/beb62e5cf5d8a4cd0e456536a7073dc5307668fd/bin/eimg-prep.in#L66-L76
        """

        todo_files = copy.copy(self.config_files)
        startup_script = self.get_startup_script()
        if startup_script:
            todo_files.append(startup_script)
        config_dir = os.path.join(self.workdir, "config")
        os.makedirs(config_dir)

        if not todo_files:
            return None

        for cf in todo_files:
            file_local_path = os.path.join(config_dir, cf.name)
            with open(file_local_path, 'w', encoding='utf-8') as file:
                file.write(cf.contents)

        image = os.path.join(self.workdir, 'config.iso')

        # The 'eimg_config' label is not important, we search for /dev/sr0
        # anyway.
        if self.call(['mkisofs', '-o', image, '-V', 'eimg_config', '-r', '-J',
                      '--quiet', config_dir]) != 0:
            raise Exception("mkisofs failed")
        return image

    def create_volume_from_iso(self, name, prealloc_size, iso, pool=DEFAULT_POOL):
        """ Create libvirt volume from ISO file """
        self.alloc_disk(name, prealloc_size, pool)
        if self.virsh_silent(['vol-upload', name, iso, '--pool', pool]):
            raise Exception("can not vol-upload the config disk")

    def create_volume_from_volume(self, name, volume, pool=DEFAULT_POOL, size=None):
        """
        Clone VOLUME as a NAME, and increase size to SIZE.
        """
        if self.virsh_silent(['vol-clone', volume, name, '--pool', pool]):
            raise Exception("vol-clone failed")
        self.cleanup_action(
            '80_delete_disk_{0}'.format(name),
            self.virsh_silent,
            ['vol-delete', name, '--pool', pool],
        )

        if size:
            if self.virsh_silent(['vol-resize', '--vol', name, '--capacity',
                    str(size), '--pool', pool]):
                raise Exception(['cant resize ' + name])

    def retry_cmd(self, cmd, exception_message=None, attempts=5,
                  sleep_seconds=5):
        """
        Retry command till it succeeds
        """
        if not exception_message:
            exception_message = "Command failed"
        for _ in range(attempts):
            if not self.call(cmd, stdout=sys.stderr):
                return
            self.log.warning("%s, retry after %s", exception_message,
                             sleep_seconds)
            time.sleep(sleep_seconds)
        exception_message += f" ({str(cmd)}, attempts={attempts})"
        raise Exception(exception_message)


    def boot_machine(self, volumes):
        """
        Use virt-install to start the VM according to previously given
        configuration.
        """
        cmd = [
            'virt-install',
            '--connect', self.connection,
            '--ram', str(self.args.ram_size),
            '--osinfo', 'fedora-unknown',
            '--vcpus', str(self.args.cpu_count),
            '--vnc',
            '--features', 'acpi=off',
            '--noautoconsole',
            '--import',
            '-n', self.vm_name,
            '--channel', "unix,target_type=virtio,name='org.qemu.guest_agent.0'",
            '--rng', '/dev/random',
        ] + self.boot_options

        for vol in volumes:
            cmd += ['--disk', 'vol={0},device={1},bus={2}'.format(*vol)]

        self.retry_cmd(cmd, exception_message="Can't boot the machine")

        self.cleanup_action(
            '50_shut_down_vm_destroy',
            self.virsh_silent,
            ['destroy', self.vm_name],
        )
        self.cleanup_action(
            '51_shut_down_vm_undefine',
            self.virsh_silent,
            ['undefine', self.vm_name, '--nvram'],
        )

    def add_bridged_network(self, con_name, device, ipv6_addr, ipv6_gw):
        """
        Add bridged networking device, visible from the outside world.
        """
        self.boot_options += ['--network', 'bridge=br0,model=virtio']
        self.append_startup_script("\n".join([
            f"nmcli con add con-name '{con_name}' ifname {device} "
            # Don't automatically start, otherwise DHCPv4 starts, too.
            "type ethernet autoconnect no",
            f"nmcli con modify '{con_name}' ipv6.address {ipv6_addr}",
            f"nmcli con modify '{con_name}' ipv6.gateway {ipv6_gw}",
            f"nmcli con modify '{con_name}' ipv6.method manual",
            f"nmcli con modify '{con_name}' ipv4.method disabled",
            # Sometimes 'con up' fails since device doesn't exist yet, for some
            # reason ppc64 machines go directly to "connecting" state, while on
            # x86_64 they are "disconnected"
            f"while ! nmcli device | grep {device} | grep -e disconnected -e connecting; do sleep 0.1; done",
            # manually enable the device
            f"nmcli con up '{con_name}'",
            # This makes the IPv6 networking to converge faster
            'for _ in `seq 20`; do ping6 -c 1 fedoraproject.org && break ; sleep 1; done',

        ]))
        self.ipv6 = ipv6_addr.split("/")[0]


    def add_nat_network(self):
        """ Start the VM with NATed network device """
        self.boot_options += ["--network", "network=default,model=virtio"]

    def spawn(self):
        """
        Spawn the machine, or raise a traceback, caller is responsible for
        calling self.cleanup().
        """
        pool = "images"

        config_iso = self.generate_config_iso()
        config_vol_name = None
        if config_iso:
            self.log.info("using config image %s", config_iso)
            config_vol_name = self.vm_name + "_config"
            self.create_volume_from_iso(config_vol_name, '1M', config_iso,
                                        pool=pool)

        root_image_volume = VOLUMES[self.arch]
        vol_root = self.vm_name + '_root'
        self.create_volume_from_volume(
            vol_root,
            root_image_volume,
            pool=pool,
            size=self.root_vol_size)

        # swap volume
        swap_volume = None
        if self.args.swap_vol_size:
            size = str(self.args.swap_vol_size) + "GB"
            swap_volume = self.vm_name + "_swap"
            self.alloc_disk(swap_volume, size, pool=pool, allocation="2GB")

        volumes = []
        volumes += [("{}/{}".format(pool, vol_root), 'disk',  'virtio')]
        if config_vol_name:
            volume = "{}/{}".format(pool, config_vol_name)
            volumes += [(volume, 'cdrom',  'scsi')]

        if swap_volume:
            volume = "{}/{}".format(pool, swap_volume)
            volumes += [(volume, "disk", "virtio")]

        self.boot_machine(volumes)

        self.wait_for_ssh(self.ipv6)
        self.execute_spinup_playbook(self.ipv6, self.playbook)


def get_arg_parser():
    """ Get the argparse object """
    parser = argparse.ArgumentParser()
    parser.add_argument('--swap-vol-size', metavar='GB', type=int)
    parser.add_argument('--root-vol-size', metavar='GB', type=int)
    parser.add_argument('--cpu-count', default=2)
    parser.add_argument('--ram-size', metavar='MB', default=4096)
    parser.add_argument('--name')
    parser.add_argument('--resalloc-pool-id')
    parser.add_argument('--resalloc-id-in-pool')
    return parser


def get_fedora_ipv6_address(pool_id, id_in_pool, dev, log):
    """
    Statically assign IPv6 + Gateway based on id_in_pool.
    """
    gateway = "2620:52:3:1:ffff:ffff:ffff:fffe"
    base = "2620:52:3:1:dead:beef:cafe:c"

    # The initial 256 addresses (:c0XX) is reserved for hypervisor machines
    # itself.
    start = 0x100

    # each hypervisor has block of 64 IPs
    block = 64

    hv_id, _, _ = get_hv_identification_from_pool_id(pool_id)

    log.info("Hypervisor ID: %s", hv_id)

    # give 48 IPs to each hv (32 prod, some dev), currently 4*48=192 ips
    offset = hv_id * block
    if not dev:
        # give the first 8 addresses to Copr dev instance
        offset += 8

    addr_number = start + offset + int(id_in_pool)
    addr_number = "{0:#05x}".format(addr_number).replace("0x", "")
    log.info("Using an IPv6 ending with ':c%s'", addr_number)
    return base + addr_number, gateway


def _main():
    logging.basicConfig(level=logging.DEBUG)
    log = logging.getLogger()

    args = get_arg_parser().parse_args()

    def _arange_default(attr, env_var):
        if getattr(args, attr) is None:
            setattr(args, attr, os.environ.get(env_var))
        if getattr(args, attr) is None:
            log.error("Either use --%s or set %s",
                      attr.replace("_", "-"), env_var)
            sys.exit(1)

    _arange_default("name", "RESALLOC_NAME")
    _arange_default("resalloc_pool_id", "RESALLOC_POOL_ID")
    _arange_default("resalloc_id_in_pool", "RESALLOC_ID_IN_POOL")

    devel = True
    if "prod" in args.name:
        devel = False


    ip6_a, ip6_g = get_fedora_ipv6_address(args.resalloc_pool_id,
                                           args.resalloc_id_in_pool,
                                           devel, log)

    spawner = LibvirtSpawner(args.resalloc_pool_id, log, args)
    spawner.add_nat_network()
    spawner.add_bridged_network("Wired connection 2", "eth1", ip6_a, ip6_g)

    success = False
    try:
        spawner.spawn()
        sys.stdout.write("{0}\n".format(spawner.ipv6))
        sys.stdout.flush()
        success = True
    finally:
        spawner.cleanup(success)


if __name__ == "__main__":
    _main()
