from server_modules.util import *
from psutil import cpu_count, cpu_times, cpu_times_percent, virtual_memory, swap_memory, disk_partitions, disk_usage


def get_cpu_info(logger):

    # cpu counts (logical, hardware)

    num_cpu = cpu_count(logical=True)
    num_cpu_hw = cpu_count(logical=False)

    cpu_results = [make_map('cpu_count', [
        make_data('integer', 'logical', num_cpu),
        make_data('integer', 'physical', num_cpu_hw)
    ])]
    logger.debug("logical CPUs: %02d, physical CPUs: %02d" % (num_cpu, num_cpu_hw))

    # cpu times and usage

    num_cpu += 1
    cpu_t = cpu_times(percpu=True)
    cpu_t.append(cpu_times(percpu=False))
    logger.debug("CPU times: %s" % str(cpu_t))
    cpu_u = cpu_times_percent(percpu=True)
    cpu_u.append(cpu_times_percent(percpu=False))
    logger.debug("CPU percentage: %s" % str(cpu_u))

    if len(cpu_u) != num_cpu:
        logger.warn("len(cpu_u) [%02d] != num_cpu + 1 [%02d]" % (len(cpu_u), num_cpu))
        return cpu_results

    if len(cpu_t) != num_cpu:
        logger.warn("len(cpu_t) [%02d] != num_cpu + 1 [%02d]" % (len(cpu_t), num_cpu))
        return cpu_results

    cpu_results.append(make_table('cpu', 'number', [
        'CPU ' + str(i) if i < num_cpu - 1 else 'total' for i in range(num_cpu)
    ], [
        make_data('integer', 'time.user', [
                  int(round(cpu_t[i].user)) for i in range(num_cpu)], 's'),
        make_data('integer', 'percentage.user', [
                  cpu_u[i].user for i in range(num_cpu)], '%'),
        make_data('integer', 'time.system', [
                  int(round(cpu_t[i].system)) for i in range(num_cpu)], 's'),
        make_data('integer', 'percentage.system', [
                  cpu_u[i].system for i in range(num_cpu)], '%'),
        make_data('integer', 'time.idle', [
                  int(round(cpu_t[i].idle)) for i in range(num_cpu)], 's'),
        make_data('integer', 'percentage.idle', [
                  cpu_u[i].idle for i in range(num_cpu)], '%')
    ]))

    return cpu_results


def get_ram_info(logger):

    ram_results = []

    # virtual memory

    ram_v = None
    try:
        ram_v = virtual_memory()
        logger.debug("virtual memory: %s" % str(ram_v))
    except Exception as e:
        logger.warn(e.message)

    if ram_v is not None:
        ram_data = []

        try:
            ram_data.append(make_data('integer', 'total', ram_v.total, 'bytes'))
        except AttributeError as e:
            logger.warn(e.message)
        try:
            ram_data.append(make_data('integer', 'used', ram_v.used, 'bytes'))
        except AttributeError as e:
            logger.warn(e.message)
        try:
            ram_data.append(make_data('integer', 'available', ram_v.available, 'bytes'))
        except AttributeError as e:
            logger.warn(e.message)
        try:
            ram_data.append(make_data('integer', 'free', ram_v.free, 'bytes'))
        except AttributeError as e:
            logger.warn(e.message)
        try:
            ram_data.append(make_data('integer', 'active', ram_v.active, 'bytes'))
        except AttributeError as e:
            logger.warn(e.message)
        try:
            ram_data.append(make_data('integer', 'inactive', ram_v.inactive, 'bytes'))
        except AttributeError as e:
            logger.warn(e.message)
        try:
            ram_data.append(make_data('integer', 'cached', ram_v.cached, 'bytes'))
        except AttributeError as e:
            logger.warn(e.message)
        try:
            ram_data.append(make_data('integer', 'shared', ram_v.shared, 'bytes'))
        except AttributeError as e:
            logger.warn(e.message)
        try:
            ram_data.append(make_data('integer', 'percent', ram_v.percent, '%'))
        except AttributeError as e:
            logger.warn(e.message)

        ram_results.append(make_map('ram', ram_data))

    # swap memory

    ram_s = None
    try:
        ram_s = swap_memory()
        logger.debug("swap memory: %s" % str(ram_s))
    except Exception as e:
        logger.warn(e.message)

    if ram_s is not None:
        ram_data = []

        try:
            ram_data.append(make_data('integer', 'total', ram_s.total, 'bytes'))
        except AttributeError as e:
            logger.warn(e.message)
        try:
            ram_data.append(make_data('integer', 'used', ram_s.used, 'bytes'))
        except AttributeError as e:
            logger.warn(e.message)
        try:
            ram_data.append(make_data('integer', 'free', ram_s.free, 'bytes'))
        except AttributeError as e:
            logger.warn(e.message)
        try:
            ram_data.append(make_data('integer', 'percent', ram_s.percent, '%'))
        except AttributeError as e:
            logger.warn(e.message)

        ram_results.append(make_map('swap', ram_data))

    return ram_results


def get_disk_info(logger):

    # partitions

    partitions = []

    for p in disk_partitions():
        mountpath = p.mountpoint

        usage = None
        try:
            usage = disk_usage(mountpath)
            logger.debug("[%s]: %s" % (mountpath, str(usage)))
        except Exception as e:
            logger.warn("[%s]: %s" % (mountpath, str(e)))

        if usage is not None:
            partitions.append({
                'mountpoint': mountpath,
                'device': p.device
            })

            try:
                partitions[-1]['total'] = {'v': usage.total, 'u': 'bytes'}
            except AttributeError as e:
                logger.warn(e.message)
            try:
                partitions[-1]['used'] = {'v': usage.used, 'u': 'bytes'}
            except AttributeError as e:
                logger.warn(e.message)
            try:
                partitions[-1]['free'] = {'v': usage.free, 'u': 'bytes'}
            except AttributeError as e:
                logger.warn(e.message)
            try:
                partitions[-1]['percent'] = {'v': usage.percent, 'u': '%'}
            except AttributeError as e:
                logger.warn(e.message)

    partitions = sorted(partitions, key=lambda x: x['device'], reverse=False)

    partition_columns = []

    for k in ['device', 'mountpoint']:
        partition_columns.append(
            make_data('string', k, [p[k] for p in partitions]))

    for k in ['total', 'used', 'free', 'percent']:
        partition_col = []
        with_units = False

        for p in partitions:
            if k in p:
                with_units = True
                partition_col.append(p[k]['v'])
            else:
                partition_col.append('N/A')

        if with_units:
            partition_columns.append(
                make_data('string', k, partition_col, p[k]['u']))
        else:
            partition_columns.append(make_data('string', k, partition_col))

    return make_table('partitions', 'number', [i for i in range(len(partitions))], partition_columns)
