#!/usr/bin/python3
# SPDX-License-Identifier: GPL-3.0+
# Copyright 2022-2022 Lukas F. Hartmann <lukas@mntre.com>
# Copyright 2024-2025 Johannes Schauer Marin Rodrigues <josch@mister-muffin.de>
#
# Do not run shellcheck on this file
# shellcheck disable=SC1071

import time, psutil, sys, getopt


def draw_bar(val, maxval):
    if val < 0:
        return "    "
    bars = ["    ", "▁", "▁", "▂", "▂", "▃", "▄", "▅", "▆", "▆", "▇", "█"]
    return bars[int((len(bars) - 1) * min(max(0, val / maxval), 1))]


def draw_chart(lst, maxval):
    chart = ""
    l = len(lst)
    for x in range(l):
        chart += draw_bar(lst[l - x - 1], maxval)
    return chart


def history_insert(lst, item, depth):
    lst.insert(0, item)
    if len(lst) > depth:
        lst.pop()
    return lst


def usage():
    print(
        "Visualize CPU usage and disk read/writes as text using box-drawing characters.",
        file=sys.stderr,
    )
    print(file=sys.stderr)
    print(
        "Usage: compstat -d <number of bars> -i <interval in seconds, float>",
        file=sys.stderr,
    )
    print(file=sys.stderr)
    print("Options:", file=sys.stderr)
    print("  --help   Display this help text and exit.", file=sys.stderr)
    print("  -d NUM   Set number of bars to NUM.", file=sys.stderr)
    print("  -i SECS  Set the update interval to SEC seconds.", file=sys.stderr)
    print("  --percpu Display the CPU utilization per CPU core.", file=sys.stderr)


def main(argv):
    depth = 5
    interval = 0.3
    percpu = False

    try:
        opts, args = getopt.getopt(argv, "hd:i:", longopts=["percpu", "help"])
    except getopt.GetoptError as e:
        print(f"E: {e.msg}", file=sys.stderr)
        print(file=sys.stderr)
        usage()
        sys.exit(2)
    for opt, arg in opts:
        if opt == "-d":
            depth = int(arg)
        elif opt == "-i":
            interval = float(arg)
        elif opt == "--percpu":
            percpu = True
        elif opt in ["-h", "--help"]:
            usage()
            exit()

    if percpu:
        cpu_history = {i: [] for i in range(psutil.cpu_count())}
    else:
        cpu_history = []
    read_history = []
    write_history = []
    ldisk_activity = psutil.disk_io_counters()

    while True:
        cpu_usage = psutil.cpu_percent(percpu=percpu)
        disk_activity = psutil.disk_io_counters()

        if percpu:
            for i, cpu in enumerate(cpu_usage):
                history_insert(cpu_history[i], cpu, depth)
        else:
            history_insert(cpu_history, cpu_usage, depth)
        history_insert(
            read_history,
            max(0, disk_activity.read_time - ldisk_activity.read_time)
            / (interval * 1000),
            depth,
        )
        history_insert(
            write_history,
            max(0, disk_activity.write_time - ldisk_activity.write_time)
            / (interval * 1000),
            depth,
        )

        print(
            "CPU "
            + (
                (
                    "".join(
                        draw_chart(cpu_history[i], 98.0)
                        for i in range(psutil.cpu_count())
                    )
                )
                if percpu
                else draw_chart(cpu_history, 98.0)
            )
            + " R/W "
            + draw_chart(read_history, 0.98)
            + " "
            + draw_chart(write_history, 0.98),
            flush=True,
        )

        ldisk_activity = disk_activity
        time.sleep(interval)


main(sys.argv[1:])
