#!/usr/bin/env python

import libvirt
import threading
from xml.dom import minidom
import time
import sys
import getopt
import os

uri = "qemu:///system"
customXMLuri = "guest-cpu.python.libvirt.org"
connectRetryTimeout = 5

def usage():
    print("usage: "+os.path.basename(sys.argv[0])+" [-h] [uri]")
    print("   uri will default to qemu:///system")
    print("   --help, -h   Print(this help message")
    print("")
    print("This service waits for the guest agent lifecycle event and reissues " +
           "guest agent calls to modify the cpu count according to the metadata " +
           "set by guest-vcpu.py example")

class workerData:
    def __init__(self):
        self.doms = list()
        self.conn = None
        self.cond = threading.Condition()

    def notify(self):
        self.cond.acquire()
        self.cond.notify()
        self.cond.release()

    def waitNotify(self):
        self.cond.acquire()
        self.cond.wait()
        self.cond.release()

    def addDomainNotify(self, dom):
        self.doms.append(dom)
        self.notify()

    def closeConnectNotify(self):
        conn = self.conn
        self.conn = None
        conn.close()
        self.notify()

    def setConnect(self, conn):
        self.conn = conn

    def hasConn(self):
        return self.conn is not None

    def hasDom(self):
        return len(self.doms) > 0

    def getDom(self):
        return self.doms.pop()

    def setDoms(self, doms):
        self.doms = doms


def virEventLoopNativeRun():
    while True:
        libvirt.virEventRunDefaultImpl()

def handleAgentLifecycleEvent(conn, dom, state, reason, opaque):
    if state == libvirt.VIR_CONNECT_DOMAIN_EVENT_AGENT_LIFECYCLE_STATE_CONNECTED:
        opaque.addDomainNotify(dom)

def handleConnectClose(conn, reason, opaque):
    print('Disconnected from ' + uri)
    opaque.closeConnectNotify()

def handleLibvirtLibraryError(opaque, error):
    pass

def processAgentConnect(dom):
    try:
        cpus = dom.metadata(libvirt.VIR_DOMAIN_METADATA_ELEMENT, customXMLuri,
                libvirt.VIR_DOMAIN_AFFECT_LIVE)
        doc = minidom.parseString(cpus)
        ncpus = int(doc.getElementsByTagName('ncpus')[0].getAttribute('count'))
    except:
        return

    try:
        dom.setVcpusFlags(ncpus, libvirt.VIR_DOMAIN_AFFECT_LIVE | libvirt.VIR_DOMAIN_VCPU_GUEST)
        print("set vcpu count for domain " + dom.name() + " to " + str(ncpus))
    except:
        print("failed to set vcpu count for domain " + dom.name())

def work():
    data = workerData()

    print("Using uri: " + uri)

    while True:
        if not data.hasConn():
            try:
                conn = libvirt.open(uri)
            except:
                print('Failed to connect to ' + uri + ', retry in ' + str(connectRetryTimeout)) + ' seconds'
                time.sleep(connectRetryTimeout)
                continue

            print('Connected to ' + uri)

            data.setConnect(conn)
            conn.registerCloseCallback(handleConnectClose, data)
            conn.setKeepAlive(5, 3)
            conn.domainEventRegisterAny(None,
                                        libvirt.VIR_DOMAIN_EVENT_ID_AGENT_LIFECYCLE,
                                        handleAgentLifecycleEvent,
                                        data)

            data.setDoms(conn.listAllDomains(libvirt.VIR_CONNECT_LIST_DOMAINS_ACTIVE))

        while data.hasConn() and data.hasDom():
            processAgentConnect(data.getDom())

        data.waitNotify()

def main():
    libvirt.virEventRegisterDefaultImpl()
    libvirt.registerErrorHandler(handleLibvirtLibraryError, None)

    worker = threading.Thread(target=work)
    worker.setDaemon(True)
    worker.start()

    eventLoop = threading.Thread(target=virEventLoopNativeRun)
    eventLoop.setDaemon(True)
    eventLoop.start()

    while True:
       time.sleep(1)

if __name__ == "__main__":
    try:
        opts, args = getopt.getopt(sys.argv[1:], "h", ["help"])
    except getopt.GetoptError as err:
        print(str(err))
        usage()
        sys.exit(2)
    for o, a in opts:
        if o in ("-h", "--help"):
            usage()
            sys.exit()

    if len(args) > 1:
        usage()
        sys.exit(1)
    elif len(args) == 1:
        uri = args[0]

    main()