
from __future__ import absolute_import

import argparse
from contextlib import contextmanager
import itertools
import libvirt
import sys
import os
import threading

from ovirt_imageio_common import directio

from vdsm.common import concurrent
from vdsm.common import libvirtconnection
from vdsm.common import time
from vdsm.common.password import ProtectedPassword

_start = None


class VMAdapter(object):
    def __init__(self, vm, src):
        self._vm = vm
        self._src = src
        self._pos = 0

    def read(self, size):
        buf = self._vm.blockPeek(self._src, self._pos, size)
        self._pos += len(buf)
        return buf

    def finish(self):
        pass


class StreamAdapter(object):
    def __init__(self, stream):
        self.read = stream.recv
        self._stream = stream

    def finish(self):
        self._stream.finish()


class Sparseness(object):
    def __init__(self, opaque, estimated_size):
        self.done = 0
        self.opaque = opaque
        self.estimated_size = estimated_size


def bytesWriteHandler(stream, buf, opaque):
    fd = opaque.opaque
    return os.write(fd, buf)


def recvSkipHandler(stream, length, opaque):
    opaque.done += length
    progress = min(99, opaque.done * 100 // opaque.estimated_size)
    write_progress(progress)
    fd = opaque.opaque
    cur = os.lseek(fd, length, os.SEEK_CUR)
    return os.ftruncate(fd, cur)


def arguments(args):
    parser = argparse.ArgumentParser()
    parser.add_argument('args')
    parser.add_argument('--uri', dest='uri', required=True,
                        help='Libvirt URI')
    parser.add_argument('--username', dest='username', default='',
                        help='Libvirt login user name')
    parser.add_argument('--password-file', dest='password_file', default='',
                        help='Libvirt login password read from a file')
    parser.add_argument('--source', dest='source', nargs='+', required=True,
                        help='Source remote volumes path')
    parser.add_argument('--dest', dest='dest', nargs='+', required=True,
                        help='Destination local volumes path')
    parser.add_argument('--storage-type', dest='storagetype', nargs='+',
                        required=True, help='Storage type (volume or path)')
    parser.add_argument('--vm-name', dest='vmname', required=True,
                        help='Libvirt source VM name')
    parser.add_argument('--bufsize', dest='bufsize', default=1048576,
                        type=int, help='Size of packets in bytes, default'
                        '1048676')
    parser.add_argument('--verbose', action='store_true',
                        help='verbose output')
    parser.add_argument('--allocation', dest='allocation', default='',
                        help='Allocation Policy')

    return parser.parse_args(args)


def write_output(msg):
    sys.stdout.write('[%7.1f] %s\n' % (time.monotonic_time() - _start, msg))
    sys.stdout.flush()


def write_error(e):
    write_output("ERROR: %s" % e)


def write_progress(progress):
    sys.stdout.write('    (%d/100%%)\r' % progress)
    sys.stdout.flush()


def volume_progress(op, done, estimated_size):
    while op.done < estimated_size:
        progress = min(99, op.done * 100 // estimated_size)
        write_progress(progress)
        if done.wait(1):
            break
    write_progress(100)


@contextmanager
def progress(op, estimated_size):
    done = threading.Event()
    th = concurrent.thread(volume_progress, args=(op, done, estimated_size))
    th.start()
    try:
        yield th
    finally:
        done.set()
        th.join()


def download_disk(adapter, estimated_size, size, dest, bufsize):
    op = directio.Receive(dest, adapter, size=size, buffersize=bufsize)
    with progress(op, estimated_size):
        op.run()
    adapter.finish()


def download_disk_sparse(stream, estimated_size, size, dest, bufsize):
    fd = os.open(dest, os.O_WRONLY | os.O_CREAT | os.O_TRUNC)
    op = Sparseness(fd, estimated_size)
    with progress(op, estimated_size):
        stream.sparseRecvAll(bytesWriteHandler, recvSkipHandler, op)
    stream.finish()
    os.close(fd)


def get_password(options):
    if not options.password_file:
        return None
    if options.verbose:
        write_output('>>> Reading password from file %s' %
                     options.password_file)
    with open(options.password_file, 'r') as f:
        return ProtectedPassword(f.read())


def handle_volume(con, diskno, src, dst, options):
    write_output('Copying disk %d/%d to %s' % (diskno, len(options.source),
                                               dst))
    vol = con.storageVolLookupByPath(src)
    _, capacity, allocation = vol.info()
    if options.verbose:
        write_output('>>> disk %d, capacity: %d allocation %d' %
                     (diskno, capacity, allocation))

    estimated_size = capacity
    stream = con.newStream()
    preallocated = True

    if options.allocation == "sparse" and \
            con.getLibVersion() >= 3004000:
        try:
            preallocated = False
            vol.download(stream, 0, 0,
                         libvirt.VIR_STORAGE_VOL_DOWNLOAD_SPARSE_STREAM)
            download_disk_sparse(stream, estimated_size, None, dst,
                                 options.bufsize)
        except libvirt.libvirtError:
            preallocated = True
            write_output('WARN: sparseness is not supported')

    if preallocated:
        vol.download(stream, 0, 0, 0)
        sr = StreamAdapter(stream)
        download_disk(sr, estimated_size, None, dst, options.bufsize)


def handle_path(con, diskno, src, dst, options):
    write_output('Copying disk %d/%d to %s' % (diskno, len(options.source),
                                               dst))
    vm = con.lookupByName(options.vmname)
    info = vm.blockInfo(src)
    physical = info[2]
    if options.verbose:
        capacity = info[0]
        write_output('>>> disk %d, capacity: %d physical %d' %
                     (diskno, capacity, physical))

    vmAdapter = VMAdapter(vm, src)
    download_disk(vmAdapter, physical, physical, dst, options.bufsize)


def validate_disks(options):
    if not (len(options.source) == len(options.dest) and
            len(options.source) == len(options.storagetype)):
        write_output('>>> source, dest, and storage-type have different'
                     ' lengths')
        sys.exit(1)
    elif not all(st in ("volume", "path") for st in options.storagetype):
        write_output('>>> unsupported storage type. (supported: volume, path)')
        sys.exit(1)
    elif not options.allocation == "sparse" and \
            not options.allocation == "preallocated":
        write_output('>>> unsupported allocation policy. (supported: sparse, '
                     'preallocated)')
        sys.exit(1)


def main(argv=None):
    global _start
    _start = time.monotonic_time()

    options = arguments(argv or sys.argv)
    validate_disks(options)

    con = libvirtconnection.open_connection(options.uri,
                                            options.username,
                                            get_password(options))

    write_output('preparing for copy')
    disks = itertools.izip(options.source, options.dest, options.storagetype)
    for diskno, (src, dst, fmt) in enumerate(disks, start=1):
        if fmt == 'volume':
            handle_volume(con, diskno, src, dst, options)
        elif fmt == 'path':
            handle_path(con, diskno, src, dst, options)
        diskno = diskno + 1
    write_output('Finishing off')
