#!/usr/bin/python2
#

import os
import sys
import subprocess
import time
import re
import xml.etree.ElementTree as ET
import psutil
import json


def _excepthook(exc_type, exc_value, exc_traceback):
    sys.stderr.write('Line %d: %s.%s: %s' % (exc_traceback.tb_lineno,
                                             exc_type.__module__,
                                             exc_type.__name__, exc_value))
    sys.exit(1)

# Setup excepthook to avoid core dump generation on failure.
sys.excepthook = _excepthook


import prlsdkapi
from prlsdkapi import consts

from vcmmd.ve_type import (VE_TYPE_SERVICE,
                           VE_TYPE_CT,
                           VE_TYPE_VM,
                           VE_TYPE_VM_LINUX,
                           VE_TYPE_VM_WINDOWS)
from vcmmd.ve_config import VEConfig, VCMMD_MEMGUARANTEE_PERCENTS
from vcmmd.rpc.dbus.client import RPCProxy
from vcmmd.error import VCMMDError, VCMMD_ERROR_VE_NAME_ALREADY_IN_USE
from vcmmd.ve.ct import _lookup_cgroup
from vcmmd.cgroup import MemoryCgroup
from vcmmd.util.libvirt import virDomainProxy

PAGE_SIZE = os.sysconf('SC_PAGESIZE')
INT64_MAX = int(2 ** 63 - 1)

def clamp(v, l, h):
    if h == -1:
        h = INT64_MAX
    return max(l, min(v, h))

VE_TYPE_MAP = {
    consts.PVT_CT: VE_TYPE_CT,
    consts.PVT_VM: VE_TYPE_VM,
}

VE_OSTYPE_MAP = {
    consts.PVS_GUEST_TYPE_WINDOWS: VE_TYPE_VM_WINDOWS,
    consts.PVS_GUEST_TYPE_LINUX: VE_TYPE_VM_LINUX,
}

DEFAULT_GUARANTEE = {
    VE_TYPE_CT: 0,
    VE_TYPE_VM: 40,
    VE_TYPE_VM_LINUX: 40,
    VE_TYPE_VM_WINDOWS: 40,
}

#Wait for virtuozzo.target, which indicates that VE list can be retrived
# via prlsdkapi.
cmd = ["systemctl", "is-active", "virtuozzo.target"]
for wait in xrange(120):
    try:
        if subprocess.check_output(cmd) == "active\n":
            break
    except subprocess.CalledProcessError:
        pass
    time.sleep(1)

helper = prlsdkapi.ApiHelper()
helper.init(consts.PRL_VERSION_7X)

srv = prlsdkapi.Server()
srv.login_local().wait()

proxy = RPCProxy()

# VStorage must be activated first
vs_service_name = 'vstorage.slice/vstorage-services.slice'
vs_cgroup_path = '/sys/fs/cgroup/memory/%s'
vs_config = '/etc/vz/vstorage-limits.conf'

vs_exc_info = None
if os.path.isdir(vs_cgroup_path % vs_service_name) and os.path.exists(vs_config):
    try:
        with open(vs_config) as f:
            j = json.loads(f.read())
        total_mem = psutil.virtual_memory().total
        kv = {}
        kv['limit'] = clamp(int(j['VStorage']['Limit']['Share'] * total_mem),
                            j['VStorage']['Limit']['Min'],
                            j['VStorage']['Limit']['Max'])
        kv['guarantee'] = clamp(int(j['VStorage']['Guarantee']['Share'] * total_mem),
                            j['VStorage']['Guarantee']['Min'],
                            j['VStorage']['Guarantee']['Max'])
        kv['swap'] = clamp(int(j['VStorage']['Swap']['Share'] * total_mem),
                           j['VStorage']['Swap']['Min'],
                           j['VStorage']['Swap']['Max'])
    except (KeyError, TypeError, ValueError):
        vs_exc_info = sys.exc_info()
    else:
        try:
            proxy.register_ve(vs_service_name, VE_TYPE_SERVICE, VEConfig(**kv), 0)
            proxy.activate_ve(vs_service_name, 0)
        except VCMMDError as e:
            if e.errno == VCMMD_ERROR_VE_NAME_ALREADY_IN_USE:
                pass
            raise

for vm in srv.get_vm_list_ex(consts.PVTF_CT | consts.PVTF_VM).wait():
    try:
        ve_type = VE_TYPE_MAP[vm.get_vm_type()]
    except KeyError:
        continue

    ostype = vm.get_os_type()
    if ve_type == VE_TYPE_VM:
        ve_type = VE_OSTYPE_MAP.get(ostype, VE_TYPE_VM)

    state = vm.get_state().wait().get_param().get_state()
    if state not in (consts.VMS_RUNNING, consts.VMS_PAUSED):
        continue

    kv = {}

    if ve_type == VE_TYPE_CT:
        uuid = vm.get_ct_id()
        memcg = _lookup_cgroup(MemoryCgroup, uuid)
        ram_size = memcg.read_mem_max()
        vram = None
        swap = memcg.read_swap_max()
    else:
        uuid = vm.get_uuid()[1:-1]
        dom = virDomainProxy(uuid)
        if vm.is_ram_hotplug_enabled():
            ram_size = vm.get_ram_size() << 20
        else:
            ram_size = dom.maxMemory() << 10
        video = ET.fromstring(dom.XMLDesc()).findall("./devices/video/model")
        vram = sum(map(lambda v: int(v.attrib.get('vram',0)), video)) << 10
        swap = None

    kv['nodelist'] = str(vm.get_node_mask())
    kv['cpulist'] = str(vm.get_cpu_mask())

    kv['limit'] = ram_size

    mem_guarantee = vm.get_mem_guarantee_size()
    guar_pct = (mem_guarantee[1] if mem_guarantee[0] == VCMMD_MEMGUARANTEE_PERCENTS else
                DEFAULT_GUARANTEE[ve_type])
    kv['guarantee'] = ram_size * guar_pct / 100
    kv['guarantee_type'] = mem_guarantee[0]

    if swap is not None:
        kv['swap'] = swap

    if vram is not None:
        kv['vram'] = vram

    try:
        proxy.register_ve(uuid, ve_type, VEConfig(**kv), 0)
    except VCMMDError as e:
        if e.errno == VCMMD_ERROR_VE_NAME_ALREADY_IN_USE:
            continue
        raise

    if state == consts.VMS_RUNNING:
        proxy.activate_ve(uuid, 0)

if vs_exc_info is not None:
    exc_class, exc, tb = vs_exc_info
    new_exc_type = type('VStorageConfigException', (Exception,), dict())
    new_exc = new_exc_type("Failed to parse VStorage config: %r: %s" % (exc_class.__name__, exc))
    raise new_exc_type, new_exc, tb
