#!/usr/bin/env python3

import argparse
import os
import signal
import subprocess
import time

import psutil


def get_free_cores(n, verbose):
  while True:
    num_physical_core = psutil.cpu_count(logical=False)
    core_usage = psutil.cpu_percent(interval=1, percpu=True)
    num_window = num_physical_core // n
    for i in range(num_window):
        window_usage = core_usage[i * n : i * n + n]
        if all(_ < 0.3 for _ in window_usage):
            return (((i * n) % 128)// 64, i * n, i * n + n - 1)
    if verbose:
        print(f"No free {n} cores found. CPU usage: {core_usage}\n")
    time.sleep(60)


def main(commands, threads, verbose):
  # We assume that CPU with more than 16 cores has NUMA nodes
  require_numa = psutil.cpu_count(logical=False) > 16
  if threads > 1 and require_numa:
    numa_node, start_core, end_core = get_free_cores(threads, verbose)
    numa_cmd = ["numactl", "-m", str(numa_node), "-C", f"{start_core}-{end_core}"]
    commands = numa_cmd + commands
  if verbose:
    print(" ".join(commands))
  try:
    proc = subprocess.Popen(commands, preexec_fn=os.setsid)
    proc.wait()
  except KeyboardInterrupt:
    os.killpg(os.getpgid(proc.pid), signal.SIGINT)


if __name__ == "__main__":
  parser = argparse.ArgumentParser(description='Runner wrapper for XiangShan emu')
  parser.add_argument('--threads', "-T", nargs='?', type=int, default=8, help='number of emu threads')
  parser.add_argument('--verbose', "-v", action='store_true', default=False, help='verbose')
  parser.add_argument('commands', nargs=argparse.REMAINDER, help='commands')

  args = parser.parse_args()

  main(args.commands, args.threads, args.verbose)
