# Copyright 2020 VMware, Inc.
# All rights reserved. -- VMware Confidential
"""
Script to extend Hard disk of an existing VM
"""

import ssl
import logging
import utils
import sys
import os
import re
import json

sys.path.append(os.environ['VMWARE_PYTHON_PATH'])
from pyVmomi import vim  #pylint: disable=E0401,E0611
from pyVim.task import WaitForTask  #pylint: disable=E0401
from l10n import msgMetadata as _T, localizedString as _
from patch_errors import UserError

ssl.match_hostname = lambda cert, hostname: True
logger = logging.getLogger(__name__)

root = '/dev/sda'


def _find_root_vmdk(vm_obj):
    """
    This method scans the Hardware , uses 'lsscsi', map the SCSI controller
    information to Virtual Device Node to determine the Hard Disk of /dev/sda
    :param vm_obj: VCSA VM Object
    :param target_size: Expected root disk size in KBs
    :return: virtual_disk_device: Hard disk details or None
    """

    logger.info("Discovering root disk")
    try:
        cmd = 'lsscsi | grep '+ root
        rc, d, err = utils.run_command(cmd)
        if rc:
            raise Exception(err)
        m = re.match(r"^\[(.+)\].+$", d)
        scsi_data = m.group(1)
        logger.info("Root SCSI controller data: [%s]", scsi_data)
        scsci_split = scsi_data.split(':')
    except Exception:
        logger.error("Unable to find scsi id for root.")
        cause = _(_T('disk.partition.resize.discover.failed',
                     'Unable to find scsi id for root. '
                     'Cannot resize.'))
        raise UserError(cause=cause)

    virtual_disk_device = None
    for dev in vm_obj.config.hardware.device:
        if "SCSI Controller".lower() in dev.deviceInfo.label.lower():
            logger.info("dev.deviceInfo.label is [%s]", dev.deviceInfo.label)
            logger.info("dev.deviceInfo.summary iKeys [%s]", dev.deviceInfo)

            for scdev in vm_obj.config.hardware.device:
                if scdev.controllerKey == dev.key:
                    logger.info(" SCSI %s:%s - %s", \
                        dev.busNumber, scdev.unitNumber, \
                        scdev.deviceInfo.label)

                    if str(scsci_split[1]) == str(dev.busNumber) and str(
                            scsci_split[2]) == str(scdev.unitNumber):
                        logger.info(
                            "Root disk /dev/sda is mapped to vmdk %s", \
                            scdev.deviceInfo.label)
                        logger.info("Rook Disk device Info : [%s]", \
                                    scdev.deviceInfo)
                        virtual_disk_device = scdev
                        break
        if virtual_disk_device:
            break

    return virtual_disk_device


def _extend_root_disk(vm_obj, virtual_disk_device, target_size):
    """
    This method builds the device spec to resize the disk. Calls Reconfigure VM
    task to extend the disk.
    :param vm_obj: VCSA VM object
    :param virtual_disk_device: Virtual Device object of type Disk
    :param target_size: Expected root disk size in KBs
    :return: None
    """

    logger.info("Extending root disk")
    virtual_disk_spec = vim.vm.device.VirtualDeviceSpec()
    virtual_disk_spec.operation = \
        vim.vm.device.VirtualDeviceSpec.Operation.edit
    virtual_disk_spec.device = virtual_disk_device
    if target_size > 0:
        logger.info("Setting value of spec prop: "
                    "virtual_disk_spec.device.capacityInKB = %s ", target_size)
        virtual_disk_spec.device.capacityInKB = target_size

    dev_changes = list()
    dev_changes.append(virtual_disk_spec)
    spec = vim.vm.ConfigSpec()
    spec.deviceChange = dev_changes
    state = WaitForTask(vm_obj.ReconfigVM_Task(spec=spec))
    if state == "success":
        logger.info("Root disk: %s is extended to %s KB",
                    virtual_disk_device.deviceInfo.label, target_size)
    else:
        cause = _(_T('disk.resize.task.failed',
                     'Root disk resize failed due to %s.'), [state])
        raise UserError(cause=cause)


def _reconfig_root_partitions():
    """
    This method modifies the root partition table as per the current size of
    /dev/sda disk
    :return:
    """
    logger.info("Configuring root disk after resizing")

    try:
        scsi_dev = os.popen('lsscsi | grep ' + root)
        scsi_id = re.match(r"^\[(.+)\].+$", scsi_dev.read())
        os.system('echo 1 > "/sys/class/scsi_device/' + scsi_id.group(1) +
                  '/device/rescan"')
    except Exception:
        logger.error("Unable to find scsi id for %s.", root)
        cause = _(_T('disk.partition.resize.reconfig.failed',
                     'Unable to find scsi id for root. '
                     'Cannot resize.'))
        raise UserError(cause=cause)

    # Get current root size in blocks
    sz_root = os.popen('blockdev --getsz ' + root)
    size_new = int(sz_root.read())
    logger.info("%s new size is %s", root, size_new)

    # Read Partition table in JSON format
    sfdisk_json = os.popen('sfdisk -J ' + root)
    sfdisk = json.loads(sfdisk_json.read())
    logger.info("Current Partition table, %s", sfdisk)

    firstlba = sfdisk['partitiontable']['firstlba']
    partitions = sfdisk['partitiontable']['partitions']

    sda3 = [partition for partition in partitions
            if partition['node'] == root + '3'][0]
    sda3Start = sda3['start']
    sda3Size = sda3['size']

    sda4 = [partition for partition in partitions
            if partition['node'] == root + '4'][0]
    sda4Start = sda4['start']
    sda4Size = sda4['size']

    if sda4Start == (sda3Start + sda3Size):
        # In Photon 1 or inplace updated system to Photon 3, /dev/sda4
        # is the last partition
        # Mapping Table according to address ordering
        # /dev/sda1 - boot partition
        # /dev/sda2 - SWAP partition
        # /dev/sda3 - root partition
        # /dev/sda4 - BIOS boot
        #
        # Applied logic adjust lastLba, /dev/sda4 start address, /dev/sda3 size

        lastlbaNew = size_new - firstlba
        sda4StartNew = lastlbaNew - sda4Size + 1
        sda3SizeNew = sda4StartNew - sda3Start

        logger.info("Calculated values : LastLba = %s, /dev/sda4 start = %s, "
                    "/dev/sda3 size = %s", lastlbaNew, \
                                             sda4StartNew, \
                                             sda3SizeNew)

        sfdisk['partitiontable']['lastlba'] = lastlbaNew
        for partition in sfdisk['partitiontable']['partitions']:
            if partition['node'] == root + '3':
                partition['size'] = sda3SizeNew
            elif partition['node'] == root + '4':
                partition['start'] = sda4StartNew

    elif sda3Start == (sda4Start + sda4Size):
        # In newly deployed Photon 3 systems, /dev/sda3 is the last partition
        # Mapping Table according to address ordering
        # /dev/sda1 - BIOS boot
        # /dev/sda2 - boot partition
        # /dev/sda4 - swap partition
        # /dev/sda3 - root partition
        #
        # Applied logic adjust lastLba and /dev/sda3 size

        lastlbaNew = size_new - firstlba
        sda3SizeNew = lastlbaNew - sda3Start + 1
        logger.info("Calculated values : LastLba = %s, /dev/sda3 size = %s", \
              lastlbaNew, sda3SizeNew)

        sfdisk['partitiontable']['lastlba'] = lastlbaNew
        for partition in sfdisk['partitiontable']['partitions']:
            if partition['node'] == root + '3':
                partition['size'] = sda3SizeNew

    else:
        msg = "Unknown partition configuration found. " \
              "Root resize not supported."
        logger.error(msg)
        cause = _(_T('disk.partition.resize.task.failed',
                     'Unknown partition configuration found. '
                     'Root resize not supported.'))
        raise UserError(cause=cause)

    with open('.disk_layout', 'w') as f:
        f.write("label: %s\n" % sfdisk['partitiontable']['label'])
        f.write("label-id: %s\n" % sfdisk['partitiontable']['id'])
        f.write("device: %s\n" % sfdisk['partitiontable']['device'])
        f.write("unit: %s\n" % sfdisk['partitiontable']['unit'])
        f.write("first-lba: %d\n" % sfdisk['partitiontable']['firstlba'])
        f.write("last-lba: %d\n\n" % sfdisk['partitiontable']['lastlba'])
        for partition in sfdisk['partitiontable']['partitions']:
            f.write("%s : start=%12.d, size=%12.d, type=%s, uuid=%s\n" %
                    (partition['node'], partition['start'], partition['size'],
                     partition['type'], partition['uuid']))

    tmp = os.popen('sfdisk ' + root + ' --force < .disk_layout')
    tmp.read()
    tmp = os.popen('rm .disk_layout')
    tmp.read()
    tmp = os.popen('partprobe')
    tmp.read()
    tmp = os.popen('resize2fs ' + root + '3')
    tmp.read()
    tmp = os.popen('grub2-install ' + root)
    tmp.read()

    # Read Partition table in JSON format
    sfdisk_json = os.popen('sfdisk -J ' + root)
    sfdisk = json.loads(sfdisk_json.read())
    logger.info("Partition table post configuration, %s", sfdisk)


def resize_root_disk(vm_obj, target_size):
    """
    This method resizes the root disk to target_size.
    target_size must be higher than existing root disk size
    :param vm_obj: VCSA VM object
    :param target_size: Expected root disk size in KBs
    :return: None
    """

    source_size = utils.get_existing_root_size()
    if source_size == 0:
        raise RuntimeError('Unable to find root size of existing VCSA')
    if int(target_size) <= int(source_size):
        logger.info('Existing root size is already higher. '
                    'Resize not required.')
        return

    # Find the root disk device
    virtual_disk_device = _find_root_vmdk(vm_obj)
    if not virtual_disk_device:
        raise RuntimeError("Unable to find root disk's vmdk mapping.")

    stripped_size = \
        virtual_disk_device.deviceInfo.summary.replace(',', '').split(' ')[0]
    logger.info("Root disk with size %s KB found. Expected "
                "%s KB.", stripped_size, target_size)
    if int(stripped_size) < int(target_size):
        logger.info("Resize required.")
    else:
        logger.info("Root resize not required")
        return

    # Root Disk is found, extend the disk
    _extend_root_disk(vm_obj, virtual_disk_device, target_size)
    # Disk is extended resize root partition
    _reconfig_root_partitions()
    logger.info("Root disk /dev/sda re-sized successfully")

