#!/bin/env python3
from MySQLdb import _mysql
import argparse
import socket
import sys
import time
import xml.etree.ElementTree as ET
import yaml

# TODO: filter by state
state_map = {
    'run':       3,
    'runn':      3,
    'running':   3,

    'poff':      8,
    'poweroff':  8,

    'done':      6,

    'suspended': 5,
    'susp':      5,
    'sus':       5,
}


def get_dict_by_keypath(idict, keypath=None):
    if keypath is None:
        return idict
    tmp = idict
    keys = keypath.split(':')
    for k in keys:
        if k in tmp:
            tmp = tmp[k]
        else:
            return {}
    return tmp

def load_config(path='/etc/one/axc_driver.conf', keypath=None):
    cfg = {}
    try:
      with open(path, 'r') as f:
        try:
            cfg = yaml.safe_load(f)
        except yaml.YAMLError as e:
            return cfg
    except FileNotFoundError as e:
        return cfg
    return get_dict_by_keypath(cfg, keypath)

def merge_with_subpath(idict, subpath=None):
    return {**idict, **get_dict_by_keypath(idict, subpath)}

cfg = load_config(keypath='mysql')

db=_mysql.connect(db=cfg['db'], passwd=cfg['password'])

def get_vm_xml(vmid: str):
    vmid = int(vmid)
    db.query(f"SELECT body from vm_pool where oid={vmid}")
    result = db.store_result()
    (xmlbody,) = result.fetch_row(maxrows=1)[0]
    root = ET.fromstring(xmlbody)
    return root

def get_disk(xmlroot, id=0):
    disks = xmlroot.findall('./TEMPLATE/DISK')
    for disk in disks:
        if id == int(disk.find('DISK_ID').text):
            return disk
    return None


def get_disk_count(xmlroot):
    disks = xmlroot.findall('./TEMPLATE/DISK')
    return len(disks)

def get_disk_field(xmlroot, field):
    ret = []
    disks = xmlroot.findall('./TEMPLATE/DISK')
    for disk in disks:
        if disk.find('CLONE') is None or 'YES' == disk.find('CLONE').text:
            if disk.find(field) is None:
                ret.append("iqn.2019-07.net.efscloud.slc:sys")
            else:
                ret.append(disk.find(field).text)
    return ret


def find_vms(id):
    id = db.escape_string(str(id)).decode()
    db.query(f'SELECT oid,name from vm_pool where (oid="{id}" OR `name`="{id}") AND state <> 6')
    r = db.store_result()
    return r


def parse_args():
    parser = argparse.ArgumentParser( prog = 'oness',
                                      description = 'Filter OpenNebula VMs by Storage')
    parser.add_argument('-c', '--column', nargs='?', type=str, choices=["id", "name", "all"], default="name", help="Column to output (default: 'name')")
    parser.add_argument('vm', type=str, help="VM id or name")
    parser.add_argument('field', type=str, nargs='?', default="DISK_ID", help="Field of disk XML to extract (Default 'DISK_ID')")
    return parser.parse_args()

def main():
    args = parse_args()
    vm = int(args.vm)
    field = args.field
    result = find_vms(vm)
    for oid, name in result.fetch_row(maxrows=0):
        x = get_vm_xml(oid)
        disks = get_disk_field(x, field)
        print(" ".join(disks))


if __name__ == '__main__':
    main()
