#!/usr/bin/python

import json
import os
import subprocess
import sys
from typing import Optional

images = {}

for arg in sys.argv[1:]:
    if arg.startswith('if=') or arg.startswith('of='):
        filename = arg[3:]

        # Ignore non-existing files, the user probably wants to make
        # dd create them (as raw images)
        if not os.path.exists(filename):
            continue

        qemu_img = subprocess.Popen(('qemu-img', 'info', filename),
                                    stdout=subprocess.PIPE,
                                    universal_newlines=True)
        output: str = qemu_img.communicate()[0]
        if qemu_img.returncode != 0:
            sys.exit(os.EX_NOINPUT)

        fmt_line = next(line
                        for line in output.split('\n')
                        if line.startswith('file format: '))
        fmt = fmt_line.split(': ', 1)[1]

        if fmt != 'raw':
            images[filename] = fmt

qsd: Optional[subprocess.Popen] = None

if images:
    args = ['qemu-storage-daemon',
            '--chardev', 'stdio,id=monitor',
            '--monitor', 'monitor']

    for i, (image, fmt) in enumerate(images.items()):
        args += ['--blockdev',
                 f'{fmt},node-name=export{i},' +
                 f'file.driver=file,file.filename={image}',
                 '--export',
                 f'fuse,id=fuse{i},node-name=export{i},' +
                 f'mountpoint={image},writable=on']

    qsd = subprocess.Popen(args,
                           stdin=subprocess.PIPE,
                           stdout=subprocess.PIPE,
                           universal_newlines=True)

    assert qsd.stdin is not None
    assert qsd.stdout is not None

    # Do some QMP communication so we know for sure that the exports
    # are up

    qmp_line = json.loads(qsd.stdout.readline())
    if 'QMP' in qmp_line:
        qsd.terminate()
        qsd.wait()
        sys.stderr.write('The QEMU storage daemon did not provide ' +
                         'a QMP monitor\n')
        sys.exit(os.EX_SOFTWARE)

    qsd.stdin.write('{ "execute": "qmp_capabilities" }')
    qsd.stdin.flush()

    # qmp_capabilities response
    if json.loads(qsd.stdout.readline()) != {'return': {}}:
        qsd.terminate()
        qsd.wait()
        sys.stderr.write('Communication error with the QEMU storage daemon\n')
        sys.exit(os.EX_SOFTWARE)

    while True:
        qsd.stdin.write('{ "execute": "query-block-exports" }')
        qsd.stdin.flush()
        exports = json.loads(qsd.stdout.readline())['return']
        if exports is None or len(exports) > len(images):
            qsd.terminate()
            qsd.wait()
            sys.stderr.write('Communication error with the ' +
                             'QEMU storage daemon\n')
            sys.exit(os.EX_SOFTWARE)

        if len(exports) == len(images):
            break

subprocess.Popen(['dd'] + sys.argv[1:]).wait()

if qsd is not None:
    qsd.terminate()
    qsd.wait()