A tool to monitor nimsuggest processes and their memory consumption

nsm.nim

# nsm.nim
# A tool to monitor nimsuggest processes and their memory consumption
# Created by FOXOMAN @ 2025

import os, osproc, strutils, strformat, parseopt, algorithm, math
import terminal

type
  ProcessInfo = object
    pid: int
    memory: int # in KB
    cpu: float
    runtime: string

  AppConfig = object
    killExcessProcesses: bool

# Track nimsuggest process IDs for cleanup
var
  nimsuggestPids: seq[int] = @[]
  config: AppConfig

proc alignCenter(text: string, width: int): string =
  let padding = max(0, width - text.len)
  let leftPad = padding div 2
  let rightPad = padding - leftPad
  return " ".repeat(leftPad) & text & " ".repeat(rightPad)

# Parse command line arguments
proc parseCommandLine(): AppConfig =
  result.killExcessProcesses = false

  # Parse command line
  for kind, key, val in getopt():
    case kind
    of cmdLongOption, cmdShortOption:
      case key
      #of "k":
      #result.killExcessProcesses = true
      #of "kill":
      #result.killExcessProcesses = true
      of "help", "h":
        echo "NimSuggest Monitor - Created by FOXOMAN @ 2025"
        echo "Usage: nsm [options]"
        echo "Options:"
        #echo "  -k, --kill    Kill excess nimsuggest processes, keeping only the last two"
        echo "  -h, --help    Show this help message"
        quit(0)
    else:
      discard

# Handle application shutdown
proc shutdownHandler() {.noconv.} =
  # Clear screen before exit
  eraseScreen()
  setCursorPos(0, 0)

  styledEcho(
    fgCyan,
    "╔════════════════════════════════════════════════════════╗",
  )
  styledEcho(fgCyan, "║              Shutting down NimSuggest Monitor          ║")
  styledEcho(
    fgCyan,
    "╚════════════════════════════════════════════════════════╝",
  )

  echo "Terminating all nimsuggest processes..."

  # Terminate all tracked nimsuggest processes
  var terminatedCount = 0
  for pid in nimsuggestPids:
    try:
      when defined(windows):
        # Windows termination
        discard execCmd(fmt"taskkill /F /PID {pid}")
      else:
        # Unix termination
        discard execCmd(fmt"kill -9 {pid}")
      terminatedCount += 1
    except:
      echo fmt"Failed to terminate process {pid}"

  echo fmt"Terminated {terminatedCount} nimsuggest processes"
  styledEcho(fgYellow, "Thank you for using NimSuggest Monitor!")
  styledEcho(fgYellow, "Created by FOXOMAN @ 2025")
  quit(0)

proc getProcessDetails(pid: int): ProcessInfo =
  when defined(windows):
    # For Windows, use WMIC
    let wmicCmd =
      fmt"wmic process where ProcessId={pid} get WorkingSetSize,UserModeTime /format:csv"
    let output = execProcess(wmicCmd).strip()
    var lines = output.splitLines()
    if lines.len >= 2:
      let parts = lines[1].split(",")
      if parts.len >= 3:
        try:
          let memBytes = parseInt(parts[2])
          let userTime = parseInt(parts[1]) / 10000000 # Convert 100ns units to seconds
          result.pid = pid
          result.memory = memBytes div 1024 # Convert to KB
          result.cpu = 0.0 # Not easily available through WMIC
          result.runtime = fmt"{userTime:.1f}s"
        except:
          discard
  elif defined(linux) or defined(macosx):
    # For Linux/macOS, use ps command
    let psCmd = fmt"ps -p {pid} -o pid,rss,%cpu,etimes"
    let output = execProcess(psCmd).strip()
    let lines = output.splitLines()
    if lines.len >= 2:
      let parts = lines[1].splitWhitespace()
      if parts.len >= 4:
        try:
          result.pid = pid
          result.memory = parseInt(parts[1]) # RSS in KB
          result.cpu = parseFloat(parts[2])
          let uptimeSecs = parseInt(parts[3])

          if uptimeSecs < 60:
            result.runtime = fmt"{uptimeSecs}s"
          elif uptimeSecs < 3600:
            result.runtime = fmt"{uptimeSecs div 60}m {uptimeSecs mod 60}s"
          else:
            let hours = uptimeSecs div 3600
            let mins = (uptimeSecs mod 3600) div 60
            result.runtime = fmt"{hours}h {mins}m"
        except:
          discard
  else:
    # Generic fallback with less information
    result.pid = pid
    result.memory = 0
    result.cpu = 0.0
    result.runtime = "unknown"

proc findNimsuggestProcesses(): seq[int] =
  when defined(windows):
    # For Windows, use WMIC
    let output = execProcess(
        "wmic process where name='nimsuggest.exe' get ProcessId /format:csv"
      )
      .strip()
    let lines = output.splitLines()
    for i in 1 ..< lines.len:
      let line = lines[i].strip()
      if line.len > 0:
        let parts = line.split(",")
        if parts.len >= 2:
          try:
            result.add(parseInt(parts[1]))
          except:
            discard
  elif defined(linux) or defined(macosx):
    # For Linux/macOS, use pgrep
    let output = execProcess("pgrep nimsuggest").strip()
    if output.len > 0:
      for pidStr in output.splitLines():
        try:
          result.add(parseInt(pidStr))
        except:
          discard
  else:
    # Generic fallback with basic checking
    echo "Unsupported platform, process detection might be limited"

    # Try to get process list in a generic way
    let output = execProcess("ps -e").strip()
    for line in output.splitLines():
      if "nimsuggest" in line:
        let parts = line.splitWhitespace()
        if parts.len > 0:
          try:
            result.add(parseInt(parts[0]))
          except:
            discard

  # Sort PIDs by value (likely correlates with start time)
  result.sort()

## FIXME: Need to find the best way
proc killExcessProcesses(pids: seq[int]) =
  if pids.len <= 2:
    return # Keep all if we have 2 or fewer

  let processesToKill = pids[0 ..^ 3] # All except last 2
  let processesToKeep = pids[^2 ..^ 1] # Last 2

  var killedCount = 0
  for pid in processesToKill:
    try:
      when defined(windows):
        # Windows termination
        discard execCmd(fmt"taskkill /F /PID {pid}")
      else:
        # Unix termination
        discard execCmd(fmt"kill -9 {pid}")
      killedCount += 1
    except:
      echo fmt"Failed to terminate process {pid}"

  if killedCount > 0:
    styledEcho(
      fgYellow,
      fmt"Killed {killedCount} excess nimsuggest processes, kept {processesToKeep}",
    )

proc formatMemory(kb: float): string =
  if kb < 1024:
    return fmt"{kb:.0f} KB"
  elif kb < 1024 * 1024:
    return fmt"{kb/1024:.1f} MB"
  else:
    return fmt"{kb/(1024.0*1024.0):.2f} GB"

proc displayHeader() =
  let width = 80
  let appTitle = "NimSuggest Monitor"
  let author = "FOXOMAN @ 2025"

  styledEcho(fgCyan, "╔" & "═".repeat(width - 2) & "╗")
  styledEcho(fgCyan, "║" & alignCenter(appTitle, width - 2) & "║")
  styledEcho(fgCyan, "║" & alignCenter(author, width - 2) & "║")
  styledEcho(fgCyan, "╚" & "═".repeat(width - 2) & "╝")

proc displayFooter(killMode: bool) =
  let width = 80

  styledEcho(fgCyan, "╔" & "═".repeat(width - 2) & "╗")
  if killMode:
    styledEcho(
      fgCyan, "║" & alignCenter("Running in kill excess mode (-k)", width - 2) & "║"
    )
  styledEcho(
    fgCyan,
    "║" &
      alignCenter(
        "Press Ctrl+C to exit and terminate all nimsuggest processes", width - 2
      ) & "║",
  )
  styledEcho(fgCyan, "╚" & "═".repeat(width - 2) & "╝")

proc displayProcessInfo(processes: seq[ProcessInfo]) =
  # Clear the screen for a cleaner display
  eraseScreen()
  setCursorPos(0, 0)

  # Display custom header
  displayHeader()
  echo ""

  if processes.len == 0:
    styledEcho(fgGreen, "No nimsuggest processes found")
    echo ""
    displayFooter(config.killExcessProcesses)
    return

  # Calculate totals
  var totalMemory = 0
  for p in processes:
    totalMemory += p.memory

  # Display summary
  styledEcho(fgCyan, "NIMSUGGEST PROCESSES: ", fgYellow, $processes.len)
  styledEcho(
    fgCyan, "TOTAL MEMORY USAGE: ", fgYellow, formatMemory(totalMemory.toFloat)
  )
  echo ""

  # Display table header
  styledEcho(
    fgWhite,
    styleBright,
    alignLeft("PID", 10),
    alignLeft("MEMORY", 15),
    alignLeft("CPU %", 10),
    alignLeft("RUNTIME", 15),
    if processes.len > 2 and config.killExcessProcesses:
      alignLeft("STATUS", 10)
    else:
      "",
  )

  # Sort processes by PID to show in order
  var sortedProcesses = processes
  sortedProcesses.sort(
    proc(x, y: ProcessInfo): int =
      cmp(x.pid, y.pid)
  )

  # Display process info
  for i, p in sortedProcesses:
    var status = ""
    if config.killExcessProcesses:
      if i < sortedProcesses.len - 2:
        status = "TO KILL"
      else:
        status = "KEEP"

    styledEcho(
      fgWhite,
      alignLeft($p.pid, 10),
      if p.memory > 500000: fgRed else: fgGreen,
      alignLeft(formatMemory(p.memory.toFloat()), 15),
      if p.cpu > 50.0: fgRed else: fgGreen,
      alignLeft(fmt"{p.cpu:.1f}%", 10),
      fgWhite,
      alignLeft(p.runtime, 15),
      if status == "TO KILL": fgRed else: fgGreen,
      alignLeft(status, 10),
    )

  echo ""
  # Display footer
  displayFooter(config.killExcessProcesses)

proc main() =
  # Parse command line arguments
  config = parseCommandLine()

  # Register the shutdown handler to catch Ctrl+C
  setControlCHook(shutdownHandler)

  echo "NimSuggest Monitor - Press Ctrl+C to exit"
  if config.killExcessProcesses:
    echo "Running in kill excess mode (-k)"
  sleep(1000)

  # Main monitoring loop
  while true:
    # Find nimsuggest processes
    let pids = findNimsuggestProcesses()

    # Update global list for cleanup
    nimsuggestPids = pids

    # Kill excess processes if enabled
    if config.killExcessProcesses and pids.len > 2:
      killExcessProcesses(pids)
      # Re-scan processes after killing
      let updatedPids = findNimsuggestProcesses()
      nimsuggestPids = updatedPids

    # Get detailed info for each process
    var processInfos: seq[ProcessInfo] = @[]
    for pid in nimsuggestPids:
      processInfos.add(getProcessDetails(pid))

    # Display information
    displayProcessInfo(processInfos)

    # Wait before next update
    sleep(2000)

when isMainModule:
  main()