#!/usr/bin/ruby

require 'getoptlong'
require 'ftools'

MODULES_DIR = '../modules'
TRACE_DIR = '/proc/kmemprof-trace'

$mod_list = [
  TRACE = "kmemprof-trace",
  CGPROF = "kmemprof-cgprof",
  OBJFILT = "kmemprof-objfilter",
  FILT = "kmemprof-filter",
  FREQCTRL = "kmemprof-freqctrl",
  PLUG = "kmemprof-plug"
  ]

$load_status = {}

IO.popen("lsmod").readlines.slice(1..-1).collect{|l| l.split[0]}.each do |m|
  $load_status[m] = 1
end
  
def loaded? (mod)
  return true if $load_status[mod] == 1
  false
end

def insmod (mod, *args)
  mod = MODULES_DIR + '/' + mod + '.o'
  comm = ['insmod', mod, *args]
  system(comm.join(' '))
end
  
def rmmod (mod, *args)
  mod = mod
  comm = ['rmmod', mod, *args]
  system(comm.join(' '))
end

def insmod_if_not_loaded(mod)
  insmod mod if ! loaded? mod
end

def rmmod_if_loaded(mod)
  rmmod mod if loaded? mod
end

########

def save_ksyms
  proc {|file| File.copy '/proc/ksyms', file.to_s }
end

def set_freq
  proc do |arg|
    rmmod_if_loaded (FREQCTRL)
    if arg.to_i != 1
      insmod FREQCTRL, 'freq='+([arg]*32).join(',')
    end
  end
end

def set_filt
  proc do |arg|
    rmmod_if_loaded (FILT)
    if arg != ''
      insmod FILT, arg.to_s
    end
  end
end

def set_objfilt
  proc do |arg|
    rmmod_if_loaded(OBJFILT)
    if arg != ''
      insmod OBJFILT, 'objfilter=' + arg.to_s
    end
  end
end

def load_cgprof
  proc do
    # if TRACE has already been loaded...

    insmod_if_not_loaded (PLUG)
    insmod CGPROF, `ruby cgprof-config.rb`
  end
end

def unload_cgprof
  proc { rmmod_if_loaded(CGPROF) }
end

def load_trace
  proc do |arg|
    if arg == ""
      nr_cpus=1
    else
      nr_cpus=arg.to_i
    end
    # if TRACE has already been loaded...

    insmod_if_not_loaded (PLUG)
    insmod TRACE, 'num_cpus=' + nr_cpus.to_s
  end
end

def unload_trace
  proc { rmmod_if_loaded(TRACE) }
end

def start
  proc { rmmod_if_loaded(PLUG) }
end

def stop
  proc { insmod_if_not_loaded(PLUG) }
end

def shutdown
  proc { $mod_list.each {|x| rmmod_if_loaded(x)} }
end

opt_handlers = {
  [ "--save-ksyms", "-l", GetoptLong::REQUIRED_ARGUMENT] =>
  save_ksyms,

  [ "--set-frequency", "-f", GetoptLong::REQUIRED_ARGUMENT] =>
  set_freq,
  [ "--set-filter", "-r", GetoptLong::OPTIONAL_ARGUMENT] =>
  set_filt,
  [ "--set-objfilter", "-o", GetoptLong::OPTIONAL_ARGUMENT] =>
  set_objfilt,

  [ "--load-cgprof", "-p", GetoptLong::NO_ARGUMENT] =>
  load_cgprof,

  [ "--unload-cgprof", "-q", GetoptLong::NO_ARGUMENT] =>
  unload_cgprof,

  [ "--load-trace", "-t", GetoptLong::OPTIONAL_ARGUMENT] => load_trace,
  [ "--unload-trace", "-u", GetoptLong::OPTIONAL_ARGUMENT] => unload_trace,

  [ "--start", "-b", GetoptLong::NO_ARGUMENT] => start,
  [ "--stop", "-e", GetoptLong::NO_ARGUMENT] => stop,
  [ "--shutdown", "-d", GetoptLong::NO_ARGUMENT] => shutdown
  }

opts = GetoptLong.new()
opts.set_options(*(opt_handlers.keys))

def expire(arg, &block)
  block.call(arg)
end

opts.each do |opt, arg|
  opt_handlers.keys.each do |k|
    if (k[0] == opt)
      expire(arg, &opt_handlers[k])
    end
  end
end

