#/usr/bin/env python3
import sys
import libvirt
import time
import re
import os
import requests
import psutil
from socket import gethostbyaddr
from xml.etree import ElementTree

POOL_CACHE_TTL = 3600

pool_cache = None  # set from parent process

def get_one_pool(dc, vmid, dev):
    key = f'{dc}:{vmid}:{dev}'
    if pool_cache is not None and key in pool_cache:
        ts, val = pool_cache[key]
        if time.time() - ts < POOL_CACHE_TTL:
            return val

    url = f'https://vmm.{dc}.efscloud.net/zzinfo/api/v1/vm/e/{vmid}'
    r = requests.get(url)
    if r.status_code == 200:
        vms = r.json()
        for vm in vms:
            for disk in vm.get('disks', []):
                if disk.get('target') == dev:
                    val = disk.get('pool')
                    if pool_cache is not None:
                        pool_cache[key] = (time.time(), val)
                    return val
    return None


def metric(name, type, value, tags):
    result = {}
    metrictags = {}
    metrictags["name"] = name
    metrictags["atlas.dstype"] = type
    result["timestamp"] = int(time.time())*1000
    result["value"] = value
    resulttags = {**metrictags, **tags}
    result["tags"] = resulttags
    return result

def get_blk_tags(deviceXML, current_tags, dc):
    ret = {}

    # Fill target
    ret["dev"] = deviceXML.find('target').get('dev')

    # Fill source & storage & pool
    srcXML = deviceXML.find('source')
    if deviceXML.get('type') == 'file' and not deviceXML.get('device') == 'cdrom':
        source = srcXML.get('file')
        ret["source"] = source.split('/').pop()
        m = re.match(r'^/nfs/([^/]+)/', source)
        if m:
            ret["storage"] = m.group(1).split('.')[0]
            ret["pool"] = m.group(1)
        else:
            ret["storage"] = "localhost"
    elif deviceXML.get('type') == 'network' and srcXML.get('protocol') == 'iscsi':
        source = srcXML.get('name')
        m = re.match(r'^.+:([^/]+)', source)
        if m:
            ret["source"] = m.group(1)
        else:
            ret["source"] = source
        storage = srcXML.find('host').get("name")
        if re.match(r'^\d+\.\d+\.\d+\.\d+$', storage):
            storage = str(gethostbyaddr(storage)[0])
        ret["storage"] = storage.split('.')[0]
        if 'one.id' in current_tags:
            one_id = current_tags['one.id']
            pool = get_one_pool(dc, one_id, ret["dev"])
            if pool is not None:
                ret["pool"] = pool
        #   r = requests.get(f'https://vmm.{dc}.efscloud.net/{one_id}/{ret["dev"]}')
        #   if r.status_code == 200:
        #     ret["pool"] = r.text
    return ret


def print_blk_stat(deviceXML, domain):
    ret = {}
    tags = get_blk_tags(deviceXML)
    ret["rdreq"], ret["rdbytes"], ret["wrreq"], ret["wrbytes"], ret["err"] = domain.blockStats(tags["dev"])
    for k in "rdreq", "rdbytes", "wrreq", "wrbytes":
        print(metric("libvirt.block.{}".format(k), "counter", ret[k], tags))
    return ret

def get_blk_metrics(deviceXML, domain, tags, dc):
    tmp = {}
    localtags = get_blk_tags(deviceXML, tags, dc)
#    flags = libvirt.VIR_DOMAIN_BLOCK_STATS_READ_TOTAL_TIMES & libvirt.VIR_DOMAIN_BLOCK_STATS_WRITE_TOTAL_TIMES
    x = domain.blockStatsFlags(localtags["dev"])
    for k in x:
        v = x[k]
        if k.endswith('times'):
            v = v / 1e6
#        print(metric("libvirt.block.{}".format(k), "counter", x[k], {**tags, **localtags}))
        yield metric("libvirt.block.{}".format(k), "counter", v, {**tags, **localtags})
#    print(x)
#    tmp["rdreq"], tmp["rdbytes"], tmp["wrreq"], tmp["wrbytes"], tmp["err"] = domain.blockStats(localtags["dev"])
#    for k in "rdreq", "rdbytes", "wrreq", "wrbytes":
#        yield metric("libvirt.block.{}".format(k), "counter", tmp[k], {**tags, **localtags})



def get_net_tags(deviceXML):
    ret = {}

    # Fill target
    ret["dev"] = deviceXML.find('target').get('dev')

    # Fill source & storage & pool
    srcXML = deviceXML.find('source')
    if deviceXML.get('type') == 'bridge':
        ret["source"] = srcXML.get('bridge')
    return ret


def get_net_metrics(deviceXML, domain, tags):
    tmp = {}
    localtags = get_net_tags(deviceXML)
    tmp["rdbytes"], tmp["rdpackets"], tmp["rderrs"], tmp["rddrops"],\
      tmp["wrbytes"], tmp["wrpackets"], tmp["wrerrs"], tmp["wrdrops"] = domain.interfaceStats(localtags["dev"])
    for k in "rdbytes", "rdpackets", "rderrs", "rddrops", "wrbytes", "wrpackets", "wrerrs", "wrdrops":
        yield metric("libvirt.net.{}".format(k), "counter", tmp[k], {**tags, **localtags})


# tmp = domain.getCPUStats(True)[0]

def get_cpu_metrics_new(deviceXML, domain, conn, tags):
    flags = libvirt.VIR_CONNECT_GET_ALL_DOMAINS_STATS_ACTIVE
    stats = libvirt.VIR_DOMAIN_STATS_VCPU
    tmp = conn.domainListGetStats([domain], stats, flags)[0][1]
    vcount = tmp['vcpu.current']
    total = 0
    for i in range(vcount):
        total = total + tmp['vcpu.{}.time'.format(i)]
    return metric("libvirt.cpu.cpu_time", "counter", total/1e9, tags)


def guess_product(name):
  if re.match(r'^(aa).+$', name):
    return {'os': 'windows', 'product': 'aa'}

  if re.match(r'^(sp-|spc-).+$', name):
    return {'os': 'windows', 'product': 'sp'}

  if re.match(r'^(ve|slc-ve-).+$', name):
    return {'os': 'windows', 'product': 've'}

  if re.match(r'^(ar-|acr-).+$', name):
    return {'os': 'windows', 'product': 'ar'}

  if re.match(r'^(dr-).+$', name):
    return {'os': 'windows', 'product': 'dr'}

  if re.match(r'(w|slc-win-|slc-addc-|slc-vac-|efs-win|efs-addc-|efs-wsus).+$', name):
    return {'os': 'windows', 'product': 'unknown'}


  if re.match(r'^cfus-prd-365-\d+$', name):
    return {'os': 'windows', 'product': 'cf'}

  if re.match(r'^rb-prd-.+\d+$', name):
    return {'os': 'linux', 'product': 'rb', 'role': 'balancer'}

  if re.match(r'^rb-.+\d+$', name):
    return {'os': 'linux', 'product': 'rb', 'role': 'vault'}

  if re.match(r'^rb-.+-mgmt$', name):
    return {'os': 'linux', 'product': 'rb', 'role': 'gmp'}

  if re.match(r'^cf(us)?-prd-index\d+$', name):
    return {'os': 'linux', 'product': 'cf', 'role': 'esnode'}

  if re.match(r'^cf(us)?-prd-index(m|p)\d+$', name):
    return {'os': 'linux', 'product': 'cf', 'role': 'esmaster'}

  if re.match(r'^(cf(us)?-).+$', name):
    return {'os': 'linux', 'product': 'cf'}

  if re.match(r'^(axc(ca|ams|atl)?-|(ylw|atl|ams|slc)-axc-).+$', name):
    return {'os': 'linux', 'product': 'axc'}

  if re.match(r'^(anc-(prd|ylw|atl|ams|slc)).+$', name):
    return {'os': 'linux', 'product': 'anc'}

  if re.match(r'^(efs-(prd|dev)-|slc-prd).+$', name):
    return {'os': 'linux', 'product': 'unknown'}


  if re.match(r'^(iaso|sec).+$', name):
    return {'os': 'linux', 'product': 'sec'}


  return {'os': 'unknown', 'product': 'unknown'}


def main(shared_cache=None):
    global pool_cache
    pool_cache = shared_cache
    conn = None
    try:
        conn = libvirt.open("qemu:///system")
    except libvirt.libvirtError as e:
        print(repr(e), file=sys.stderr)
        exit(1)
    
    stats = conn.getCPUStats(libvirt.VIR_NODE_CPU_STATS_ALL_CPUS)
    
    domainIDs = conn.listDomainsID()
    if domainIDs == None:
        print('Failed to get a list of domain IDs', file=sys.stderr)
    
    _domain = {}
    metrics = []
    
    _host = conn.getHostname().split('.')[0]
    dc = _host[:3]

    cputotal = stats['user']/1e9 + stats['kernel']/1e9 + stats['iowait']/1e9 
    
    metrics.append(metric("libvirt.host.mem",        "gauge",   psutil.virtual_memory().total, {}))
    metrics.append(metric("libvirt.host.cpu.user",   "counter", stats['user']/1e9,   {}))
    metrics.append(metric("libvirt.host.cpu.kernel", "counter", stats['kernel']/1e9, {}))
    metrics.append(metric("libvirt.host.cpu.idle",   "counter", stats['idle']/1e9,   {}))
    metrics.append(metric("libvirt.host.cpu.iowait", "counter", stats['iowait']/1e9, {}))
    metrics.append(metric("libvirt.host.cpu.total",  "counter", cputotal, {}))
    metrics.append(metric("libvirt.host.cpu.cores",  "gauge", os.cpu_count(), {}))
    metrics.append(metric("libvirt.host.load1",      "gauge", os.getloadavg()[0] / os.cpu_count(), {}))
    metrics.append(metric("libvirt.domains.count",   "gauge", len(domainIDs), {}))
    
    if len(domainIDs) == 0:
        pass
    else:
        for domainID in domainIDs:
            tags = {}
            dom = conn.lookupByID(domainID)
            tree = ElementTree.fromstring(dom.XMLDesc())
            iface = tree.findall('devices/disk')
    
            (_domain["state"], _domain["maxmem"], _domain["mem"], _domain["vcpus"], _domain["cputime"]) = dom.info()
            tags["vmname"] = tree.find('name').text
            vmn = re.match(r'^one-(\d+)$', tags["vmname"])
            if vmn:
                tags["vmname"] = tree.find('title').text
                tags["one.id"] = int(vmn.group(1))
                tags["vmm"] = "one"
            else:
                tags["vmm"] = "spxops"
    
            product_tags = guess_product(tags["vmname"])
            tags["has_win"] = 0
            tags["has_rbvault"] = 0
            if product_tags['os'] == 'windows':
                tags["has_win"] = 1
            if product_tags["product"] == 'rb' and product_tags["role"] == 'vault':
                tags["has_rbvault"] = 1
            tags = {**tags, **product_tags}
            #print(tags)
            for _if in iface:
                for m in get_blk_metrics(_if, dom, tags, dc):
                    metrics.append(m)
    
            ifcs = tree.findall('devices/interface')
            for ifc in ifcs:
                for m in get_net_metrics(ifc, dom, tags):
                    metrics.append(m)
            metrics.append(get_cpu_metrics_new(tree, dom, conn, tags))
    
            metrics.append(metric("libvirt.mem.size", "gauge", _domain["mem"]*1024, tags))
            metrics.append(metric("libvirt.cpu.cores", "gauge", _domain["vcpus"], tags))
    
    r = requests.post('https://efs-prd-atlas1.slc.efscloud.net/api/v1/publish', json={"tags":{"host":_host, "datacenter": dc, "role": "vmc"}, "metrics": metrics})
    if not r.status_code == 200:
        print(r.text, file=sys.stderr)

if __name__ == "__main__":
    from multiprocessing import Process, Manager

    manager = Manager()
    shared_cache = manager.dict()

    while True:
        p = Process(target=main, args=(shared_cache,))
        p.start()
        p.join(60)

        if p.is_alive():
            p.terminate()
            p.join(5)
            if p.is_alive():
                p.kill()
                p.join()
            print("Collection timeout, process killed", file=sys.stderr)

        time.sleep(15)
