#!/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 (physical footprint via top, joined with ps) ---
#
# We use physical footprint (top's MEM column) rather than ps RSS. RSS counts
# every resident page, including shared read-only pages like the dyld shared
# cache and frameworks — so summing it across many processes that share those
# pages (e.g. 300+ iOS simulator daemons) overcounts wildly. Footprint charges
# shared pages once and matches what Activity Monitor reports. top truncates
# command names though, so we pull full paths from ps and join on PID.

def parse_mem_mb(str)
  m = str.match(/\A([\d.]+)([KMGB]?)/)
  return nil unless m

  val = m[1].to_f
  case m[2]
  when "G" then val * 1024
  when "M" then val
  when "K" then val / 1024.0
  else val / 1024.0 / 1024.0 # bytes
  end
end

footprints = {}
run("top -l 1 -stats pid,mem").each_line do |line|
  parts = line.strip.split(/\s+/)
  next unless parts.size == 2 && parts[0].match?(/\A\d+\z/)

  mb = parse_mem_mb(parts[1])
  footprints[parts[0].to_i] = mb if mb
end

commands = {}
run("ps -axo pid=,comm=").each_line do |line|
  pid, comm = line.strip.split(/\s+/, 2)
  commands[pid.to_i] = comm if pid && comm
end

processes = footprints.filter_map do |pid, mem_mb|
  next if mem_mb < 0.1

  command = commands[pid]
  next unless command

  { pid:, mem_mb:, command: }
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] = { mem_mb: 0.0, count: 0 } }
processes.each do |proc|
  name = group_name(proc[:command])
  groups[name][:mem_mb] += proc[:mem_mb]
  groups[name][:count] += 1
end

top_groups = groups.sort_by { -it[1][:mem_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, footprint_mb: info[:mem_mb].round(1), process_count: info[:count] }
    end,
    note: "Physical footprint (matches Activity Monitor); shared pages are charged once.",
  }
  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", "MEM", "PROCS"]
top_groups.each do |name, info|
  puts "  %-28s %10s %8d" % [name, human_bytes(info[:mem_mb]), info[:count]]
end
puts
puts "#{dim}Note: physical footprint (matches Activity Monitor); shared pages charged once.#{reset}"