#!/usr/bin/env python3
# coding: utf-8
#
# Copyright © 2023 - Rennes Physics Institute
#
# This file is part of lsusers.
#
# lsusers 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 3 of the License, or
# (at your option) any later version.
# lsusers 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 lsusers.  If not, see <http://www.gnu.org/licenses/>.


import os
from subprocess import Popen, PIPE, STDOUT

NCPUS = os.cpu_count()
TOTALMEM = os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES") / (1024**3)

def get_users_list():
    cmd = Popen(['loginctl', 'list-users', '--no-legend'], stdin=PIPE, stdout=PIPE, stderr=STDOUT)
    output = cmd.stdout.read().decode('utf8').split()[1::2]

    cmd = Popen(['users'], stdin=PIPE, stdout=PIPE, stderr=STDOUT)
    users_output = cmd.stdout.read().decode('utf8').split()
    for name in users_output:
        if name not in output:
            output.append(name)
    return output


def get_resources_percent(username=None):
    try:
        if username is None:
            user_opt = ['-e']
        else:
            user_opt = ['-u', username]
        cmd = Popen(['ps', '-o', 'pcpu', '--no-headers'] + user_opt, stdin=PIPE, stdout=PIPE, stderr=STDOUT)
        output = list(map(float, cmd.stdout.read().decode('utf8').split()))
        cpu_percent = sum(output) / NCPUS
        cmd = Popen(['ps', '-o', 'pmem', '--no-headers'] + user_opt, stdin=PIPE, stdout=PIPE, stderr=STDOUT)
        output = map(float, cmd.stdout.read().decode('utf8').split())
        mem_percent = sum(output)
    except:
        cpu_percent = mem_percent = 0
    return cpu_percent, mem_percent


def get_free_resources():
    cpu, mem = get_resources_percent()
    free_cpu = (100 - cpu) * NCPUS / 100
    free_mem = (100 - mem) * TOTALMEM / 100
    return free_cpu, free_mem


def get_resources():
    users = get_users_list()
    resources = list(map(get_resources_percent, users))
    cpu_percent, mem_percent = list(zip(*resources))
    cpu_usage = [int(_ / 100 * NCPUS) + 1 for _ in cpu_percent]
    mem_usage = [_ /100 * TOTALMEM for _ in mem_percent]
    free_cpu, free_mem = get_free_resources()
    return users, cpu_percent, cpu_usage, mem_percent, mem_usage, free_cpu, free_mem


def print_output(users, cpu_percent, cpu_usage, mem_percent, mem_usage, free_cpu, free_mem):
    header = "| {:<20s} | {:<20s} | {:>23s} |".format("Name ({:d} users)".format(len(users)), "CPU".center(20), "Memory")
    double_line = "+" + "="*(len(header) - 2) + "+"
    simple_line = "|" + "-"*(len(header) - 2) + "|"
    row_lines = []
    for i, user in enumerate(users):
        s = "| {:<20s} | {:5.1f}% [{:2d}/{:d} cores] | {:5.1f}% [{:7.3f}/{:d} GB] |".format(
                user, cpu_percent[i], cpu_usage[i], NCPUS, mem_percent[i], mem_usage[i], int(TOTALMEM))
        row_lines.append(s)
    free_line = "| {:<20s} | {:13.1f} cores  | {:19.0f} GB  |".format("FREE", free_cpu, free_mem)

    s  = double_line + '\n'
    s += header + '\n'
    s += simple_line + '\n'
    s += '\n'.join(row_lines) + '\n'
    s += simple_line + '\n'
    s += free_line + '\n'
    s += double_line

    print(s)


if __name__ == "__main__":
    resources = get_resources()
    print_output(*resources)