#!/usr/bin/env ruby

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

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"
@path = ENV["path"]

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)
  path = Shellwords.shellescape(@path)
  result = `#{@groonga} #{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

case command
when "autoconf", "detect"
  success, body = run("status")
  if success
    puts "yes"
    exit(true)
  else
    puts "no (#{body})"
    exit(false)
  end
when "config"
  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

path = Shellwords.shellescape(@path)
du_result = `#{@du} -B1 #{path}*`
unless $?.success?
  puts("error: #{du_result}")
  exit(false)
end
usages = parse_du_result(du_result)
usage = compute_size(usages, @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
