# Copyright 2015 TrilioData Inc.
# All Rights Reserved.

import copy
import os
import uuid
import json
import time
import rbd
import rados
import re

from tempfile import mkstemp

try:
    from oslo_log import log as logging
except ImportError:
    from nova.openstack.common import log as logging

try:
    from oslo_utils import encodeutils
except ImportError:
    from nova.openstack.common import strutils as encodeutils

try:
    from oslo_config import cfg
except ImportError:
    from oslo.config import cfg

try:
    from oslo_utils import fileutils
except ImportError:
    from nova.openstack.common import fileutils

from keystoneclient.auth.identity import v2 as v2_auth
from keystoneclient import session

from nova import utils

try:
    import nova.virt.libvirt.rbd_utils as rbd_utils
except BaseException:
    from nova.virt.libvirt.storage import rbd_utils

import nova.virt.libvirt.utils as libvirt_utils
import nova.virt.libvirt.imagebackend as imagebackend
from contego.nova.extension.driver import qemuimages
from contego.nova.extension.driver import vault
from contego import utils as contego_utils

import rbd_base


LOG = logging.getLogger(__name__)
CONF = cfg.CONF


class RBDBootBackend(rbd_base.RBDBaseBackend):
    key_file = None
    key_user = None

    def __init__(self, **kwargs):
        super(RBDBootBackend, self).__init__(**kwargs)
        pass

    def create_snapshot(self, devices, **kwargs):
        """Creates an rbd snapshot."""
        disks_info = []
        for device in devices:
            disk_type = device.get('type')
            if disk_type == 'network':
                backend = device.find('source').get('protocol')
                disk_path = backend + ':' + device.find('source').get('name')
            elif disk_type == 'volume':
                backend = device.find('source').get('pool')
                disk_path = backend + ':' + device.find('source').get('volume')

            volume_name = disk_path.split('/')[1]
            pool_name = disk_path.split(':')[1].split('/')[0]
            with rbd_utils.RBDVolumeProxy(self.driver,
                                          volume_name,
                                          pool=pool_name) as volume:
                snapshot_name = "triliovault-" + str(uuid.uuid4())
                snap = encodeutils.safe_encode(snapshot_name)
                volume.create_snap(snap)

                try:
                    volume.protect_snap(snap)
                except rbd.FunctionNotSupported:
                    pass

                vsize = self._size(disk_path.split(':')[1])
                disk_info = {'dev': device.find('target').get('dev'),
                             'type': disk_type,
                             'path': disk_path,
                             'size': vsize,
                             'backend': backend + 'boot',
                             'backings': [{'path': disk_path, 'size': vsize}],
                             'snapshot_name': snapshot_name,
                             'volume_id': None}
                disks_info.append(disk_info)

        return disks_info

    def delete_snapshot(self, disk_info, **kwargs):
        """Deletes an rbd snapshot."""
        volume_name = disk_info['path'].split('/')[1]
        pool_name = disk_info['path'].split(':')[1].split('/')[0]
        volume_name = encodeutils.safe_encode(volume_name)

        # delete all snapshots except the current
        curr_snapshot_name = encodeutils.safe_encode(
            disk_info['snapshot_name'])
        with rbd_utils.RBDVolumeProxy(self.driver,
                                      volume_name,
                                      pool=pool_name) as volume:

            if kwargs['params']['workload_failed']:
                all_snaps_minus_curr = set(
                    [s['name'] for s in volume.list_snaps()
                     if 'triliovault-' in s['name']])
            else:
                all_snaps_minus_curr = set(
                    [s['name'] for s in volume.list_snaps()
                     if s['name'] != curr_snapshot_name and
                     'triliovault-' in s['name']])

            for snap_to_delete in all_snaps_minus_curr:
                try:
                    volume.unprotect_snap(snap_to_delete)
                except Exception as ex:
                    LOG.exception(ex)
                    pass

                try:
                    volume.remove_snap(snap_to_delete)
                except Exception as ex:
                    LOG.exception(ex)
                    pass

    def check_prev_snapshot(self, disk_info, **kwargs):

        status = {'result': 'invalid'}

        volume_name = disk_info['path'].split('/')[1]
        pool_name = disk_info['path'].split(':')[1].split('/')[0]
        volume_name = encodeutils.safe_encode(volume_name)

        if 'prev_disk_info' in disk_info and disk_info['prev_disk_info']:
            prev_snapshot_name = encodeutils.safe_encode(
                disk_info['prev_disk_info']['snapshot_name'])
            with rbd_utils.RBDVolumeProxy(
                self.driver, volume_name, pool=pool_name) as volume:  # noqa

                allsnaps = set([s['name'] for s in volume.list_snaps()])
                if prev_snapshot_name in allsnaps:
                    status = {'result': 'success'}

        return status

    def upload_snapshot(self, disk_info, **kwargs):
        context = kwargs['context']
        params = kwargs['params']

        snapshot_full_name = disk_info['path'].split(':')[1] + "@" + \
            disk_info['snapshot_name']
        volume_name = disk_info['path'].split('/')[1]
        pool_name = disk_info['path'].split(':')[1].split('/')[0]
        volume_name = encodeutils.safe_encode(volume_name)
        progress_tracker_path = \
            vault.get_progress_tracker_path(params['metadata'])

        # prev_snapshot_name should be available in disk_info:
        if 'prev_disk_info' in disk_info and disk_info['prev_disk_info']:
            prev_snapshot_name = encodeutils.safe_encode(
                disk_info['prev_disk_info']['snapshot_name'])

            # verify if the snapshot still exists, otherwise default
            # to full snapshot
            with rbd_utils.RBDVolumeProxy(
                    self.driver, volume_name, pool=pool_name) as volume:
                if prev_snapshot_name not in set(
                        [s['name'] for s in volume.list_snaps()]):
                    prev_snapshot_name = None
        else:
            prev_snapshot_name = None
            # this is a full snapshot request.
            # Delete all other previous snapshots
            # triliovault created.
            with rbd_utils.RBDVolumeProxy(
                self.driver, volume_name, pool=pool_name) as volume:  # noqa
                for snap in volume.list_snaps():
                    if 'triliovault' not in snap['name']:
                        continue

                    if snap['name'] in snapshot_full_name:
                        continue

                    try:
                        volume.unprotect_snap(snap['name'])
                    except rbd.FunctionNotSupported:
                        pass

                    try:
                        volume.remove_snap(snap['name'])
                    except rbd.ImageBusy:
                        raise Exception("Snapshot %s is busy" % snap['name'])

        # diffsize = self.diff_size(snapshot_full_name,
                        # from_snap=prev_snapshot_name)
        vsize = self._size(disk_info['path'].split(':')[1])

        prev_devname = None
        extentsfile = None
        qcow2file = None

        try:
            extentsfile = self._get_temp_file_path(True)

            if CONF.vault_storage_type.lower() in ('nfs', 'swift-s', 's3'):
                qcow2file = vault.get_snapshot_vm_disk_resource_path(
                     params['metadata'])  # noqa
                head, tail = os.path.split(qcow2file)
                fileutils.ensure_tree(head)
            else:
                qcow2file = self._get_temp_file_path(True)

            self._export_diff(snapshot_full_name, extentsfile,
                              from_snap=prev_snapshot_name)
            try:
                libvirt_utils.create_cow_image(prev_devname, qcow2file, vsize)
            except:
                # s3 sometimes may throw exception in create_cow_images
                # due its eventual consistency nature
                # give sometime for S3 to propagate the changes
                time.sleep(30)
                qemuimages.qemu_img_info(qcow2file)

            if prev_snapshot_name:
                backing_image = disk_info['prev_disk_info']['path'] + '@' +\
                    prev_snapshot_name
                backing_image = backing_image.split('rbd:')[1]
            else:
                backing_image = pool_name + '/' + volume_name
            self._create_qcow2_from_extents(qcow2file, extentsfile,
                                            backing_image, snapshot_full_name,
                                            progress_tracker_path)
        finally:
            if extentsfile and os.path.exists(extentsfile):
                os.remove(extentsfile)

        # create a new token with new expiration date to upload the image
        if qcow2file != vault.get_snapshot_vm_disk_resource_path(
               params['metadata']):  # noqa
            newcontext = rbd_base.get_new_context(context)
            vault.upload_snapshot_vm_disk_resource(
                newcontext, params['metadata'], qcow2file)
            os.remove(qcow2file)

    def reset_snapshot(self, devices, **kwargs):
        disks_info = []
        if kwargs['instance_ref']['deleted']:
            volume_name = kwargs['instance_uuid'] + '_disk'
            pool_name = CONF.libvirt.images_rbd_pool
            volume_name = encodeutils.safe_encode(volume_name)
            with rbd_utils.RBDVolumeProxy(
                self.driver, volume_name, pool=pool_name) as volume:  # noqa
                for snap in volume.list_snaps():
                    if 'triliovault' not in snap['name']:
                        continue

                    try:
                        volume.unprotect_snap(snap['name'])
                    except rbd.FunctionNotSupported:
                        pass

                    try:
                        volume.remove_snap(snap['name'])
                    except Exception:
                        raise Exception("Snapshot %s is busy" % snap['name'])
            with rbd_utils.RADOSClient(self.driver, pool=pool_name) as client:
                try:
                    rbd.RBD().remove(client.ioctx, volume_name)
                except Exception:
                    raise Exception("Volume %s is busy" % volume_name)
        else:
            for device in devices:
                disk_type = device.get('type')
                if disk_type == 'network':
                    backend = device.find('source').get('protocol')
                    disk_path = backend + ':' + device.find('source').get('name')
                elif disk_type == 'volume':
                    backend = device.find('source').get('pool')
                    disk_path = backend + ':' + device.find('source').get('volume')

                volume_name = disk_path.split('/')[1]
                pool_name = disk_path.split(':')[1].split('/')[0]
                volume_name = encodeutils.safe_encode(volume_name)
                with rbd_utils.RBDVolumeProxy(
                    self.driver, volume_name, pool=pool_name) as volume:  # noqa
                    for snap in volume.list_snaps():
                        if 'triliovault' not in snap['name']:
                            continue

                        try:
                            volume.unprotect_snap(snap['name'])
                        except rbd.FunctionNotSupported:
                            pass

                        try:
                            volume.remove_snap(snap['name'])
                        except Exception:
                            raise Exception("Snapshot %s is busy" % snap['name'])

        return disks_info

    def transfer_qemu_image_to_volume(self, rbd_path, backup_image_file_path,
                                      progress_tracking_file_path):

        def ceph_rename_volume(source, target):
            args = [source]
            args += [target]
            kwargs = {'run_as_root': True}
            out, err = self.rbd_keyring_search_and_execute(
                'mv', *args, **kwargs)

            if err != '':
                raise Exception("Cannot run rbd rename %s to %s. Error %s" %
                                (source, target, err))

            return

        def ceph_delete_volume(volume_name):
            args = [volume_name]
            kwargs = {'run_as_root': True}
            out, err = self.rbd_keyring_search_and_execute(
                'rm', *args, **kwargs)

            return

        def ceph_volume_info(volume_name):
            args = [volume_name]
            args += ['--format=json']
            kwargs = {'run_as_root': True}
            out, err = self.rbd_keyring_search_and_execute(
                'info', *args, **kwargs)

            if err != '':
                raise Exception("Cannot run rbd info on %s. Error %s" %
                                (volume_name, err))

            return json.loads(out)

        def create_volume_from_file(volume_name, backup_image_file_path,
                                    ceph_conf_file):

            volume_endpoint = 'rbd:' + volume_name
            volume_endpoint += ':id=' + self.key_user.split('.')[1]
            volume_endpoint += ':conf=' + ceph_conf_file

            cmdspec = ['qemu-img', 'convert', '-p', '-t', 'none', '-O',
                       'raw', backup_image_file_path, volume_endpoint]

            cmd = " ".join(cmdspec)
            LOG.debug(('transfer_qemu_image_to_volume cmd %s ' % cmd))

            self._execute_qemu_img_and_track_progress(
                     cmdspec, volume_name, backup_image_file_path,
                     progress_tracking_file_path)

        # import image to rbd volume

        # load the self.key_user
        pool_name = rbd_path.split('/')[0]
        args = [pool_name]
        kwargs = {}
        self.rbd_keyring_search_and_execute('ls', *args, **kwargs)

        if not self.key_user:
            raise Exception(
                "Could not find any valid ceph key to use. "
                "Please make sure 'rbd ls -l' can be run successfully")

        temp_volume = pool_name + '/' + "boot-" + \
            str(uuid.uuid4().hex) + "-TrilioVault"

        with rbd_base.make_copy_ceph_conf(self) as temp_path:
            try:
                create_volume_from_file(temp_volume, backup_image_file_path,
                                        temp_path)
            except Exception as ex:
                LOG.exception(ex)
                time.sleep(30)
                ceph_delete_volume(temp_volume)
                raise

        # delete volume created by cinder
        rbd_boot_name = rbd_path
        ceph_delete_volume(rbd_boot_name)

        # rename the other volume to cinder created volume name
        ceph_rename_volume(temp_volume, rbd_boot_name)

        # Get the info on the newly restored volume
        statinfo = ceph_volume_info(rbd_boot_name)

        LOG.debug("Transferred backup image to ceph volume %s. "
                  "Transferred size %s" %
                  (rbd_path, statinfo.get('object_size', "Not Found")))

        return

    def copy_backup_image_to_volume(self, context, instance_uuid,
                                    instance_name, params):
        rbd_path = params['path']
        backup_path = params['image_overlay_file_path']
        progress_tracking_file_path = params['progress_tracking_file_path']

        self.transfer_qemu_image_to_volume(rbd_path, backup_path,
                                           progress_tracking_file_path)


"""
driver = ContegoRBDDriver("volumes", "/etc/ceph/ceph.conf", "cinder")
driver.create_snapshot(snapshot)
fileh, extentsfile = mkstemp()
close(fileh)
remove(extentsfile)
fileh, qcow2file = mkstemp()
close(fileh)
remove(qcow2file)
driver.export_diff("volumes/"+snapshot['volume_name']+"@"+snapshot['name'],
                   extentsfile)
vsize = driver.volume_size(snapshot['volume_name'])
libutils_utils.create_image("qcow2", qcow2file, vsize)
driver.create_qcow2_from_extents(qcow2file, extentsfile)
self.delete_snapshot(snapshot)
"""