#!/usr/bin/python
# ovirt-imageio
# Copyright (C) 2015-2016 Red Hat, Inc.
#
# 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; either version 2 of the License, or
# (at your option) any later version.

import argparse
import sys
import threading

from ovirt_imageio import ops
from ovirt_imageio import util

MB = 1024 * 1024


def send(args):
    op = ops.Send(
        args.path,
        sys.stdout,
        args.size,
        offset=args.offset,
        buffersize=args.buffersize)
    run(op)


def recv(args):
    op = ops.Receive(
        args.path,
        sys.stdin,
        args.size,
        offset=args.offset,
        buffersize=args.buffersize)
    run(op)


def zero(args):
    op = ops.Zero(
        args.path,
        args.size,
        offset=args.offset,
        flush=args.flush,
        buffersize=args.buffersize)
    run(op)


def run(op):
    done = threading.Event()
    t = util.start_thread(progress, args=(op, done))
    op.run()
    done.set()
    t.join()


def progress(op, done, stream=sys.stderr):
    start = util.monotonic_time()

    while not done.wait(0.1):
        elapsed = util.monotonic_time() - start
        if op.size:
            progress = "%6.2f%%" % (float(op.done) / op.size * 100)
        else:
            progress = " ----- "
        rate = op.done / elapsed / MB
        stream.write("[ %s ] %5.02f MiB/s %5.02fs\r"
                     % (progress, rate, elapsed))
        stream.flush()

    stream.write("[   done  ]\n")


def megabyte(s):
    return int(s) * MB


def add_arguments(parser):
    parser.add_argument("path", help="path to existing image")
    parser.add_argument("-s", "--size", type=megabyte,
                        help="amount of data to copy (in MiB)")
    parser.add_argument("-o", "--offset", dest="offset", type=int,
                        help="offset in image (in bytes)")
    parser.add_argument("-b", "--buffer-size", dest="buffersize", type=int,
                        help="copy buffer size (in bytes)")
    parser.set_defaults(size=None, offset=0, buffersize=ops.BUFFERSIZE)


parser = argparse.ArgumentParser(
    description="oVirt image daemon benchmark tool")
subparsers = parser.add_subparsers(title="commands")

send_parser = subparsers.add_parser("send",
                                    help="send image data to stdout")
send_parser.set_defaults(command=send)
add_arguments(send_parser)

recv_parser = subparsers.add_parser("recv",
                                    help="receive image data from stdin")
recv_parser.set_defaults(command=recv)
add_arguments(recv_parser)

zero_parser = subparsers.add_parser("zero",
                                    help="zero image")
zero_parser.set_defaults(command=zero)
add_arguments(zero_parser)
zero_parser.add_argument("-f", "--flush", dest="flush", action="store_true",
                         help="wait until data is flushed to disk")

args = parser.parse_args()
args.command(args)
