# vim: tabstop=4 shiftwidth=4 softtabstop=4

# Copyright (c) 2013 TrilioData, Inc.
# All Rights Reserved.

"""
Handling of VM disk images.
"""

import json
import os
import re
import time
import subprocess
import threading

from oslo_config import cfg
from oslo_concurrency import processutils

from workloadmgr import exception
from workloadmgr.image import glance
from workloadmgr.openstack.common import log as logging
from workloadmgr import utils
from workloadmgr import async_utils
from workloadmgr import autolog

LOG = logging.getLogger(__name__)
Logger = autolog.Logger(LOG)

image_opts = [
    cfg.BoolOpt('force_raw_images',
                default=True,
                help='Force backing images to raw format'),
]

CONF = cfg.CONF
CONF.register_opts(image_opts)

QEMU_IMG = 'qemu-img'


class QemuImgInfo(object):
    BACKING_FILE_RE = re.compile((r"^(.*?)\s*\(actual\s+path\s*:"
                                  r"\s+(.*?)\)\s*$"), re.I)
    TOP_LEVEL_RE = re.compile(r"^([\w\d\s\_\-]+):(.*)$")
    SIZE_RE = re.compile(r"\(\s*(\d+)\s+bytes\s*\)", re.I)

    def __init__(self, cmd_output=None):
        details = self._parse(cmd_output)
        self.image = details.get('image')
        self.backing_file = details.get('backing_file')
        self.file_format = details.get('file_format')
        self.virtual_size = details.get('virtual_size')
        self.cluster_size = details.get('cluster_size')
        self.disk_size = details.get('disk_size')
        self.snapshots = details.get('snapshot_list', [])
        self.encryption = details.get('encryption')

    def __str__(self):
        lines = [
            'image: %s' % self.image,
            'file_format: %s' % self.file_format,
            'virtual_size: %s' % self.virtual_size,
            'disk_size: %s' % self.disk_size,
            'cluster_size: %s' % self.cluster_size,
            'backing_file: %s' % self.backing_file,
        ]
        if self.snapshots:
            lines.append("snapshots: %s" % self.snapshots)
        return "\n".join(lines)

    def _canonicalize(self, field):
        # Standardize on underscores/lc/no dash and no spaces
        # since qemu seems to have mixed outputs here... and
        # this format allows for better integration with python
        # - ie for usage in kwargs and such...
        field = field.lower().strip()
        for c in (" ", "-"):
            field = field.replace(c, '_')
        return field

    def _extract_bytes(self, details):
        # Replace it with the byte amount
        real_size = self.SIZE_RE.search(details)
        if real_size:
            details = real_size.group(1)
        try:
            details = utils.to_bytes(details)
        except (TypeError, ValueError):
            pass
        return details

    def _extract_details(self, root_cmd, root_details, lines_after):
        consumed_lines = 0
        real_details = root_details
        if root_cmd == 'backing_file':
            # Replace it with the real backing file
            backing_match = self.BACKING_FILE_RE.match(root_details)
            if backing_match:
                real_details = backing_match.group(2).strip()
        elif root_cmd in ['virtual_size', 'cluster_size', 'disk_size']:
            # Replace it with the byte amount (if we can convert it)
            real_details = self._extract_bytes(root_details)
        elif root_cmd == 'file_format':
            real_details = real_details.strip().lower()
        elif root_cmd == 'snapshot_list':
            # Next line should be a header, starting with 'ID'
            if not lines_after or not lines_after[0].startswith("ID"):
                msg = _("Snapshot list encountered but no header found!")
                raise ValueError(msg)
            consumed_lines += 1
            possible_contents = lines_after[1:]
            real_details = []
            # This is the sprintf pattern we will try to match
            # "%-10s%-20s%7s%20s%15s"
            # ID TAG VM SIZE DATE VM CLOCK (current header)
            for line in possible_contents:
                line_pieces = line.split(None)
                if len(line_pieces) != 6:
                    break
                else:
                    # Check against this pattern in the final position
                    # "%02d:%02d:%02d.%03d"
                    date_pieces = line_pieces[5].split(":")
                    if len(date_pieces) != 3:
                        break
                    real_details.append({
                        'id': line_pieces[0],
                        'tag': line_pieces[1],
                        'vm_size': line_pieces[2],
                        'date': line_pieces[3],
                        'vm_clock': line_pieces[4] + " " + line_pieces[5],
                    })
                    consumed_lines += 1
        return (real_details, consumed_lines)

    def _parse(self, cmd_output):
        # Analysis done of qemu-img.c to figure out what is going on here
        # Find all points start with some chars and then a ':' then a newline
        # and then handle the results of those 'top level' items in a separate
        # function.
        #
        # TODO(harlowja): newer versions might have a json output format
        #                 we should switch to that whenever possible.
        #                 see: http://bit.ly/XLJXDX
        if not cmd_output:
            cmd_output = ''
        contents = {}
        lines = cmd_output.splitlines()
        i = 0
        line_am = len(lines)
        while i < line_am:
            line = lines[i]
            if not line.strip():
                i += 1
                continue
            consumed_lines = 0
            top_level = self.TOP_LEVEL_RE.match(line)
            if top_level:
                root = self._canonicalize(top_level.group(1))
                if not root:
                    i += 1
                    continue
                root_details = top_level.group(2).strip()
                details, consumed_lines = self._extract_details(root,
                                                                root_details,
                                                                lines[i + 1:])
                contents[root] = details
            i += consumed_lines + 1
        return contents


@autolog.log_method(logger=Logger)
def qemu_img_info(path, run_as_root=False):
    """Return an object containing the parsed output from qemu-img info."""
    # TODO(giri): check if the remote file exists
    # if not os.path.exists(path):
    #    return QemuImgInfo()

    # if encrypted image has backing_file format in json
    text = ['json', 'encrypt.key-secret', 'filename']
    if all([each in path for each in text]):
        LOG.info('backing_file has json format and image is encrypted.')
        path = utils.parse_encrypted_image_backing_file(path)
    head, tail = os.path.split(path)
    os.listdir(head)
    cmd = ('qemu-img', 'info', path)
    out, err = utils.execute(*cmd, run_as_root=run_as_root)

    return QemuImgInfo(out)


@async_utils.run_async
def convert_image(source, dest, out_format, run_as_root=False):
    """Async Convert image to other format."""
    cmd = ('qemu-img', 'convert', '-t', 'none', '-O',
           out_format, source, dest)
    utils.execute(*cmd, run_as_root=run_as_root)


def create_image(qemu_img_path, size, img_format='qcow2', run_as_root=False):
    """Create qemu-img in the format."""
    cmd = ('qemu-img', 'create', '-f',
           img_format, qemu_img_path, str(size))
    utils.execute(*cmd, run_as_root=run_as_root)


def fetch(context, image_href, path, _user_id, _project_id):
    # TODO(vish): Improve context handling and add owner and auth data
    #             when it is added to glance.  Right now there is no
    #             auth checking in glance, so we assume that access was
    #             checked before we got here.
    (image_service, image_id) = glance.get_remote_image_service(context,
                                                                image_href)
    with utils.remove_path_on_error(path):
        with open(path, "wb") as image_file:
            image_service.download(context, image_id, image_file)


def fetch_to_raw(context, image_href, path, user_id, project_id):
    path_tmp = "%s.part" % path
    fetch(context, image_href, path_tmp, user_id, project_id)

    with utils.remove_path_on_error(path_tmp):
        data = qemu_img_info(path_tmp)

        fmt = data.file_format
        if fmt is None:
            raise exception.ImageUnacceptable(
                reason=_("'qemu-img info' parsing failed."),
                image_id=image_href)

        backing_file = data.backing_file
        if backing_file is not None:
            raise exception.ImageUnacceptable(
                image_id=image_href,
                reason=_("fmt=%(fmt)s backed by: %(backing_file)s") %
                locals())

        if fmt != "raw" and CONF.force_raw_images:
            staged = "%s.converted" % path
            LOG.debug("%s was %s, converting to raw" % (image_href, fmt))
            with utils.remove_path_on_error(staged):
                convert_image(path_tmp, staged, 'raw')
                os.unlink(path_tmp)

                data = qemu_img_info(staged)
                if data.file_format != "raw":
                    raise exception.ImageUnacceptable(
                        image_id=image_href,
                        reason=_("Converted to raw, but format is now %s") %
                        data.file_format)

                os.rename(staged, path)
        else:
            os.rename(path_tmp, path)


@autolog.log_method(logger=Logger)
def rebase_qcow2(backing_file_base, backing_file_top, run_as_root=False, payload=None):
    """rebase the backing_file_top to backing_file_base using unsafe mode
    :param backing_file_base: backing file to rebase to
    :param backing_file_top: top file to rebase
    """
    version = processutils.execute("qemu-img",
                    "--version")[0].split('version')[1].split()[0]
    version = tuple(version.split('.'))

    def _rebase_unencrypted_image(backing_file_base, backing_file_top, run_as_root=False, payload=None):
         # inner method for rebasing unencrypted images.
        out, err = utils.execute( 'qemu-img', 'rebase', '-u',
            backing_file_top, run_as_root=run_as_root)
        if backing_file_base:
            if version >= ('2','10','0'):
                out, err = utils.execute( 'qemu-img', 'rebase', '-u',
                '-F', qcow2.file_format, '-b', backing_file_base,
                backing_file_top, run_as_root=run_as_root)
            else:
                out, err = utils.execute( 'qemu-img', 'rebase', '-u',
                    '-b', backing_file_base, backing_file_top,
                    run_as_root=run_as_root)
        return out, err


    def _rebase_encrypted_image(backing_file_base, backing_file_top, run_as_root=False, payload=None):
        # inner method for rebasing encrypted images.
        # -F backing_format is compatible with qemu versions(2.12, 4.2.0, 6.2.0)
        base_cmd, rebase, unsafe_opts, backing_format_opts, backing_format = QEMU_IMG, 'rebase', '-u', '-F', 'qcow2'
        objectdef, secret_opts = '--object', 'secret,id=sec0,data={0}'.format(payload)
        image_opts, driver_opts = '--image-opts', 'driver=qcow2,encrypt.key-secret=sec0,file.filename={0}'.format(backing_file_top)
        
        out, err = utils.execute(base_cmd, rebase, unsafe_opts, 
                        objectdef, 
                        secret_opts, 
                        image_opts, 
                        driver_opts,
                        '-p', 
                        run_as_root=run_as_root)

        if backing_file_base:
            backing, backing_opts = '-b', 'json:{"encrypt.key-secret": "sec0","driver": "qcow2","file": {"driver": "file","filename": "'+backing_file_base+'"}}'
            out, err = utils.execute(base_cmd, rebase, unsafe_opts,
                            backing_format_opts,
                            backing_format,
                            backing,
                            backing_opts,
                            objectdef,
                            secret_opts,
                            image_opts,
                            driver_opts,
                            '-p',
                            run_as_root=run_as_root)
        return out, err



    rebase_opts = {
                    0:_rebase_unencrypted_image,
                    1:_rebase_encrypted_image
                }

    for i in range(3):
        qcow2 = qemu_img_info(backing_file_top)
        try:
            out, err = rebase_opts.get(bool(payload),
                            LOG.exception('Invalid option for rebasing.'))(
                            backing_file_base, 
                            backing_file_top, 
                            run_as_root=run_as_root,
                            payload=payload)
            if not err:
                LOG.info('Rebasing of {0} is successful.'.format(backing_file_top))
        except Exception as ex:
            LOG.exception(ex)

        try:
            qcow2 = qemu_img_info(backing_file_top)
            LOG.debug("QCOW2: %s, BACKING_FILE: %s" %
                (backing_file_top, utils.sanitize_message(qcow2.backing_file)))
            if not qcow2.backing_file and \
                    qcow2.backing_file == backing_file_base:
                        return
            if utils.parse_encrypted_image_backing_file(
                    qcow2.backing_file) == backing_file_base:
                return
        except Exception as ex:
            LOG.info("qemu-img rebase -b %s %s failed with error '%s'. Retrying..." % 
                (backing_file_base, backing_file_top, err))
            time.sleep(15)

    raise Exception("qemu-img rebase -b %s %s failed with error '%s'" % 
        (backing_file_base, backing_file_top, err))


def commit_qcow2(backing_file_top, run_as_root=False):
    """rebase the backing_file_top to backing_file_base
     :param backing_file_top: top file to commit from to its base
    """
    utils.execute(
        'qemu-img',
        'commit',
        backing_file_top,
        run_as_root=run_as_root)


def resize_image(path, new_size, run_as_root=False):
    """rebase the backing_file_top to backing_file_base
     :param backing_file_top: top file to commit from to its base
    """
    utils.execute(
        'qemu-img',
        'resize',
        path,
        new_size,
        run_as_root=run_as_root)


def get_disk_backing_file(path, basename=True):
    """Get the backing file of a disk image

    :param path: Path to the disk image
    :returns: a path to the image's backing store
    """
    backing_file = qemu_img_info(path).backing_file
    if backing_file and basename:
        backing_file = os.path.basename(backing_file)

    return backing_file


@autolog.log_method(logger=Logger)
def get_effective_size(path, run_as_root=False):
    """rebase the backing_file_top to backing_file_base
     :param backing_file_top: top file to commit from to its base
    """
    # Sometimes after uploading the image from contego side
    # It takes time to reflect on Tvault storage backend.
    # In that case waiting for 60 seconds.
    for i in range(0, 12):
        if not os.path.exists(path):
            time.sleep(5)
        else:
            break

    try:
        qemuinfo = qemu_img_info(path)
    except Exception as ex:
        LOG.exception(_("qemu-img Errored. Retrying: %s"), ex)
        time.sleep(10)
        qemuinfo = qemu_img_info(path)

    return qemuinfo.virtual_size