#!/usr/bin/ruby

#%# family=auto
#%# capabilities=autoconf suggest

require 'shellwords'
begin
  require 'json'
rescue LoadError
  require 'rubygems'
  require 'json'
end
require 'English'

label = ENV["label"]
@groonga = ENV["groonga"] || "groonga"
@du = ENV["du"] || "du"

command = ARGV.shift

def parse(success, result)
  if success
    begin
      status, body = JSON.parse(result)
      return_code, start_time, elapsed, error_message = status
      if return_code.zero?
        [success, body]
      else
        [false, error_message]
      end
    rescue JSON::ParserError
      [false, $!.message]
    end
  else
    [success, result]
  end
end

def run(command, *args)
  database_path = Shellwords.shellescape(@database_path)
  result = `#{@groonga} #{database_path} #{command} #{args.join(' ')} 2>&1`
  parse($?.success?, result)
end

def parse_list(header, list)
  list.collect do |item|
    parsed_item = {}
    header.each_with_index do |(name, type), i|
      parsed_item[name] = item[i]
    end
    parsed_item
  end
end

def schema
  tables = []
  success, table_list_body = run("table_list")
  unless success
    puts("error: #{table_list_body}")
    exit(false)
  end
  parse_list(table_list_body[0], table_list_body[1..-1]).each do |table|
    table_name = table["name"]
    table["key"] = "table_#{table_name}"
    success, column_list_body = run("column_list", table_name)
    unless success
      puts("error: #{column_list_body}")
      exit(false)
    end
    table["columns"] = parse_list(column_list_body[0], column_list_body[1..-1])
    table["columns"].each do |column|
      column["key"] = "column_#{table_name}_#{column['name']}"
      column["full_name"] = "#{table_name}.#{column['name']}"
    end
    tables << table
  end
  tables
end

def parse_du_result(result)
  usages = {}
  result.each_line do |line|
    if /\A(\d+)\s+/ =~ line
      usage = $1
      path = $POSTMATCH.strip
      usages[path] = usage.to_i
    end
  end
  usages
end

def compute_size(usages, base_path)
  usage = 0
  return usage if base_path.empty?

  usages.each do |path, size|
    usage += size if path.start_with?(base_path)
  end
  usage
end

def setup_database_path
  if /_([^_]+)\z/ =~ $0
    service_type = $1
    database_path_variable_name = "#{service_type}_database_path"
    database_path = ENV[database_path_variable_name]
  else
    database_path_variable_name = "database_path"
    database_path = ENV[database_path_variable_name]
    # For backward compatibility. Remove me when 5.0.0.
    database_path ||= ENV["path"]
  end

  if database_path.nil?
    key = "env.#{database_path_variable_name}"
    $stderr.puts("Database path isn't specified by #{key}")
    exit(false)
  end

  unless File.exist?(database_path)
    $stderr.puts("Database path doesn't exist: #{database_path}")
    exit(false)
  end

  @database_path = database_path
end

def autoconf
  database_path_variable_names = [
    "database_path",
    "path", # For backward compatibility. Remove me when 5.0.0.
    "http_database_path",
    "httpd_database_path",
    "gqtp_database_path",
  ]
  database_paths = database_path_variable_names.collect do |variable_name|
    ENV[variable_name]
  end
  database_paths = database_paths.compact
  if database_paths.empty?
    variable_names = database_path_variable_names.collect do |name|
      "env.#{name}"
    end
    variable_names_label = variable_names[0..-2].join(", ")
    variable_names_label << " and/or #{variable_names.last}"
    puts("no (No database path is specified. Specify #{variable_names_label}.)")
    exit(false)
  end
  database_paths.each do |database_path|
    next unless File.exist?(database_path)
    puts("yes")
    exit(true)
  end
  database_paths_label = database_paths.join(", ")
  puts("no (All database paths don't exist: #{database_paths_label})")
  exit(false)
end

def suggest
  env_path_exist_p = lambda do |variable_name|
    path = ENV[variable_name]
    path and File.exist?(path)
  end
  exist_p = lambda do |database_path_name, pid_file_path_name|
    env_path_exist_p.call(database_path_name) and
      env_path_exist_p.call(pid_file_path_name)
  end
  if exist_p.call("http_database_path", "http_pid_path")
    puts("http")
  end
  if exist_p.call("httpd_database_path", "httpd_pid_path")
    puts("httpd")
  end
  if exist_p.call("gqtp_database_path", "gqtp_pid_path")
    puts("gqtp")
  end
  exit(true)
end

case command
when "autoconf", "detect"
  autoconf
when "suggest"
  suggest
when "config"
  setup_database_path
  if label.nil?
    title = "groonga: disk usage"
  else
    title = "groonga: #{label}: disk usage"
  end
  puts(<<EOF)
graph_title #{title}
graph_vlabel Bytes
graph_category groonga
graph_info disk usage in groonga tables and columns
graph_args --base 1024
graph_total Total

database.label Database
database.draw AREA
EOF
  schema.each do |table|
    table_key = table["key"]
    table_name = table["name"]
    puts(<<EOF)

#{table_key}.label #{table_name}
#{table_key}.draw STACK
EOF
    table["columns"].each do |column|
      column_key = column["key"]
      column_name = column["full_name"]
      puts(<<EOF)

#{column_key}.label #{column_name}
#{column_key}.draw STACK
EOF
    end
  end
  exit(true)
end

setup_database_path
database_path = Shellwords.shellescape(@database_path)
du_result = `#{@du} -B1 #{database_path}*`
unless $?.success?
  $stderr.puts("error: #{du_result}")
  exit(false)
end
usages = parse_du_result(du_result)
usage = compute_size(usages, @database_path)
puts(<<EOF)
database.value #{usage}
EOF
schema.each do |table|
  table_key = table["key"]
  table_name = table["name"]
  usage = compute_size(usages, table["path"])
  puts(<<EOF)
#{table_key}.value #{usage}
EOF
  table["columns"].each do |column|
    column_key = column["key"]
    usage = compute_size(usages, column["path"])
    puts(<<EOF)
#{column_key}.value #{usage}
EOF
  end
end