
from __future__ import absolute_import
from __future__ import division

from collections import defaultdict, namedtuple
import xml.etree.cElementTree as ET

from vdsm import taskset
from vdsm.common import cache
from vdsm.common import commands
from vdsm.common import libvirtconnection
from vdsm.common.cmdutils import CommandPath


NumaTopology = namedtuple('NumaTopology', 'topology, distances, cpu_topology')
CpuTopology = namedtuple('CpuTopology', 'sockets, cores, threads, online_cpus')


_SYSCTL = CommandPath("sysctl", "/sbin/sysctl", "/usr/sbin/sysctl")


AUTONUMA_STATUS_DISABLE = 0
AUTONUMA_STATUS_ENABLE = 1
AUTONUMA_STATUS_UNKNOWN = 2


def topology(capabilities=None):
    return _numa(capabilities).topology


def distances():
    return _numa().distances


def cpu_topology(capabilities=None):
    return _numa(capabilities).cpu_topology


@cache.memoized
def autonuma_status():
    out = commands.run([_SYSCTL.cmd, '-n', '-e', 'kernel.numa_balancing'])

    if not out:
        return AUTONUMA_STATUS_UNKNOWN
    elif out[0] == '0':
        return AUTONUMA_STATUS_DISABLE
    elif out[0] == '1':
        return AUTONUMA_STATUS_ENABLE
    else:
        return AUTONUMA_STATUS_UNKNOWN


def memory_by_cell(index):
    conn = libvirtconnection.get()
    meminfo = conn.getMemoryStats(index, 0)
    meminfo['total'] = str(meminfo['total'] // 1024)
    meminfo['free'] = str(meminfo['free'] // 1024)
    return meminfo


@cache.memoized
def _numa(capabilities=None):
    if capabilities is None:
        capabilities = _get_libvirt_caps()

    topology = defaultdict(dict)
    distances = defaultdict(dict)
    sockets = set()
    siblings = set()
    online_cpus = []

    caps = ET.fromstring(capabilities)
    cells = caps.findall('.host//cells/cell')

    for cell in cells:
        cell_id = cell.get('id')
        if len(cells) == 1:
            idx = -1
        else:
            idx = int(cell_id)
        meminfo = memory_by_cell(idx)
        topology[cell_id]['totalMemory'] = meminfo['total']
        topology[cell_id]['cpus'] = []
        distances[cell_id] = []

        for cpu in cell.findall('cpus/cpu'):
            topology[cell_id]['cpus'].append(int(cpu.get('id')))
            if cpu.get('siblings') and cpu.get('socket_id'):
                online_cpus.append(cpu.get('id'))
                sockets.add(cpu.get('socket_id'))
                siblings.add(cpu.get('siblings'))

        if cell.find('distances') is not None:
            for sibling in cell.find('distances').findall('sibling'):
                distances[cell_id].append(int(sibling.get('value')))

    cpu_topology = CpuTopology(len(sockets), len(siblings),
                               len(online_cpus), online_cpus)

    if not cells:
        hostcputop = caps.find('.host/cpu/topology')

        if hostcputop is not None:
            socketnum = int(hostcputop.get('sockets'))
            corenum = int(hostcputop.get('cores')) * socketnum
            threadnum = int(hostcputop.get('threads')) * corenum

            online_cpus = taskset.online_cpus()

            cpu_topology = CpuTopology(socketnum, corenum,
                                       threadnum, online_cpus)

    return NumaTopology(topology, distances, cpu_topology)


def _get_libvirt_caps():
    conn = libvirtconnection.get()
    return conn.getCapabilities()
