#!/usr/bin/env ruby
#
# Copyright (C) 2011  Kouhei Sutou <kou@clear-code.com>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

require 'pathname'
require 'time'
require 'optparse'

class Top
  def initialize
    @pids = []
  end

  def <<(pid)
    @pids << pid
  end

  def [](pid)
    processes[pid]
  end

  private
  def processes
    @processes ||= retrieve
  end

  def split_pids
    splitted_pids = []
    @pids.each do |pid|
      splitted_pids << pid
      if splitted_pids.size >= 20
        yield(splitted_pids)
        splitted_pids.clear
      end
    end
    yield(splitted_pids) unless splitted_pids.empty?
  end

  def retrieve
    _processes = {}
    split_pids do |pids|
      result = `env LANG=C top -n 1 -b -p #{pids.join(',')}`
      header, body = result.split(/\n\n/, 2)
      labels = nil
      body.each_line do |line|
        values = line.strip.split(/\s+/, 12)
        if labels.nil?
          labels = values[1..-1]
        else
          pid = values.shift
          _processes[pid] = process = {}
          labels.each_with_index do |label, i|
            process[label] = values[i]
          end
        end
      end
    end
    _processes
  end
end

class MilterStatus
  attr_accessor :top
  def initialize(proc_path)
    @proc_path = proc_path
    @status = parse_status(read_proc_path("status"))
    @top = nil
  end

  def target?(targets)
    _name = name
    return false if _name.empty?
    no_lt_name = _name.sub(/\Alt-/, '')
    targets.any? {|target| _name == target or no_lt_name == target}
  end

  def name
    @name ||= File.basename(command_line[0] || '')
  end

  def pid
    @status["Pid"]
  end

  def vss
    @status["VmSize"]
  end

  def rss
    @status["VmRSS"]
  end

  def cpu_time
    @top["TIME+"]
  end

  def cpu_percent
    @top["%CPU"]
  end

  def n_file_descriptors
    @n_file_descriptors ||= entries("fd").size
  end

  def command_line
    @command_line ||= read_proc_path("cmdline").split(/\0/)
  end

  private
  def read_proc_path(type)
    begin
      (@proc_path + type).read
    rescue SystemCallError
      ""
    end
  end

  def entries(type)
    begin
      (@proc_path + type).entries
    rescue SystemCallError
      []
    end
  end

  def parse_status(status_text)
    status = {}
    status_text.each_line do |line|
      key, value = line.chomp.split(/\s*:\s*/, 2)
      status[key] = value
    end
    status
  end
end

class MilterStatisticsReporter
  def initialize
    @targets = []
    @filters = []
    @interval = 1
  end

  def parse(argv=ARGV)
    @targets = option_parser.parse!(argv)
    if @targets.empty?
      puts option_parser
      exit(false)
    end
  end

  def run
    show_header
    loop do
      begin
        start = Time.now
        report
        report_time = Time.now - start
        sleep(@interval - report_time) if report_time < @interval
      rescue Interrupt
        break
      end
    end
  end

  def show_header
    show("Time", "PID", "VSS", "RSS", "%CPU", "CPU time", "#FD", "command")
  end

  def report
    reported = false
    top = Top.new
    statuses = []
    Pathname.glob("/proc/[0-9]*") do |proc_path|
      status = MilterStatus.new(proc_path)
      next if status.pid.to_i == Process.pid
      next unless status.target?(@targets)
      command_line = status.command_line.join(" ")
      next unless @filters.all? {|filter| filter =~ command_line}
      statuses << status
      top << status.pid
    end
    statuses.each do |status|
      status.top = top[status.pid] || {}
      show(time_stamp,
           status.pid,
           status.vss,
           status.rss,
           status.cpu_percent || "",
           status.cpu_time || "",
           status.n_file_descriptors,
           status.command_line.join(" "))
    end
    puts("%8s not found" % time_stamp) if statuses.empty?
  end

  private
  def option_parser
    @option_parser ||= create_option_parser
  end

  def create_option_parser
    OptionParser.new do |parser|
      parser.banner += " TARGET1 TARGET2 ..."

      parser.on("--filter=REGEXP",
                "Filter report targets by REGEXP.",
                "Multiple --filter options are accepted.") do |regexp|
        @filters << /#{regexp}/i
      end

      parser.on("--interval=INTERVAL", Float,
                "Report each INTERVAL second.") do |interval|
        @interval = interval
      end
    end
  end

  def show(time, pid, vss, rss, cpu_percent, cpu_time, n_fds, command_line)
    items = [time, pid, vss, rss, cpu_percent, cpu_time, n_fds, command_line]
    puts("%8s %6s %9s %9s %5s %8s %5s %s" % items)
  end

  def time_stamp
    Time.now.strftime("%H:%M:%S")
  end
end

if __FILE__ == $0
  reporter = MilterStatisticsReporter.new
  reporter.parse(ARGV)
  reporter.run
end

# vi:ts=2:nowrap:ai:expandtab:sw=2
