
from __future__ import absolute_import

import os
import logging
import threading
import uuid
from contextlib import contextmanager

from vdsm import constants
from vdsm import utils
from vdsm import virtsparsify
from vdsm.config import config
from vdsm.common import cmdutils
from vdsm.common import logutils
from vdsm.common.threadlocal import vars
from vdsm.storage import constants as sc
from vdsm.storage import exception as se
from vdsm.storage import imageSharing
from vdsm.storage import misc
from vdsm.storage import qemuimg
from vdsm.storage import resourceManager as rm
from vdsm.storage import sd
from vdsm.storage import volume
from vdsm.storage import workarounds
from vdsm.storage.sdc import sdCache

from vdsm.common.exception import ActionStopped

log = logging.getLogger('storage.Image')

SYNC_VOLUMES_ALL = 'ALL'
SYNC_VOLUMES_INTERNAL = 'INTERNAL'
SYNC_VOLUMES_LEAF = 'LEAF'

UNKNOWN_OP = 0
COPY_OP = 1
MOVE_OP = 2
OP_TYPES = {UNKNOWN_OP: 'UNKNOWN', COPY_OP: 'COPY', MOVE_OP: 'MOVE'}

RENAME_RANDOM_STRING_LEN = 8


def _deleteImage(dom, imgUUID, postZero, discard):
    allVols = dom.getAllVolumes()
    imgVols = sd.getVolsOfImage(allVols, imgUUID)
    if not imgVols:
        log.warning("No volumes found for image %s. %s", imgUUID, allVols)
        return
    elif postZero:
        dom.zeroImage(dom.sdUUID, imgUUID, imgVols, discard)
    else:
        dom.deleteImage(dom.sdUUID, imgUUID, imgVols)


class Image:
    log = logging.getLogger('storage.Image')
    _fakeTemplateLock = threading.Lock()

    def __init__(self, repoPath):
        self._repoPath = repoPath

    @property
    def repoPath(self):
        return self._repoPath

    def _run_qemuimg_operation(self, operation):
        self.log.debug('running qemu-img operation')
        with vars.task.abort_callback(operation.abort):
            operation.run()
        self.log.debug('qemu-img operation has completed')

    def deletedVolumeName(self, uuid):
        randomStr = misc.randomStr(RENAME_RANDOM_STRING_LEN)
        return "%s%s_%s" % (sc.REMOVED_IMAGE_PREFIX, randomStr, uuid)

    def estimate_qcow2_size(self, src_vol_params, dst_sd_id):
        qemu_measure = qemuimg.measure(
            image=src_vol_params['path'],
            format=sc.fmt2str(src_vol_params['volFormat']),
            output_format=qemuimg.FORMAT.QCOW2)

        chunk_size_mb = config.getint("irs", "volume_utilization_chunk_mb")
        chunk_size = chunk_size_mb * constants.MEGAB
        required = (qemu_measure["required"] + chunk_size)
        vol_class = sdCache.produce(dst_sd_id).getVolumeClass()
        max_size = vol_class.max_size(src_vol_params['capacity'],
                                      sc.COW_FORMAT)
        allocation = min(required, max_size)

        self.log.debug("Estimated allocation for qcow2 volume:"
                       "%d", allocation)
        return allocation

    def estimateChainSize(self, sdUUID, imgUUID, volUUID, capacity):
        chain = self.getChain(sdUUID, imgUUID, volUUID)
        log_str = logutils.volume_chain_to_str(vol.volUUID for vol in chain)
        self.log.info("chain=%s ", log_str)

        chain_allocation = 0
        template = chain[0].getParentVolume()
        if template:
            chain_allocation = template.getVolumeSize()
        for vol in chain:
            chain_allocation += vol.getVolumeSize()
        if chain_allocation > capacity:
            chain_allocation = capacity
        chain_allocation = int(chain_allocation * sc.COW_OVERHEAD)
        return chain_allocation

    def getChain(self, sdUUID, imgUUID, volUUID=None):
        chain = []
        volclass = sdCache.produce(sdUUID).getVolumeClass()

        if volUUID:
            srcVol = volclass(self.repoPath, sdUUID, imgUUID, volUUID)

            if srcVol.isShared():
                return [srcVol]

        else:
            uuidlist = volclass.getImageVolumes(sdUUID, imgUUID)

            if not uuidlist:
                raise se.ImageDoesNotExistInSD(imgUUID, sdUUID)

            srcVol = volclass(self.repoPath, sdUUID, imgUUID, uuidlist[0])

            if len(uuidlist) == 1 and srcVol.isShared():
                return [srcVol]

            for vol in uuidlist:
                srcVol = volclass(self.repoPath, sdUUID, imgUUID, vol)

                if srcVol.isLeaf():
                    break

                srcVol = None

            if not srcVol:
                self.log.error("There is no leaf in the image %s", imgUUID)
                raise se.ImageIsNotLegalChain(imgUUID)

        seen = set()

        while not srcVol.isShared():
            chain.insert(0, srcVol)
            seen.add(srcVol.volUUID)

            parentUUID = srcVol.getParent()
            if parentUUID == sc.BLANK_UUID:
                break

            if parentUUID in seen:
                self.log.error("Image %s volume %s has invalid parent UUID %s",
                               imgUUID, srcVol.volUUID, parentUUID)
                raise se.ImageIsNotLegalChain(imgUUID)

            srcVol = srcVol.getParentVolume()

        return chain

    def getTemplate(self, sdUUID, imgUUID):
        tmpl = None
        chain = self.getChain(sdUUID, imgUUID)
        pvol = chain[0].getParentVolume()
        if pvol:
            tmpl = pvol
        elif chain[0].isShared():
            tmpl = chain[0]

        return tmpl

    def createFakeTemplate(self, sdUUID, volParams):
        with self._fakeTemplateLock:
            try:
                destDom = sdCache.produce(sdUUID)
                volclass = destDom.getVolumeClass()
                volclass(self.repoPath, sdUUID, volParams['imgUUID'],
                         volParams['volUUID'])
            except (se.VolumeDoesNotExist, se.ImagePathError):
                try:
                    destDom.createVolume(
                        imgUUID=volParams['imgUUID'],
                        capacity=volParams['capacity'],
                        volFormat=sc.COW_FORMAT,
                        preallocate=sc.SPARSE_VOL,
                        diskType=volParams['disktype'],
                        volUUID=volParams['volUUID'], desc="Fake volume",
                        srcImgUUID=sc.BLANK_UUID,
                        srcVolUUID=sc.BLANK_UUID)

                    vol = destDom.produceVolume(imgUUID=volParams['imgUUID'],
                                                volUUID=volParams['volUUID'])
                    vol.setLegality(sc.FAKE_VOL)
                    vol.setShared()
                    destDom.templateRelink(volParams['imgUUID'],
                                           volParams['volUUID'])

                    self.log.debug("Succeeded to create fake image %s in "
                                   "domain %s", volParams['imgUUID'],
                                   destDom.sdUUID)
                except Exception:
                    self.log.error("Failure to create fake image %s in domain "
                                   "%s", volParams['imgUUID'], destDom.sdUUID,
                                   exc_info=True)

    def isLegal(self, sdUUID, imgUUID):
        try:
            legal = True
            volclass = sdCache.produce(sdUUID).getVolumeClass()
            vollist = volclass.getImageVolumes(sdUUID, imgUUID)
            self.log.info("image %s in domain %s has vollist %s", imgUUID,
                          sdUUID, str(vollist))
            for v in vollist:
                vol = volclass(self.repoPath, sdUUID, imgUUID, v)
                if not vol.isLegal() or vol.isFake():
                    legal = False
                    break
        except:
            legal = False
        return legal

    def __cleanupMove(self, srcVol, dstVol):
        try:
            if srcVol:
                srcVol.teardown(sdUUID=srcVol.sdUUID, volUUID=srcVol.volUUID)
            if dstVol:
                dstVol.teardown(sdUUID=dstVol.sdUUID, volUUID=dstVol.volUUID)
        except Exception:
            self.log.error("Unexpected error", exc_info=True)

    def _createTargetImage(self, destDom, srcSdUUID, imgUUID):
        try:
            srcChain = self.getChain(srcSdUUID, imgUUID)
            log_str = logutils.volume_chain_to_str(
                vol.volUUID for vol in srcChain)
            self.log.info("Source chain=%s ", log_str)
        except se.StorageException:
            self.log.error("Unexpected error", exc_info=True)
            raise
        except Exception as e:
            self.log.error("Unexpected error", exc_info=True)
            raise se.SourceImageActionError(imgUUID, srcSdUUID, str(e))

        fakeTemplate = False
        pvol = srcChain[0].getParentVolume()
        if pvol:
            volParams = pvol.getVolumeParams()
            if destDom.isBackup():
                fakeTemplate = True

        @contextmanager
        def justLogIt(img):
            self.log.debug("You don't really need lock parent of image %s",
                           img)
            yield

        dstImageResourcesNamespace = rm.getNamespace(sc.IMAGE_NAMESPACE,
                                                     destDom.sdUUID)
        with rm.acquireResource(dstImageResourcesNamespace, pimg, rm.SHARED) \
                if pimg != sc.BLANK_UUID else justLogIt(imgUUID):
            if fakeTemplate:
                self.createFakeTemplate(destDom.sdUUID, volParams)

            dstChain = []
            for srcVol in srcChain:
                try:
                    volParams = srcVol.getVolumeParams()

                    if (destDom.supportsSparseness or
                            volParams['volFormat'] != sc.RAW_FORMAT):
                        tmpVolPreallocation = sc.SPARSE_VOL
                    else:
                        tmpVolPreallocation = sc.PREALLOCATED_VOL

                    destDom.createVolume(
                        imgUUID=imgUUID,
                        capacity=volParams['capacity'],
                        volFormat=volParams['volFormat'],
                        preallocate=tmpVolPreallocation,
                        diskType=volParams['disktype'],
                        volUUID=srcVol.volUUID,
                        desc=volParams['descr'],
                        srcImgUUID=pimg,
                        srcVolUUID=volParams['parent'])

                    dstVol = destDom.produceVolume(imgUUID=imgUUID,
                                                   volUUID=srcVol.volUUID)

                    dstVol.extend(volParams['apparentsize'])

                    if volParams['prealloc'] == sc.PREALLOCATED_VOL \
                            and tmpVolPreallocation != sc.PREALLOCATED_VOL:
                        dstVol.setType(sc.PREALLOCATED_VOL)

                    dstChain.append(dstVol)
                except se.StorageException:
                    self.log.error("Unexpected error", exc_info=True)
                    raise
                except Exception as e:
                    self.log.error("Unexpected error", exc_info=True)
                    raise se.DestImageActionError(imgUUID, destDom.sdUUID,
                                                  str(e))

                pimg = imgUUID

        return {'srcChain': srcChain, 'dstChain': dstChain}

    def _interImagesCopy(self, destDom, srcSdUUID, imgUUID, chains):
        srcLeafVol = chains['srcChain'][-1]
        dstLeafVol = chains['dstChain'][-1]
        try:
            srcLeafVol.prepare(rw=False)
            dstLeafVol.prepare(rw=True, chainrw=True, setrw=True)
        except Exception:
            self.log.error("Unexpected error", exc_info=True)
            self.__cleanupMove(srcLeafVol, dstLeafVol)
            raise

        try:
            for srcVol in chains['srcChain']:
                try:
                    dstVol = destDom.produceVolume(imgUUID=imgUUID,
                                                   volUUID=srcVol.volUUID)

                    if workarounds.invalid_vm_conf_disk(srcVol):
                        srcFormat = dstFormat = qemuimg.FORMAT.RAW
                    else:
                        srcFormat = sc.fmt2str(srcVol.getFormat())
                        dstFormat = sc.fmt2str(dstVol.getFormat())

                    parentVol = dstVol.getParentVolume()

                    if parentVol is not None:
                        backing = volume.getBackingVolumePath(
                            imgUUID, parentVol.volUUID)
                        backingFormat = sc.fmt2str(parentVol.getFormat())
                    else:
                        backing = None
                        backingFormat = None

                    if (destDom.supportsSparseness and
                            dstVol.getType() == sc.PREALLOCATED_VOL):
                        preallocation = qemuimg.PREALLOCATION.FALLOC
                    else:
                        preallocation = None

                    operation = qemuimg.convert(
                        srcVol.getVolumePath(),
                        dstVol.getVolumePath(),
                        srcFormat=srcFormat,
                        dstFormat=dstFormat,
                        dstQcow2Compat=destDom.qcow2_compat(),
                        backing=backing,
                        backingFormat=backingFormat,
                        preallocation=preallocation,
                        unordered_writes=destDom.recommends_unordered_writes(
                            dstVol.getFormat()))
                    with utils.stopwatch("Copy volume %s"
                                         % srcVol.volUUID):
                        self._run_qemuimg_operation(operation)
                except ActionStopped:
                    raise
                except se.StorageException:
                    self.log.error("Unexpected error", exc_info=True)
                    raise
                except Exception:
                    self.log.error("Copy image error: image=%s, src domain=%s,"
                                   " dst domain=%s", imgUUID, srcSdUUID,
                                   destDom.sdUUID, exc_info=True)
                    raise se.CopyImageError()
        finally:
            self.__cleanupMove(srcLeafVol, dstLeafVol)

    def _finalizeDestinationImage(self, destDom, imgUUID, chains, force):
        for srcVol in chains['srcChain']:
            try:
                dstVol = destDom.produceVolume(imgUUID=imgUUID,
                                               volUUID=srcVol.volUUID)
                if srcVol.isShared():
                    dstVol.setShared()
                elif srcVol.isInternal():
                    dstVol.setInternal()
            except se.StorageException:
                self.log.error("Unexpected error", exc_info=True)
                raise
            except Exception as e:
                self.log.error("Unexpected error", exc_info=True)
                raise se.DestImageActionError(imgUUID, destDom.sdUUID, str(e))

    def move(self, srcSdUUID, dstSdUUID, imgUUID, vmUUID, op, postZero, force,
             discard):
        self.log.info("srcSdUUID=%s dstSdUUID=%s imgUUID=%s vmUUID=%s op=%s "
                      "force=%s postZero=%s discard=%s", srcSdUUID, dstSdUUID,
                      imgUUID, vmUUID, OP_TYPES[op], str(force), str(postZero),
                      discard)

        destDom = sdCache.produce(dstSdUUID)
        if not self.isLegal(destDom.sdUUID, imgUUID):
            force = True
        if force:
            self.log.info("delete image %s on domain %s before overwriting",
                          imgUUID, destDom.sdUUID)
            _deleteImage(destDom, imgUUID, postZero, discard)

        chains = self._createTargetImage(destDom, srcSdUUID, imgUUID)
        self._interImagesCopy(destDom, srcSdUUID, imgUUID, chains)
        self._finalizeDestinationImage(destDom, imgUUID, chains, force)
        if force:
            leafVol = chains['dstChain'][-1]
            destDom.templateRelink(imgUUID, leafVol.volUUID)

        vars.task.clearRecoveries()
        if op == MOVE_OP:
            try:
                dom = sdCache.produce(srcSdUUID)
                _deleteImage(dom, imgUUID, postZero, discard)
            except se.StorageException:
                self.log.warning("Failed to remove img: %s from srcDom %s: "
                                 "after it was copied to: %s", imgUUID,
                                 srcSdUUID, dstSdUUID)

        self.log.info("%s task on image %s was successfully finished",
                      OP_TYPES[op], imgUUID)
        return True

    def _getSparsifyVolume(self, sdUUID, imgUUID, volUUID):
        volClass = sdCache.produce(sdUUID).getVolumeClass()
        return volClass(self.repoPath, sdUUID, imgUUID, volUUID)

    def sparsify(self, tmpSdUUID, tmpImgUUID, tmpVolUUID, dstSdUUID,
                 dstImgUUID, dstVolUUID):
        self.log.info("tmpSdUUID=%s, tmpImgUUID=%s, tmpVolUUID=%s, "
                      "dstSdUUID=%s, dstImgUUID=%s, dstVolUUID=%s", tmpSdUUID,
                      tmpImgUUID, tmpVolUUID, dstSdUUID, dstImgUUID,
                      dstVolUUID)

        tmpVolume = self._getSparsifyVolume(tmpSdUUID, tmpImgUUID, tmpVolUUID)
        dstVolume = self._getSparsifyVolume(dstSdUUID, dstImgUUID, dstVolUUID)

        if not dstVolume.isSparse():
            raise se.VolumeNotSparse()

        srcVolume = self._getSparsifyVolume(tmpSdUUID, tmpImgUUID,
                                            tmpVolume.getParent())

        tmpVolume.prepare()
        try:
            dstVolume.prepare()
            try:
                tmpVolume.extend(tmpVolume.getCapacity())
                dstVolume.extend(tmpVolume.getCapacity())

                srcFormat = sc.fmt2str(srcVolume.getFormat())
                dstFormat = sc.fmt2str(dstVolume.getFormat())

                virtsparsify.sparsify(srcVolume.getVolumePath(),
                                      tmpVolume.getVolumePath(),
                                      dstVolume.getVolumePath(),
                                      src_format=srcFormat,
                                      dst_format=dstFormat)
            except Exception:
                self.log.exception('Unexpected error sparsifying %s',
                                   tmpVolUUID)
                raise se.CannotSparsifyVolume(tmpVolUUID)
            finally:
                dstVolume.teardown(sdUUID=dstSdUUID, volUUID=dstVolUUID)
        finally:
            tmpVolume.teardown(sdUUID=tmpSdUUID, volUUID=tmpVolUUID)

        self._shrinkVolumeToOptimalSize(tmpVolume)
        self._shrinkVolumeToOptimalSize(dstVolume)

    def cloneStructure(self, sdUUID, imgUUID, dstSdUUID):
        self._createTargetImage(sdCache.produce(dstSdUUID), sdUUID, imgUUID)

    def syncData(self, sdUUID, imgUUID, dstSdUUID, syncType):
        srcChain = self.getChain(sdUUID, imgUUID)
        log_str = logutils.volume_chain_to_str(vol.volUUID for vol in srcChain)
        self.log.info("Source chain=%s ", log_str)

        dstChain = self.getChain(dstSdUUID, imgUUID)
        log_str = logutils.volume_chain_to_str(vol.volUUID for vol in dstChain)
        self.log.info("Dest chain=%s ", log_str)

        if syncType == SYNC_VOLUMES_INTERNAL:
            try:
                del srcChain[-1], dstChain[-1]
            except IndexError:
                raise se.ImageIsNotLegalChain()
        elif syncType == SYNC_VOLUMES_LEAF:
            try:
                del srcChain[:-1], dstChain[:-1]
            except IndexError:
                raise se.ImageIsNotLegalChain()
        elif syncType != SYNC_VOLUMES_ALL:
            raise se.MiscNotImplementedException()

        if len(srcChain) != len(dstChain):
            raise se.DestImageActionError(imgUUID, dstSdUUID)

        for i, v in enumerate(srcChain):
            if v.volUUID != dstChain[i].volUUID:
                raise se.DestImageActionError(imgUUID, dstSdUUID)

        dstDom = sdCache.produce(dstSdUUID)

        self._interImagesCopy(dstDom, sdUUID, imgUUID,
                              {'srcChain': srcChain, 'dstChain': dstChain})
        self._finalizeDestinationImage(dstDom, imgUUID,
                                       {'srcChain': srcChain,
                                        'dstChain': dstChain}, False)

    def __cleanupCopy(self, srcVol, dstVol):
        try:
            if srcVol:
                srcVol.teardown(sdUUID=srcVol.sdUUID, volUUID=srcVol.volUUID)
            if dstVol:
                dstVol.teardown(sdUUID=dstVol.sdUUID, volUUID=dstVol.volUUID)
        except Exception:
            self.log.error("Unexpected error", exc_info=True)

    def validateVolumeChain(self, sdUUID, imgUUID):
        if not self.isLegal(sdUUID, imgUUID):
            raise se.ImageIsNotLegalChain(imgUUID)
        chain = self.getChain(sdUUID, imgUUID)
        log_str = logutils.volume_chain_to_str(vol.volUUID for vol in chain)
        self.log.info("Current chain=%s ", log_str)

        pvol = chain[0].getParentVolume()
        if pvol:
            if not pvol.isLegal() or pvol.isFake():
                raise se.ImageIsNotLegalChain(imgUUID)

    def copyCollapsed(self, sdUUID, vmUUID, srcImgUUID, srcVolUUID, dstImgUUID,
                      dstVolUUID, descr, dstSdUUID, volType, volFormat,
                      preallocate, postZero, force, discard):
        self.log.info("sdUUID=%s vmUUID=%s srcImgUUID=%s srcVolUUID=%s "
                      "dstImgUUID=%s dstVolUUID=%s dstSdUUID=%s volType=%s "
                      "volFormat=%s preallocate=%s force=%s postZero=%s "
                      "discard=%s",
                      sdUUID, vmUUID, srcImgUUID, srcVolUUID, dstImgUUID,
                      dstVolUUID, dstSdUUID, volType,
                      sc.type2name(volFormat), sc.type2name(preallocate),
                      str(force), str(postZero), discard)
        try:
            srcVol = dstVol = None

            if dstSdUUID == sd.BLANK_UUID:
                dstSdUUID = sdUUID
            volclass = sdCache.produce(sdUUID).getVolumeClass()
            destDom = sdCache.produce(dstSdUUID)

            try:
                srcVol = volclass(self.repoPath, sdUUID, srcImgUUID,
                                  srcVolUUID)
            except se.StorageException:
                raise
            except Exception as e:
                self.log.error(e, exc_info=True)
                raise se.SourceImageActionError(srcImgUUID, sdUUID, str(e))

            try:
                srcVol.prepare(rw=False)

                volParams = srcVol.getVolumeParams()

                if volFormat in [sc.COW_FORMAT, sc.RAW_FORMAT]:
                    dstVolFormat = volFormat
                else:
                    dstVolFormat = volParams['volFormat']

                dst_vol_allocation = self.calculate_vol_alloc(
                    sdUUID, volParams, dstSdUUID, dstVolFormat)

                if preallocate in [sc.PREALLOCATED_VOL, sc.SPARSE_VOL]:
                    volParams['prealloc'] = preallocate

                initial_size = self.calculate_initial_size(
                    destDom.supportsSparseness,
                    dstVolFormat,
                    volParams['prealloc'],
                    dst_vol_allocation)

                self.log.info(
                    "Copy source %s:%s:%s to destination %s:%s:%s size=%s, "
                    "initial size=%s",
                    sdUUID,
                    srcImgUUID,
                    srcVolUUID,
                    dstSdUUID,
                    dstImgUUID,
                    dstVolUUID,
                    volParams['capacity'],
                    initial_size)

                if not self.isLegal(dstSdUUID, dstImgUUID):
                    force = True

                if force:
                    self.log.info("delete image %s on domain %s before "
                                  "overwriting", dstImgUUID, dstSdUUID)
                    _deleteImage(destDom, dstImgUUID, postZero, discard)

                destDom.createVolume(
                    imgUUID=dstImgUUID,
                    capacity=volParams['capacity'],
                    volFormat=dstVolFormat,
                    preallocate=volParams['prealloc'],
                    diskType=volParams['disktype'],
                    volUUID=dstVolUUID,
                    desc=descr,
                    srcImgUUID=sc.BLANK_UUID,
                    srcVolUUID=sc.BLANK_UUID,
                    initial_size=initial_size)

                dstVol = sdCache.produce(dstSdUUID).produceVolume(
                    imgUUID=dstImgUUID, volUUID=dstVolUUID)

            except se.StorageException:
                self.log.error("Unexpected error", exc_info=True)
                raise
            except Exception as e:
                self.log.error("Unexpected error", exc_info=True)
                raise se.CopyImageError("Destination volume %s error: %s" %
                                        (dstVolUUID, str(e)))

            try:
                dstVol.prepare(rw=True, setrw=True)

                if (destDom.supportsSparseness and
                        dstVol.getType() == sc.PREALLOCATED_VOL):
                    preallocation = qemuimg.PREALLOCATION.FALLOC
                else:
                    preallocation = None

                try:
                    operation = qemuimg.convert(
                        volParams['path'],
                        dstVol.getVolumePath(),
                        srcFormat=sc.fmt2str(volParams['volFormat']),
                        dstFormat=sc.fmt2str(dstVolFormat),
                        dstQcow2Compat=destDom.qcow2_compat(),
                        preallocation=preallocation,
                        unordered_writes=destDom.recommends_unordered_writes(
                            dstVolFormat))
                    with utils.stopwatch("Copy volume %s"
                                         % srcVol.volUUID):
                        self._run_qemuimg_operation(operation)
                except ActionStopped:
                    raise
                except cmdutils.Error as e:
                    self.log.exception('conversion failure for volume %s',
                                       srcVol.volUUID)
                    raise se.CopyImageError(str(e))

                if volType == sc.SHARED_VOL:
                    dstVol.setShared()

                dstVol.setLegality(sc.LEGAL_VOL)

                if force:
                    destDom.templateRelink(dstImgUUID, dstVolUUID)
            except se.StorageException:
                self.log.error("Unexpected error", exc_info=True)
                raise
            except Exception as e:
                self.log.error("Unexpected error", exc_info=True)
                raise se.CopyImageError("src image=%s, dst image=%s: msg=%s" %
                                        (srcImgUUID, dstImgUUID, str(e)))

            self.log.info("Finished copying %s:%s -> %s:%s", sdUUID,
                          srcVolUUID, dstSdUUID, dstVolUUID)
            return dstVolUUID
        finally:
            self.__cleanupCopy(srcVol=srcVol, dstVol=dstVol)

    def calculate_initial_size(self, is_file, format, prealloc,
                               estimate):
        if is_file:
            if format == sc.RAW_FORMAT and prealloc == sc.PREALLOCATED_VOL:
                return 0
        else:
            if format == sc.COW_FORMAT and prealloc == sc.SPARSE_VOL:
                return estimate

        return None

    def calculate_vol_alloc(self, src_sd_id, src_vol_params,
                            dst_sd_id, dst_vol_format):
        if dst_vol_format == sc.RAW_FORMAT:
            return src_vol_params['capacity']
        else:
            if src_vol_params['volFormat'] == sc.COW_FORMAT:
                if src_vol_params['parent'] != sc.BLANK_UUID:
                    if src_vol_params['prealloc'] != sc.SPARSE_VOL:
                        raise se.IncorrectFormat(self)
                    return self.estimateChainSize(
                        src_sd_id,
                        src_vol_params['imgUUID'],
                        src_vol_params['volUUID'],
                        src_vol_params['capacity'])
                else:
                    return self.estimate_qcow2_size(src_vol_params, dst_sd_id)
            else:
                return self.estimate_qcow2_size(src_vol_params, dst_sd_id)

    def markIllegalSubChain(self, sdDom, imgUUID, chain):
        if not chain:
            raise se.InvalidParameterException("chain", str(chain))

        volclass = sdDom.getVolumeClass()
        ancestor = chain[0]
        successor = chain[-1]
        tmpVol = volclass(self.repoPath, sdDom.sdUUID, imgUUID, successor)
        dstParent = volclass(self.repoPath, sdDom.sdUUID, imgUUID,
                             ancestor).getParent()

        while tmpVol and dstParent != tmpVol.volUUID:
            vol = tmpVol.getParentVolume()
            tmpVol.setLegality(sc.ILLEGAL_VOL)
            tmpVol = vol

    def __teardownSubChain(self, sdUUID, imgUUID, chain):
        if not chain:
            raise se.InvalidParameterException("chain", str(chain))

        volclass = sdCache.produce(sdUUID).getVolumeClass()
        ancestor = chain[0]
        successor = chain[-1]
        srcVol = volclass(self.repoPath, sdUUID, imgUUID, successor)
        dstParent = volclass(self.repoPath, sdUUID, imgUUID,
                             ancestor).getParent()

        while srcVol and dstParent != srcVol.volUUID:
            try:
                self.log.info("Teardown volume %s from image %s",
                              srcVol.volUUID, imgUUID)
                vol = srcVol.getParentVolume()
                srcVol.teardown(sdUUID=srcVol.sdUUID, volUUID=srcVol.volUUID,
                                justme=True)
                srcVol = vol
            except Exception:
                self.log.info("Failure to teardown volume %s in subchain %s "
                              "-> %s", srcVol.volUUID, ancestor, successor,
                              exc_info=True)

    def removeSubChain(self, sdDom, imgUUID, chain, postZero, discard):
        if not chain:
            raise se.InvalidParameterException("chain", str(chain))

        volclass = sdDom.getVolumeClass()
        ancestor = chain[0]
        successor = chain[-1]
        srcVol = volclass(self.repoPath, sdDom.sdUUID, imgUUID, successor)
        dstParent = volclass(self.repoPath, sdDom.sdUUID, imgUUID,
                             ancestor).getParent()

        while srcVol and dstParent != srcVol.volUUID:
            self.log.info("Remove volume %s from image %s", srcVol.volUUID,
                          imgUUID)
            vol = srcVol.getParentVolume()
            srcVol.delete(postZero=postZero, force=True, discard=discard)
            chain.remove(srcVol.volUUID)
            srcVol = vol

    def _internalVolumeMerge(self, sdDom, srcVolParams, volParams, new_size,
                             chain):
        srcVol = sdDom.produceVolume(imgUUID=srcVolParams['imgUUID'],
                                     volUUID=srcVolParams['volUUID'])
        srcVol.extend(new_size)

        srcVol.prepare(rw=True, chainrw=True, setrw=True)
        try:
            backingVolPath = volume.getBackingVolumePath(
                srcVolParams['imgUUID'], volParams['volUUID'])
            srcVol.rebase(volParams['volUUID'], backingVolPath,
                          volParams['volFormat'], unsafe=False, rollback=True)
        finally:
            srcVol.teardown(sdUUID=srcVol.sdUUID, volUUID=srcVol.volUUID)

        chain.remove(srcVolParams['volUUID'])
        self.__teardownSubChain(sdDom.sdUUID, srcVolParams['imgUUID'], chain)

        return chain

    def _baseCowVolumeMerge(self, sdDom, srcVolParams, volParams, new_size,
                            chain, discard):
        srcVol = sdDom.produceVolume(imgUUID=srcVolParams['imgUUID'],
                                     volUUID=srcVolParams['volUUID'])
        srcVol.extend(new_size)
        newUUID = str(uuid.uuid4())
        sdDom.createVolume(
            imgUUID=srcVolParams['imgUUID'],
            capacity=volParams['capacity'],
            volFormat=volParams['volFormat'],
            preallocate=sc.SPARSE_VOL,
            diskType=volParams['disktype'],
            volUUID=newUUID,
            desc="New base volume",
            srcImgUUID=sc.BLANK_UUID,
            srcVolUUID=sc.BLANK_UUID)

        tmpVol = sdDom.produceVolume(imgUUID=srcVolParams['imgUUID'],
                                     volUUID=newUUID)
        tmpVol.prepare(rw=True, justme=True, setrw=True)

        srcVol.prepare(rw=True, chainrw=True, setrw=True)
        try:
            backingVolPath = volume.getBackingVolumePath(
                srcVolParams['imgUUID'], newUUID)
            srcVol.rebase(newUUID, backingVolPath, volParams['volFormat'],
                          unsafe=False, rollback=True)
        finally:
            srcVol.teardown(sdUUID=srcVol.sdUUID, volUUID=srcVol.volUUID)

        srcVol.prepare(rw=True, chainrw=True, setrw=True)
        try:
            srcVol.rebase(sc.BLANK_UUID, "", volParams['volFormat'],
                          unsafe=True, rollback=False)
        finally:
            srcVol.teardown(sdUUID=srcVol.sdUUID, volUUID=srcVol.volUUID)

        tmpVol.teardown(sdUUID=tmpVol.sdUUID, volUUID=tmpVol.volUUID,
                        justme=True)
        tmpVol.delete(postZero=False, force=True, discard=discard)

        chain.remove(srcVolParams['volUUID'])
        self.__teardownSubChain(sdDom.sdUUID, srcVolParams['imgUUID'], chain)

        return chain

    def _baseRawVolumeMerge(self, sdDom, srcVolParams, volParams, chain):
        srcVol = chain[-1]
        with srcVol.scopedPrepare(rw=True, chainrw=True, setrw=True):
            chList = srcVolParams['children']
            newUUID = srcVol.volUUID + "_MERGE"
            sdDom.createVolume(
                imgUUID=srcVolParams['imgUUID'],
                capacity=srcVolParams['capacity'],
                volFormat=volParams['volFormat'],
                preallocate=volParams['prealloc'],
                diskType=volParams['disktype'],
                volUUID=newUUID,
                desc=srcVolParams['descr'],
                srcImgUUID=sc.BLANK_UUID,
                srcVolUUID=sc.BLANK_UUID)

            newVol = sdDom.produceVolume(imgUUID=srcVolParams['imgUUID'],
                                         volUUID=newUUID)
            with newVol.scopedPrepare(rw=True, justme=True, setrw=True):
                try:
                    operation = qemuimg.convert(
                        srcVolParams['path'],
                        newVol.getVolumePath(),
                        srcFormat=sc.fmt2str(srcVolParams['volFormat']),
                        dstFormat=sc.fmt2str(volParams['volFormat']),
                        dstQcow2Compat=sdDom.qcow2_compat(),
                        unordered_writes=sdDom.recommends_unordered_writes(
                            volParams['volFormat']))
                    with utils.stopwatch("Copy volume %s"
                                         % srcVol.volUUID):
                        self._run_qemuimg_operation(operation)
                except cmdutils.Error:
                    self.log.exception('conversion failure for volume %s',
                                       srcVol.volUUID)
                    raise se.MergeSnapshotsError(newUUID)

        if chList:
            newVol.setInternal()

        tmpUUID = self.deletedVolumeName(srcVol.volUUID)
        srcVol.rename(tmpUUID)
        newVol.rename(srcVolParams['volUUID'])

        for ch in chList:
            ch.prepare(rw=True, chainrw=True, setrw=True, force=True)
            backingVolPath = volume.getBackingVolumePath(
                srcVolParams['imgUUID'], srcVolParams['volUUID'])
            try:
                ch.rebase(srcVolParams['volUUID'], backingVolPath,
                          volParams['volFormat'], unsafe=True, rollback=True)
            finally:
                ch.teardown(sdUUID=ch.sdUUID, volUUID=ch.volUUID)
            ch.recheckIfLeaf()

        rmChain = [vol.volUUID for
                   vol in chain if vol.volUUID != srcVolParams['volUUID']]
        rmChain.append(tmpUUID)

        return rmChain

    def subChainSizeCalc(self, ancestor, successor, vols):
        chain = []
        accumulated_chain_size = 0
        currVolName = successor
        while (currVolName != endVolName):
            chain.insert(0, currVolName)
            vol_size = vols[currVolName].getVolumeSize()
            accumulated_chain_size += vol_size
            currVolName = vols[currVolName].getParent()

        return accumulated_chain_size, chain

    def syncVolumeChain(self, sdUUID, imgUUID, volUUID, actualChain):
        curChain = self.getChain(sdUUID, imgUUID, volUUID)
        log_str = logutils.volume_chain_to_str(vol.volUUID for vol in curChain)
        self.log.info("Current chain=%s ", log_str)

        subChain = []
        for vol in curChain:
            if vol.volUUID not in actualChain:
                subChain.insert(0, vol.volUUID)
            elif len(subChain) > 0:
                break
        if len(subChain) == 0:
            return
        self.log.info("Unlinking subchain: %s", subChain)

        sdDom = sdCache.produce(sdUUID=sdUUID)
        dstParent = sdDom.produceVolume(imgUUID, subChain[0]).getParent()
        subChainTailVol = sdDom.produceVolume(imgUUID, subChain[-1])
        if subChainTailVol.isLeaf():
            self.log.info("Leaf volume %s is being removed from the chain. "
                          "Marking it ILLEGAL to prevent data corruption",
                          subChainTailVol.volUUID)
            subChainTailVol.setLegality(sc.ILLEGAL_VOL)
        else:
            for childID in subChainTailVol.getChildren():
                self.log.info("Setting parent of volume %s to %s",
                              childID, dstParent)
                sdDom.produceVolume(imgUUID, childID). \
                    setParentMeta(dstParent)

    def reconcileVolumeChain(self, sdUUID, imgUUID, leafVolUUID):
        dom = sdCache.produce(sdUUID)
        allVols = dom.getAllVolumes()
        imgVolumes = sd.getVolsOfImage(allVols, imgUUID).keys()
        dom.activateVolumes(imgUUID, imgVolumes)

        actualVolumes = []
        volUUID = leafVolUUID
        while volUUID is not None:
            actualVolumes.insert(0, volUUID)
            vol = dom.produceVolume(imgUUID, volUUID)
            qemuImgFormat = sc.fmt2str(vol.getFormat())
            imgInfo = qemuimg.info(vol.volumePath, qemuImgFormat)
            backingFile = imgInfo.get('backingfile')
            if backingFile is not None:
                volUUID = os.path.basename(backingFile)
            else:
                volUUID = None

        vol = dom.produceVolume(imgUUID, leafVolUUID)
        if vol.getLegality() == sc.ILLEGAL_VOL:
            actualVolumes.remove(leafVolUUID)

        self.syncVolumeChain(sdUUID, imgUUID, actualVolumes[-1], actualVolumes)

        dom.deactivateImage(imgUUID)
        return actualVolumes

    def merge(self, sdUUID, vmUUID, imgUUID, ancestor, successor, postZero,
              discard):
        self.log.info("sdUUID=%s vmUUID=%s imgUUID=%s ancestor=%s successor=%s"
                      "postZero=%s discard=%s",
                      sdUUID, vmUUID, imgUUID, ancestor, successor,
                      str(postZero), discard)
        sdDom = sdCache.produce(sdUUID)
        allVols = sdDom.getAllVolumes()
        volsImgs = sd.getVolsOfImage(allVols, imgUUID)
        vols = {}
        for vName in volsImgs.iterkeys():
            vols[vName] = sdDom.produceVolume(imgUUID, vName)

        srcVol = vols[successor]
        srcVolParams = srcVol.getVolumeParams()
        srcVolParams['children'] = []
        for vName, vol in vols.iteritems():
            if vol.getParent() == successor:
                srcVolParams['children'].append(vol)
        dstVol = vols[ancestor]
        dstParentUUID = dstVol.getParent()
        if dstParentUUID != sd.BLANK_UUID:
            volParams = vols[dstParentUUID].getVolumeParams()
        else:
            volParams = dstVol.getVolumeParams()

        acc_size, chain = self.subChainSizeCalc(ancestor, successor, vols)
        reqSize = min(acc_size, volParams['capacity']) * sc.COW_OVERHEAD
        try:
            if dstParentUUID != sd.BLANK_UUID:
                self.log.info("Internal volume merge: src = %s dst = %s",
                              srcVol.getVolumePath(), dstVol.getVolumePath())
                chainToRemove = self._internalVolumeMerge(
                    sdDom, srcVolParams, volParams, reqSize, chain)
            elif volParams['volFormat'] == sc.RAW_FORMAT:
                self.log.info("merge with convert: src = %s dst = %s",
                              srcVol.getVolumePath(), dstVol.getVolumePath())
                chainToRemove = self._baseRawVolumeMerge(
                    sdDom, srcVolParams, volParams,
                    [vols[vName] for vName in chain])
            else:
                self.log.info("4 steps merge: src = %s dst = %s",
                              srcVol.getVolumePath(), dstVol.getVolumePath())
                chainToRemove = self._baseCowVolumeMerge(
                    sdDom, srcVolParams, volParams, reqSize, chain, discard)

            vars.task.clearRecoveries()
            self.markIllegalSubChain(sdDom, imgUUID, chainToRemove)
        except ActionStopped:
            raise
        except se.StorageException:
            self.log.error("Unexpected error", exc_info=True)
            raise
        except Exception as e:
            self.log.error(e, exc_info=True)
            raise se.SourceImageActionError(imgUUID, sdUUID, str(e))

        try:
            self.removeSubChain(sdDom, imgUUID, chainToRemove, postZero,
                                discard)
        except Exception:
            self.log.error("Failure to remove subchain %s -> %s in image %s",
                           ancestor, successor, imgUUID, exc_info=True)

        newVol = sdDom.produceVolume(imgUUID=srcVolParams['imgUUID'],
                                     volUUID=srcVolParams['volUUID'])
        try:
            self._shrinkVolumeToOptimalSize(newVol)
        except cmdutils.Error:
            self.log.warning("Auto shrink after merge failed", exc_info=True)

        self.log.info("Merge src=%s with dst=%s was successfully finished.",
                      srcVol.getVolumePath(), dstVol.getVolumePath())

    def _shrinkVolumeToOptimalSize(self, vol):
        if not vol.chunked():
            return

        vol.prepare()
        try:
            optimal_size = vol.optimal_size()
        finally:
            vol.teardown(vol.sdUUID, vol.volUUID)

        vol.reduce(optimal_size)

    def _activateVolumeForImportExport(self, domain, imgUUID, volUUID=None):
        chain = self.getChain(domain.sdUUID, imgUUID, volUUID)
        log_str = logutils.volume_chain_to_str(vol.volUUID for vol in chain)
        self.log.info("chain=%s ", log_str)

        template = chain[0].getParentVolume()
        if template or len(chain) > 1:
            self.log.error("Importing and exporting an image with more "
                           "than one volume is not supported")
            raise se.CopyImageError()

        domain.activateVolumes(imgUUID, volUUIDs=[chain[0].volUUID])
        return chain[0]

    def upload(self, methodArgs, sdUUID, imgUUID, volUUID=None):
        domain = sdCache.produce(sdUUID)

        vol = self._activateVolumeForImportExport(domain, imgUUID, volUUID)
        try:
            imageSharing.upload(vol.getVolumePath(), methodArgs)
        finally:
            domain.deactivateImage(imgUUID)

    def download(self, methodArgs, sdUUID, imgUUID, volUUID=None):
        domain = sdCache.produce(sdUUID)

        vol = self._activateVolumeForImportExport(domain, imgUUID, volUUID)
        try:
            vol.extend(imageSharing.getSize(methodArgs))
            imageSharing.download(vol.getVolumePath(), methodArgs)
        finally:
            domain.deactivateImage(imgUUID)

    def copyFromImage(self, methodArgs, sdUUID, imgUUID, volUUID):
        domain = sdCache.produce(sdUUID)

        vol = self._activateVolumeForImportExport(domain, imgUUID, volUUID)
        try:
            imageSharing.copyFromImage(vol.getVolumePath(), methodArgs)
        finally:
            domain.deactivateImage(imgUUID)

    def copyToImage(self, methodArgs, sdUUID, imgUUID, volUUID=None):
        domain = sdCache.produce(sdUUID)

        vol = self._activateVolumeForImportExport(domain, imgUUID, volUUID)
        try:
            vol.extend(imageSharing.getLengthFromArgs(methodArgs))
            imageSharing.copyToImage(vol.getVolumePath(), methodArgs)
        finally:
            domain.deactivateImage(imgUUID)
