
from __future__ import absolute_import

from vdsm.common import cmdutils
from vdsm.common import commands
from . import constants


AUTOMATIC = "auto"

_SYS_ONLINE_CPUS = "/sys/devices/system/cpu/online"


def get(pid):
    command = [constants.EXT_TASKSET, '--pid', str(pid)]

    rc, out, err = commands.execCmd(command, resetCpuAffinity=False)

    if rc != 0:
        raise cmdutils.Error(command, rc, out, err)

    return _cpu_set_from_output(out[-1])


def set(pid, cpu_set, all_tasks=False):
    command = [constants.EXT_TASKSET]
    if all_tasks:
        command.append("--all-tasks")

    command.extend((
                   '--pid',
                   '--cpu-list', ','.join(str(i) for i in cpu_set),
                   str(pid)
                   ))

    rc, out, err = commands.execCmd(command, resetCpuAffinity=False)

    if rc != 0:
        raise cmdutils.Error(command, rc, out, err)


def online_cpus():
    with open(_SYS_ONLINE_CPUS, 'r') as src:
        return cpulist_parse(src.readline())


def pick_cpu(cpu_set):
    cpu_list = sorted(cpu_set)
    return cpu_list[:2][-1]


def _cpu_set_from_output(line):
    hexmask = line.decode().rsplit(":", 1)[1].strip()
    mask = int(hexmask, 16)
    return frozenset(i for i in range(mask.bit_length()) if mask & (1 << i))


def cpulist_parse(cpu_range):
    cpus = []
    for item in cpu_range.split(','):
        if '-' in item:
            begin, end = item.split('-', 1)
            cpus.extend(range(int(begin), int(end) + 1))
        else:
            cpus.append(int(item))
    return frozenset(cpus)
