#!/usr/bin/env ruby
#
# Copyright (C) 2010  Kouhei Sutou <kou@clear-code.com>
#
# This library is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This library 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 Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with this library.  If not, see <http://www.gnu.org/licenses/>.
#
#%# family=auto
#%# capabilities=autoconf suggest

type, milter = File.basename($0).gsub(/milter_manager_/, '').split(/_/, 2)
milter = nil if milter and milter.empty?

mode = ARGV[0]

log_dir = ENV["logdir"] || File.expand_path("~milter-manager/public_html/log")

def rrdtool(command, *arguments)
  command_line = "rrdtool #{command} #{arguments.join(' ')}"
  result = `#{command_line}`
  if $?.success?
    result
  else
    $stderr.puts("failed to run rrdtool: <#{command_line}>: <#{result}>")
    nil
  end
end

def normalize_name(name)
  name.gsub(/-/, "_")
end

def fetch(rrd)
  result = rrdtool("fetch", "#{rrd} MAX -r 1 -s -15min")
  return [[], []] if result.nil?
  lines = result.split(/\r?\n/)
  header = lines[0].split.collect {|name| normalize_name(name)}
  values = nil
  lines[2..-1].reverse_each do |line|
    time, *_values = lines[2].split
    if _values[0] != "nan"
      values = _values
      break
    end
  end
  values ||= header.collect {"nan"}
  [header, values]
end

def data_sets(rrd)
  result = rrdtool("info", rrd)
  return result if result.nil?
  sets = {}
  name_order = []
  result.each_line do |line|
    case line
    when /\Ads\[(.+)\].(.+) = (.+)$/
      name, attribute, value = $1, $2, $3
      name = normalize_name(name)
      name_order << name
      sets[name] ||= {}
      sets[name][attribute] = value
    end
  end
  ordered_sets = []
  name_order.uniq.each do |name|
    ordered_sets << [name, sets[name]]
  end
  ordered_sets
end

def rrd_path(log_dir, type, milter)
  if milter.nil?
    base_name = "milter-manager.#{type}.rrd"
  else
    base_name = "milter.#{type}.#{milter}.rrd"
  end
  File.join(log_dir, base_name)
end

def config(log_dir, type, milter, parameters)
  sets = data_sets(rrd_path(log_dir, type, milter))
  return if sets.nil? or sets.empty?
  title = parameters[:title]
  vlabel = parameters[:vlabel]
  category = parameters[:category]
  info = parameters[:info]
  labels = parameters[:labels]
  order = (parameters[:order] || sets.collect {|key,| key}).reject do |key|
    sets.assoc(key).nil?
  end
  unordered_keys = []
  sets.each do |key,|
    unordered_keys << key unless order.include?(key)
  end
  order.concat(unordered_keys)
  case type
  when "status"
    order << "smtp"
    order << "disconnected"
    order << "concurrent"
  end
  print(<<-EOM)
graph_title #{title}
graph_order #{order.join(' ')}
graph_vlabel #{vlabel}
graph_category #{category}
graph_info #{info}

EOM

  sets.each_with_index do |(key, attributes), i|
    puts("#{key}.label #{labels[key] || key}")
    draw_type = i.zero? ? "AREA" : "STACK"
    puts("#{key}.draw #{draw_type}")
    attributes.each do |name, value|
      if ["min", "max", "type"].include?(name)
        if /\A\"(.*)\"\z/ =~ value
          value = $1
        end
        puts("#{key}.#{name} #{value}")
      end
    end
  end
  puts
  case type
  when "status"
    puts("smtp.label SMTP")
    puts("smtp.draw LINE2")
    puts("smtp.type GAUGE")
    puts
    puts("disconnected.label Disconnected")
    puts("disconnected.draw LINE2")
    puts("disconnected.type GAUGE")
    puts
    puts("concurrent.label Concurrent")
    puts("concurrent.draw LINE2")
    puts("concurrent.type GAUGE")
  end
end

def report(log_dir, type, milter)
  header, values = fetch(rrd_path(log_dir, type, milter))
  return if header.empty?
  header.each_with_index do |name, index|
    puts("#{name}.value #{values[index]}")
  end
  case type
  when "status"
    header, values = fetch(File.join(log_dir, "session.rrd"))
    ["smtp", "disconnected", "concurrent"].each do |key|
      puts("#{key}.value #{values[header.index(key)]}")
    end
  end
end

def invalid_file_name_message
  "must be 'milter_manager_status_XXX' or " +
    "'milter_manager_report_XXX': #{File.basename($0)}"
end

case mode
when "auto", "autoconf", "detect"
  if File.exist?(log_dir)
    puts "yes"
    exit(true)
  else
    puts "no (log directory not found: #{log_dir})"
    exit(false)
  end
when "suggest"
  type_order = ["status", "report"]
  milters = []
  Dir.glob(File.join(log_dir, "milter.{status,report}.*.rrd")) do |path|
    if /\Amilter\.(status|report)\.(.+?).rrd\z/ =~ File.basename(path)
      type, name = $1, $2
      milters << [name, type]
    end
  end

  type_order.each do |type|
    puts type
  end
  milters.sort_by do |name, type|
    [name, type_order.index(type)]
  end.each do |name, type|
    puts "#{type}_#{name}"
  end
  exit(true)
when "config"
  if milter.nil?
    target = "milter manager"
    title_prefix = target
  else
    target = milter
    title_prefix = "milter: #{target}"
  end

  case type
  when "status"
    config(log_dir, type, milter,
           :title => "#{title_prefix}: status",
           :vlabel => "sessions/min",
           :category => "milter-manager",
           :info => "This graph shows returned statuses by #{target}",
           :labels => {
             "pass" => "Pass",
             "accept" => "Accept",
             "reject" => "Reject",
             "discard" => "Discard",
             "temporary_failure" => "Temp-Fail",
             "quarantine" => "Quarantine",
             "abort" => "Abort",
             "error" => "Error",
             "stop" => "Stop",
           })
  when "report"
    config(log_dir, type, milter,
           :title => "#{title_prefix}: report",
           :vlabel => "mails/min",
           :category => "milter-manager",
           :info => "This graph reports mail kind of #{target}",
           :labels => {
             "spam" => "spam",
             "virus" => "Virus",
             "uribl" => "URIBL",
             "greylisting_pass" => "Greylisting (pass)",
             "tarpitting_pass" => "Tarpitting (pass)",
             "spf_pass" => "SPF (pass)",
             "spf_fail" => "SPF (fail)",
             "sender_id_pass" => "Sender ID (pass)",
             "sender_id_fail" => "Sender ID (fail)",
             "dkim_pass" => "DKIM (pass)",
             "dkim_fail" => "DKIM (fail)",
             "dkim_adsp_pass" => "DKIM ADSP (pass)",
             "dkim_adsp_fail" => "DKIM ADSP (fail)",
           },
           :order => ["spam",
                      "virus",
                      "uribl",
                      "greylisting_pass",
                      "tarpitting_pass",
                      "spf_pass",
                      "spf_fail",
                      "sender_id_pass",
                      "sender_id_fail",
                      "dkim_pass",
                      "dkim_fail",
                      "dkim_adsp_pass",
                      "dkim_adsp_fail",
                     ])
  else
    puts invalid_file_name_message
    exit(false)
  end
  exit(true)
end

case type
when "status", "report"
  report(log_dir, type, milter)
else
  puts invalid_file_name_message
  exit(false)
end
