#!/usr/bin/python2
#
# Copyright (c) 2016-2017, Parallels International GmbH
# Copyright (c) 2017-2019, Virtuozzo International GmbH, All rights reserved
#
# This file is part of OpenVZ. OpenVZ is free software; you can redistribute
# it and/or modify it under the terms of the GNU General Public License as
# published by the Free Software Foundation; either version 2 of the License,
# or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
# 02110-1301, USA.
#
# Our contact details: Virtuozzo International GmbH, Vordergasse 59, 8200
# Schaffhausen, Switzerland.

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


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
from libvirt import libvirtError

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))

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

is_avail_prlsdkapi = True
try:
    import prlsdkapi
    from prlsdkapi import consts
    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,
    }
except ImportError:
    print "Failed to import prlsdkapi"
    is_avail_prlsdkapi = False

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

proxy = RPCProxy()

# Services must be activated first
cgroup_path = '/sys/fs/cgroup/memory/%s'
config = '/etc/vz/vstorage-limits.conf'
known_params = {'Limit', 'Guarantee', 'Swap', 'Path'}

svs_exc_info = None
try:
    with open(config) as f:
        j = json.loads(f.read())
    to_register = []
    for name, service in j.iteritems():
        #VStorage config may exist without Path for backward compatibility
        if 'Path' not in service and name == 'VStorage':
            service['Path'] = 'vstorage.slice/vstorage-services.slice'
            print "Assuming that VStorage is located at %s" % service['Path']
        #check that cgroup exists
        service_name = service['Path']
        service_path = cgroup_path % service_name
        if not os.path.isdir(service_path):
            print "Memory cgroup %s not found" % service_path
            continue
        # TODO removeme PSBM-64068
        os.system("echo 1 > %s/memory.disable_cleancache" % service_path)
        # PSBM-89802
        os.system("echo 0 > %s/memory.swappiness" % service_path)
        #read config for limit, guarantee and swap
        total_mem = psutil.virtual_memory().total
        kv = {}
        unknown_params = set(service.iterkeys()) - known_params
        if unknown_params:
            raise Exception("Unknown fields in %s: %s" % (service_name, unknown_params))
        for Param in known_params - {'Path'}:
            param = Param.lower()
            try:
                unknown_params = (set(service[Param].iterkeys()) -
                                  {'Share','Min','Max'})
                if unknown_params:
                    raise Exception("Unknown fields in %s: %s" % (service_name, unknown_params))
                kv[param] = clamp(int(service[Param]['Share'] * total_mem),
                                  int(service[Param].get('Min', 0)),
                                  int(service[Param].get('Max', -1)))
            except (KeyError, TypeError, ValueError):
                raise Exception("Error parsing %s.%s" %(service_name, Param))
        #register and activate service later, if no exception occured
        to_register.append((service_name,VEConfig(**kv)))
    for service in to_register:
        proxy.register_ve(service[0], VE_TYPE_SERVICE, service[1], 0)
        proxy.activate_ve(service[0], 0)
except Exception as e:
    svs_exc_info = sys.exc_info()

vm_list = ()

if is_avail_prlsdkapi:
    helper = prlsdkapi.ApiHelper()
    helper.init(consts.PRL_VERSION_7X)

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

    vm_list = srv.get_vm_list_ex(consts.PVTF_CT | consts.PVTF_VM).wait()

for vm in vm_list:
    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]

        try:
            dom = virDomainProxy(uuid)
        except libvirtError:
            dom = None

        if vm.is_ram_hotplug_enabled() or not dom:
            ram_size = vm.get_ram_size() << 20
        else:
            ram_size = dom.maxMemory() << 10

        if dom:
            video = ET.fromstring(dom.XMLDesc()).findall("./devices/video/model")
            vram = sum(map(lambda v: int(v.attrib.get('vram',0)), video)) << 10
        else:
            vram = vm.get_video_ram_size() << 20

        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 svs_exc_info is not None:
    exc_class, exc, tb = svs_exc_info
    new_exc_type = type('ServicesConfigException', (Exception,), dict())
    new_exc = new_exc_type("Failed to parse services config: %r: %s" % (exc_class.__name__, exc))
    raise new_exc_type, new_exc, tb
else:
    print "VZ script finished successfully"
