#!/usr/bin/env python3


import argparse
import logging
import math
import subprocess
import time
from collections import Counter
from typing import List


logging.basicConfig(format='%(asctime)-15s [%(levelname)s] %(message)s')
LOG = logging.getLogger(__name__)
LOG.setLevel('INFO')


def get_pid(slot: int) -> int:
    """Get PID of the process running in the slot. Note that this naively assumes
    the process is the worker. If it's a wrapper, the pid obtained here won't
    be the right one for monitoring usage."""
    # Assuming only boinc has such path pattern. Or we should prepend with full
    # boinc data directory path.
    sp = subprocess.run(
        ["pgrep", "-f", f"/slots/{slot}/"],
        stdout=subprocess.PIPE,
    )
    output = sp.stdout.decode("utf-8")
    try:
        pid = int(output)
        LOG.info(f"pid of slot {slot}: {pid}")
        return pid
    except:
        LOG.error(f"Failed to get pid for slot {slot}: {output}")
        raise


def collect_once(pid: int) -> int:
    sp = subprocess.run(
        ["ps", "--no-headers", "-o", "rss", "--pid", str(pid)],
        stdout=subprocess.PIPE,
    )
    try:
        rss = int(sp.stdout.decode("utf-8"))
        LOG.debug(f"RSS usage of {pid}: {rss}")
    except:
        # Not worth crashing if process happen to exit
        LOG.warning(f"Failed to get RSS for pid {pid}")
        rss = 0
    return rss


def print_histogram(numbers: List[int], num_buckets: int) -> None:
    numbers = [n for n in numbers if n != 0]
    min_val = min(numbers)
    max_val = max(numbers)
    step = (max_val - min_val + 1) / num_buckets
    buckets = Counter()
    for n in numbers:
        buckets[int((n - min_val) / step)] += 1

    count_per_star = max(1, int(max(buckets.values()) / 32))

    for i in range(num_buckets):
        lower = math.ceil(min_val + i * step)
        upper = math.floor(min_val + (i + 1) * step)
        stars = "*" * int(buckets[i] / count_per_star)
        print(f"{lower} - {upper}: {stars} ({buckets[i]})")


def parse_args():
    parser = argparse.ArgumentParser(description="Collect memory usage of a boinc task")
    parser.add_argument("--slot", type=int, required=True, help="Task slot to monitor")
    parser.add_argument(
        "--duration", type=int, default=300, help="Duration of capture in seconds",
    )
    parser.add_argument(
        "--buckets", type=int, default=16, help="Number of buckets for histogram",
    )
    parser.add_argument("--debug", action="store_true", help="Enable deubg logging")
    return parser.parse_args()


def main():
    args = parse_args()
    if args.debug:
        LOG.setLevel('DEBUG')

    pid = get_pid(args.slot)
    rsses = []
    for _ in range(args.duration):
        rsses.append(collect_once(pid))
        time.sleep(1)

    print_histogram(rsses, args.buckets)


if __name__ == "__main__":
    main()