#!/usr/bin/python2
#
# 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 xml.etree.ElementTree as ET


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_config import VEConfig
from vcmmd.rpc.dbus.client import RPCProxy
from vcmmd.error import VCMMDError, VCMMD_ERROR_VE_NAME_ALREADY_IN_USE
from vcmmd.cgroup import MemoryCgroup
from vcmmd.util.libvirt import virDomainProxy, virConnectionProxy
from libvirt import VIR_CONNECT_LIST_DOMAINS_RUNNING, \
                    VIR_CONNECT_LIST_DOMAINS_PAUSED

from vcmmd.ve_type import VE_TYPE_VM, VE_TYPE_SERVICE

from vcmmd.cgroup import MemoryCgroup


VSTORAGE_SLICE = 'vstorage.slice'
MEMORY_CGROUP_PATH = '/sys/fs/cgroup/memory/'
VSTORAGE_CGROUP_PATH = os.path.join(MEMORY_CGROUP_PATH, VSTORAGE_SLICE)

proxy = RPCProxy()

memcgrps = ['system.slice', 'user.slice']
for service_slice in ('vstorage-compute.slice',
                      'vstorage-compute.slice/vstorage-compute-storage.slice',
                      'vstorage-services.slice',
                      'vstorage-target.slice',
                      'vstorage-ui.slice',
                     ):
    full_path = os.path.join(VSTORAGE_CGROUP_PATH, service_slice)
    if not os.path.isdir(full_path):
        continue
    name = '/'.join([VSTORAGE_SLICE, service_slice])
    memcgrps.append(name)

for service in ['system.slice/postgresql.service', 'system.slice/nginx.service']:
    full_path = os.path.join(MEMORY_CGROUP_PATH, service)
    if not os.path.isdir(full_path):
        continue
    memcgrps.append(service)

for name in memcgrps:
    memcgrp = MemoryCgroup(name)
    kv = {}
    kv['guarantee'] = memcgrp.read_mem_low()
    kv['limit'] = memcgrp.read_mem_max()
    if not memcgrp.read_swappiness():
        kv['swap'] = 0
    kv['cache'] = memcgrp.read_cache_limit_in_bytes()

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

    proxy.activate_ve(name, 0)

for vm in virConnectionProxy().list_all_vm(VIR_CONNECT_LIST_DOMAINS_RUNNING | \
                                           VIR_CONNECT_LIST_DOMAINS_PAUSED):
    kv = {}

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

    kv['nodelist'] = str() # TODO
    kv['cpulist'] = str() # TODO

    ram_size = vm.maxMemory() << 10

    kv['limit'] = ram_size

    kv['guarantee_type'] = 0

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

    uuid = vm.UUIDString()

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

    if not vm.isActive():
        continue

    proxy.activate_ve(uuid, 0)
