#!/usr/bin/python3
import subprocess
import json
import os

class SilentException(Exception):
    pass

def list_slots():
    slots = {}
    current = None
    dump = subprocess.check_output(['smbiosDump'])
    if not isinstance(dump, str):
        dump = dump.decode('utf8')
    for line in dump.split('\n'):
        if 'System Slot (Type 9):' in line:
            current = {}
        elif current is not None:
            stripped = line.strip()
            if stripped.startswith('Designation:'):
                current['designation'] = stripped.split(':', 1)[1].strip().strip('"')
            if stripped.startswith('Type:'):
                current['type_label'] = stripped.split(':', 1)[1].strip()
                current['type_number'] = int(stripped.split(':', 1)[1].strip().split(' ')[0], 16)
                # 0x14 through 0x17 are M.2 slots, inclusively
                if current['type_number'] in (0x14, 0x15, 0x16, 0x17):
                    current['is_m2'] = True
                else:
                    current['is_m2'] = False
            elif stripped.startswith('PCI Address:'):
                addr = stripped.split(':', 1)[1].strip()
                current['pci_address'] = addr
                slots[addr] = current
                current = None
    return slots

def list_m2_slots():
    slots = list_slots()
    m2slots = {}
    for addr in slots:
        slot = slots[addr]
        if slot.get('is_m2', False):
            m2slots[addr] = slot
    return m2slots

def get_deployment_storage():
    apiclient = None
    for clientpath in ('/opt/confluent/bin/apiclient', '/etc/confluent/apiclient'):
        if os.path.exists(clientpath):
            apiclient = clientpath
            break
    if not apiclient:
        return None
    for attriburl in ('/confluent-api/self/myattrib', '/confluent-api/self/myattribs'):
        try:
            attribs = subprocess.check_output(
                ['python3', apiclient, attriburl, '-j'], stderr=subprocess.DEVNULL)
            if not isinstance(attribs, str):
                attribs = attribs.decode('utf8')
            attribs = json.loads(attribs)
            if not isinstance(attribs, dict):
                continue
            storage = attribs.get('deployment.storage', None)
            if isinstance(storage, str) and storage.strip():
                return storage.strip()
        except Exception:
            continue
    return None

# List of models that are generally used for OS install disk (generally 2-disk mirroring)
PRIORITY_MODELS = ('m.2 nvme 2-bay raid kit', 
                   'thinksystem_m.2_vd',
                   'thinksystem m.2',
                   'thinksystem_m.2',
                   'thinksystem 7mm',
                   'thinksystem_7mm',
                   'thinksystem_7mm_vd',                  
                   'b540p-2hs',  # 7mm or M.2 in rear
                   'b540d-2hs',  # M.2 in E3.S or 2.5inch
                   'b550d-2hs',  # Front M.2
                   'b550p-2hs',  # Rear M.2
                   'raid b540p-2hs',
                   'raid b540d-2hs',
                   'raid b550d-2hs',
                   'raid b550p-2hs',
                   'raid b540i-2i',
                   'raid_b540i-2i',
                   'raid b545i-2i',
                   'raid b545i-2i',
                   'raid_b545-2i')

class DiskInfo(object):
    def __init__(self, devname, devinfo, storagespec=None):
        self.name = devname
        self.path = '/dev/' + devname
        self.wwn = None
        self.model = devinfo.get('model', 'Unknown')
        self.driver = devinfo.get('adapter_driver', 'Unknown')
        self.size = devinfo.get('size', 0)  # in MiB
        self.m2_slot = devinfo.get('m2_slot', None)
        if not devinfo.get('is_local', False):
            raise SilentException("Not local")
        if devinfo.get('is_removable', False):
            raise SilentException("Removable")
        if devinfo.get('is_usb', False):
            raise SilentException("USB device")
        if devinfo.get('type', '').lower() in ('cd-rom',):
            raise SilentException("CD-ROM device")
        if self.size < 2048:
            raise SilentException("Too small")
        if storagespec and storagespec == 'm2':
            if self.model.lower() not in PRIORITY_MODELS and not devinfo.get('m2_slot', None):
                raise SilentException("Not an M.2 disk")
               
            


    @property
    def priority(self):
        if self.m2_slot:
            return 0
        if self.model.lower() in PRIORITY_MODELS:
            return 0
        if self.driver == 'vmw_ahci':
            return 2
        if self.driver == 'nvme_pcie':
            return 3
        return 99

    def __repr__(self):
        return repr({
            'name': self.name,
            'path': self.path,
            'wwn': self.wwn,
            'driver': self.driver,
            'size': self.size,
            'model': self.model,
        })

def list_disks():
    current_dev = None
    m2slots = list_m2_slots()
    disks = {}
    devlist = subprocess.check_output(['localcli', 'storage', 'core', 'device', 'list'])
    if not isinstance(devlist, str):
        devlist = devlist.decode('utf8')
    devbyadp = {}
    for line in devlist.split('\n'):
        if not line.strip():
            continue
        if not line.startswith(' '):
            current_dev = line.rsplit(':', 1)[0]
            if current_dev not in disks:
                disks[current_dev] = {}
        elif current_dev:
            if '   Model:' in line:
                disks[current_dev]['model'] = ' '.join(line.split()[1:])
            elif '   Driver:' in line:
                disks[current_dev]['driver'] = ' '.join(line.split()[1:])
            elif '   Is Local:' in line:
                disks[current_dev]['is_local'] = ' '.join(line.split()[2:]).lower() == 'true'
            elif '   Is Removable:' in line:
                disks[current_dev]['is_removable'] = ' '.join(line.split()[2:]).lower() == 'true'
            elif '   Size:' in line:  # in MiB
                disks[current_dev]['size'] = int(line.split()[1])
            elif '   Is SSD:' in line:
                disks[current_dev]['is_ssd'] = ' '.join(line.split()[2:]).lower() == 'true'
            elif '   Is USB:' in line:
                disks[current_dev]['is_usb'] = ' '.join(line.split()[2:]).lower() == 'true'
            elif '   Is Removable:' in line:
                disks[current_dev]['is_removable'] = ' '.join(line.split()[2:]).lower() == 'true'
            elif 'Device Type:' in line:
                disks[current_dev]['type'] = ' '.join(line.split()[2:])
    for dev in disks:
        pathlist = subprocess.check_output(['localcli', 'storage', 'core', 'path', 'list', '--device', dev])
        if not isinstance(pathlist, str):
            pathlist = pathlist.decode('utf8')
        for line in pathlist.split('\n'):
            if not line.strip():
                continue
            if not line.startswith(' '):
                continue
            if '   Adapter Identifier:' in line:
                adpname = ' '.join(line.split()[2:])
                disks[dev]['adapter_id'] = adpname
            elif '   Adapter:' in line:
                adp = ' '.join(line.split()[1:])
                disks[dev]['adapter'] = adp
                devbyadp.setdefault(adp, []).append(dev)
    adapterlist = subprocess.check_output(['localcli', 'storage', 'core', 'adapter', 'list'])
    if not isinstance(adapterlist, str):
        adapterlist = adapterlist.decode('utf8')
    adapters = {}
    curradpinfo = {}
    headline = None
    headerdone = False
    descidx = 0
    for line in adapterlist.split('\n'):
        if not line.strip():
            continue
        if not headline:
            headline = line
            descidx = headline.find('Description')
            if descidx < 0:
                raise Exception("Unexpected adapter list format: {0}".format(headline))
        elif not headerdone:
            if line.startswith('-'):
                headerdone = True
            continue
        fields = line.split(maxsplit=2)
        if len(fields) < 2:
            raise Exception("Unexpected adapter list format: {0}".format(line))
        currname, currdriver = fields[0], fields[1]
        currdesc = line[descidx:].strip()
        pcieaddr = None
        if currdesc.startswith('('):
            pcieaddr = currdesc.split('(')[1].split(')')[0]
        curradpinfo = {
            'name': currname,
            'driver': currdriver,
            'description': currdesc,
            'pcieaddress': pcieaddr
        }
        adapters[currname] = curradpinfo
    for adp in devbyadp:
        driver = adapters.get(adp, {}).get('driver', 'Unknown')
        pcieaddr = adapters.get(adp, {}).get('pcieaddress', None)
        for dev in devbyadp[adp]:
            disks[dev]['adapter_driver'] = driver
            disks[dev]['adapter_pcieaddress'] = pcieaddr
            if pcieaddr and pcieaddr in m2slots:
                disks[dev]['m2_slot'] = m2slots[pcieaddr]
    return disks
        
def main():
    disks = []
    storagespec = get_deployment_storage()
    try:
        alldisks = list_disks()
    except Exception as e:
        print("Error listing disks: {0}".format(str(e)))
        alldisks = {}
    for disk in alldisks:
        try:
            disks.append(DiskInfo(disk, alldisks[disk], storagespec=storagespec))
        except SilentException:
            pass
        except Exception as e:
            print("Skipping {0}: {1}".format(disk, str(e)))
    nd = [x.name for x in sorted(disks, key=lambda x: [x.priority, x.size])]
    if nd:
        with open('/tmp/storagecfg', 'w') as sc:
            sc.write(f'clearpart --all --drives={nd[0]} --overwritevmfs\n')
            sc.write(f'install --drive={nd[0]} --overwritevmfs\n')
    else:
        with open('/tmp/storagecfg', 'w') as sc:
            sc.write('clearpart --firstdisk --overwritevmfs\n')
            sc.write('install --firstdisk --overwritevmfs\n')


if __name__ == '__main__':
    main()
