
from __future__ import absolute_import

import os
import collections
import threading

from vdsm import osinfo
from vdsm.common import cache
from vdsm.common import cpuarch
from vdsm.common import supervdsm
from vdsm.config import config


_PATH = '/sys/kernel/mm/hugepages'
_VM = '/proc/sys/vm/'

lock = threading.Lock()

DEFAULT_HUGEPAGESIZE = {
    cpuarch.X86_64: 2048,
    cpuarch.PPC64LE: 16384,
}


class NonContiguousMemory(Exception):
    """Raised when the memory is too fragmented to allocate hugepages"""


@cache.memoized
def supported(path=_PATH):
    return state(path).keys()


def alloc(count, size=None,
          path='/sys/kernel/mm/hugepages/hugepages-{}kB/nr_hugepages'
          ):
    return _alloc(count, size, path)


def dealloc(count, size=None,
            path='/sys/kernel/mm/hugepages/hugepages-{}kB/nr_hugepages'
            ):
    return -(_alloc(-count, size, path))


def _alloc(count, size, path):
    if size is None:
        size = DEFAULT_HUGEPAGESIZE[cpuarch.real()]

    path = path.format(size)

    ret = supervdsm.getProxy().hugepages_alloc(count, path)
    if ret != count:
        supervdsm.getProxy().hugepages_alloc(-ret, path)
        raise NonContiguousMemory

    return ret


def state(path=_PATH):
    sizes = collections.defaultdict(dict)
    for size in os.listdir(path):
        for key in (
                'free_hugepages', 'nr_hugepages',
                'nr_hugepages_mempolicy', 'nr_overcommit_hugepages',
                'resv_hugepages', 'surplus_hugepages'):
            size_in_kb = _size_from_dir(size)
            with open(os.path.join(path, size, key)) as f:
                sizes[size_in_kb][key] = int(f.read())

            sizes[size_in_kb]['vm.free_hugepages'] = max(
                int(sizes[size_in_kb]['free_hugepages']) -
                _reserved_hugepages(key), 0
            )

    return sizes


def calculate_required_allocation(cif, vm_hugepages, vm_hugepagesz):
    if vm_hugepagesz == 0:
        return 0

    if not config.getboolean('performance', 'use_preallocated_hugepages'):
        return vm_hugepages

    all_vm_hugepages = _all_vm_hugepages(cif, vm_hugepages, vm_hugepagesz)
    system_hugepages = state()
    free_hugepages = int(system_hugepages[vm_hugepagesz]['free_hugepages'])
    nr_hugepages = int(system_hugepages[vm_hugepagesz]['nr_hugepages'])

    really_free_hugepages = min(
        free_hugepages,
        nr_hugepages - all_vm_hugepages - _reserved_hugepages(vm_hugepagesz)
    )

    really_free_hugepages = max(really_free_hugepages, 0)

    to_allocate = max(vm_hugepages - really_free_hugepages, 0)

    return to_allocate


def calculate_required_deallocation(vm_hugepages, vm_hugepagesz):
    if vm_hugepagesz == 0:
        return 0

    if not config.getboolean('performance', 'use_preallocated_hugepages'):
        return vm_hugepages

    nr_hugepages = int(state()[vm_hugepagesz]['nr_hugepages'])

    to_deallocate = min(
        vm_hugepages,
        nr_hugepages - max(_reserved_hugepages(vm_hugepagesz),
                           _preallocated_hugepages(vm_hugepagesz))
    )

    return to_deallocate


def _all_vm_hugepages(cif, vm_hugepages, vm_hugepagesz):
    return sum(
        [vm.nr_hugepages for vm in cif.getVMs().values() if
         vm.hugepagesz == vm_hugepagesz]
    ) - vm_hugepages


def _preallocated_hugepages(vm_hugepagesz):
    kernel_args = osinfo.kernel_args_dict()
    if 'hugepagesz' not in kernel_args:
        hugepagesz = DEFAULT_HUGEPAGESIZE[cpuarch.real()]
    else:
        hugepagesz = _cmdline_hugepagesz_to_kb(
            kernel_args['hugepagesz']
        )

    preallocated_hugepages = 0
    if ('hugepages' in kernel_args and
            hugepagesz == vm_hugepagesz):
        preallocated_hugepages = int(kernel_args['hugepages'])

    return preallocated_hugepages


def _reserved_hugepages(hugepagesz):
    reserved_hugepages = 0
    if config.getboolean('performance', 'use_preallocated_hugepages'):
        reserved_hugepages = (
            config.getint('performance', 'reserved_hugepage_count') if
            config.get('performance', 'reserved_hugepage_size') ==
            str(hugepagesz) else 0
        )

    return reserved_hugepages


def _cmdline_hugepagesz_to_kb(cmdline):
    return {
        '1GB': 1048576,
        '1G': 1048576,
        '2M': 2048,
        '2MB': 2048,
    }[cmdline]


def _size_from_dir(path):
    return int(path[10:-2])
