#!/usr/bin/env ruby
# frozen_string_literal: true

require "json"
require "optparse"
require "open3"

options = { group_top: 20, json: false, color: true }

OptionParser.new do |opts|
  opts.banner = "Usage: mem-report [options]"
  opts.on("--group-top N", Integer, "Number of top groups to show (default: 20)") { options[:group_top] = it }
  opts.on("--json", "Output as JSON") { options[:json] = true }
  opts.on("--no-color", "Disable colored output") { options[:color] = false }
end.parse!

def run(cmd) = `#{cmd}`.strip

def human_bytes(mb)
  if mb >= 1024
    format("%.1f GB", mb / 1024.0)
  else
    format("%.0f MB", mb)
  end
end

# --- System memory summary ---

page_size = run("sysctl -n hw.pagesize").to_i
phys_mem_bytes = run("sysctl -n hw.memsize").to_i
phys_mem_mb = phys_mem_bytes / 1024.0 / 1024.0

vm_stat_output = run("vm_stat")
vm = {}
vm_stat_output.each_line do |line|
  case line
  when /Pages free:\s+([\d.]+)/          then vm[:free] = $1.to_i
  when /Pages active:\s+([\d.]+)/        then vm[:active] = $1.to_i
  when /Pages inactive:\s+([\d.]+)/      then vm[:inactive] = $1.to_i
  when /Pages speculative:\s+([\d.]+)/   then vm[:speculative] = $1.to_i
  when /Pages wired down:\s+([\d.]+)/    then vm[:wired] = $1.to_i
  when /Pages occupied by compressor:\s+([\d.]+)/ then vm[:compressor] = $1.to_i
  when /File-backed pages:\s+([\d.]+)/   then vm[:file_backed] = $1.to_i
  when /Pages purgeable:\s+([\d.]+)/     then vm[:purgeable] = $1.to_i
  end
end

to_mb = ->(pages) { (pages * page_size) / 1024.0 / 1024.0 }

used_mb = to_mb.call((vm[:active] || 0) + (vm[:wired] || 0) + (vm[:compressor] || 0) + (vm[:inactive] || 0))
cache_mb = to_mb.call((vm[:file_backed] || 0) + (vm[:purgeable] || 0))
app_mb = [used_mb - cache_mb, 0].max

system_summary = {
  physical_memory: phys_mem_mb.round(0),
  used_memory: used_mb.round(0),
  filesystem_cache: cache_mb.round(0),
  app_memory: app_mb.round(0),
  free_memory: to_mb.call((vm[:free] || 0) + (vm[:speculative] || 0)).round(0),
}

# --- Per-process memory (via ps) ---

ps_output = run("ps -axo pid=,user=,rss=,comm=")
processes = ps_output.each_line.filter_map do |line|
  parts = line.strip.split(/\s+/, 4)
  next if parts.size < 4

  pid, user, rss_kb, comm = parts
  rss_mb = rss_kb.to_f / 1024.0
  next if rss_mb < 0.1

  { pid: pid.to_i, user:, rss_mb:, command: comm }
end

# --- Grouped memory ---

GROUP_RULES = [
  [/com\.apple\.WebKit|Safari/i, "Safari / WebKit"],
  [/Simulator|CoreSimulator|simctl/i, "Simulator"],
  [/Xcode|SourceKit|clang|swift-frontend|IBAgent|XCBBuildService|devicectl/i, "Xcode Toolchain"],
  [/OrbStack|orbctl|orb-/i, "OrbStack"],
  [/Google Chrome|chrome_/i, "Google Chrome"],
  [/Arc|arc\./i, "Arc Browser"],
  [/Firefox|firefox/i, "Firefox"],
  [/Slack/i, "Slack"],
  [/Discord/i, "Discord"],
  [/Spotify/i, "Spotify"],
  [/docker|Docker/i, "Docker"],
  [/Terminal|iTerm|Alacritty|kitty|ghostty|WezTerm/i, "Terminal"],
  [/claude/i, "Claude"],
  [/WindowServer/, "WindowServer"],
  [/kernel_task/, "kernel_task"],
]

def group_name(command)
  GROUP_RULES.each do |pattern, name|
    return name if command.match?(pattern)
  end
  File.basename(command).sub(/\.\w+$/, "")
end

groups = Hash.new { |h, k| h[k] = { rss_mb: 0.0, count: 0 } }
processes.each do |proc|
  name = group_name(proc[:command])
  groups[name][:rss_mb] += proc[:rss_mb]
  groups[name][:count] += 1
end

top_groups = groups.sort_by { -it[1][:rss_mb] }.first(options[:group_top])

# --- Output ---

if options[:json]
  data = {
    system: system_summary.transform_values(&:to_i),
    top_groups: top_groups.map do |name, info|
      { group: name, rss_mb: info[:rss_mb].round(1), process_count: info[:count] }
    end,
    note: "RSS values and group sums can overcount shared memory (shared libraries, frameworks, mmap regions).",
  }
  puts JSON.pretty_generate(data)
  exit
end

bold  = options[:color] ? "\e[1m" : ""
dim   = options[:color] ? "\e[2m" : ""
cyan  = options[:color] ? "\e[36m" : ""
reset = options[:color] ? "\e[0m" : ""

puts "#{bold}=== System Memory ==#{reset}"
puts
puts "  Physical memory:    #{cyan}#{human_bytes(system_summary[:physical_memory])}#{reset}"
puts "  Used memory:        #{human_bytes(system_summary[:used_memory])}"
puts "  Filesystem cache:   #{human_bytes(system_summary[:filesystem_cache])}"
puts "  App memory:         #{cyan}#{human_bytes(system_summary[:app_memory])}#{reset}"
puts "  Free memory:        #{human_bytes(system_summary[:free_memory])}"
puts

puts "#{bold}=== Top #{options[:group_top]} Grouped Programs ==#{reset}"
puts
puts "  #{dim}%-28s %10s %8s#{reset}" % ["GROUP", "RSS", "PROCS"]
top_groups.each do |name, info|
  puts "  %-28s %10s %8d" % [name, human_bytes(info[:rss_mb]), info[:count]]
end
puts
puts "#{dim}Note: RSS and group sums can overcount shared memory (shared libs, frameworks, mmap).#{reset}"