#!/usr/bin/env ruby

require 'rubygems'
begin
  gem 'instrumental_agent'
rescue Gem::LoadError
  puts "Requires the Instrumental Agent gem:\n"
  puts '  gem install instrumental_agent'
  exit 1
end
require 'instrumental_agent'
require 'pidly'
require 'tmpdir'
require 'optparse'

class SystemInspector
  TYPES = [:gauges, :incrementors]
  attr_accessor *TYPES

  def self.memory
    @memory ||= Memory.new
  end

  def initialize
    @gauges = {}
    @incrementors = {}
    @platform =
      case RUBY_PLATFORM
      when /linux/
        Linux
      when /darwin/
        OSX
      else
        raise "unsupported OS"
      end
  end

  def load_all
    self.class.memory.cycle

    load @platform.load_cpu
    load @platform.load_memory
    load @platform.load_disks
    load @platform.load_filesystem
  end

  def load(stats)
    @gauges.merge!(stats[:gauges] || {})
  end

  module OSX
    def self.load_cpu
      { :gauges => top }
    end

    def self.top
      lines = []
      processes = date = load = cpu = nil
      IO.popen('top -l 1 -n 0') do |top|
        processes = top.gets.split(': ')[1]
        date = top.gets
        load = top.gets.split(': ')[1]
        cpu = top.gets.split(': ')[1]
      end

      user, system, idle = cpu.split(", ").map { |v| v.to_f }
      load1, load5, load15 = load.split(", ").map { |v| v.to_f }
      total, running, stuck, sleeping, threads = processes.split(", ").map { |v| v.to_i }

      {
        'cpu.user' => user,
        'cpu.system' => system,
        'cpu.idle' => idle,
        'load.1min' => load1,
        'load.5min' => load5,
        'load.15min' => load15,
        'processes.total' => total,
        'processes.running' => running,
        'processes.stuck' => stuck,
        'processes.sleeping' => sleeping,
        'threads' => threads,
      }
    end

    def self.load_memory
      # TODO: swap
      { :gauges => vm_stat }
    end

    def self.vm_stat
      header, *rows = `vm_stat`.split("\n")
      page_size = header.match(/page size of (\d+) bytes/)[1].to_i
      sections = ["free", "active", "inactive", "wired", "speculative", "wired down"]
      output = {}
      total = 0.0
      rows.each do |row|
        if match = row.match(/Pages (.*):\s+(\d+)\./)
          section, value = match[1, 2]
          if sections.include?(section)
            value = value.to_f * page_size / 1024 / 1024
            output["memory.#{section.gsub(' ', '_')}_mb"] = value
            total += value
          end
        end
      end
      output["memory.free_percent"] = output["memory.free_mb"] / total * 100 # TODO: verify
      output
    end

    def self.load_disks
      { :gauges => df }
    end

    def self.df
      output = {}
      `df -k`.split("\n").grep(%r{^/dev/}).each do |line|
        device, total, used, available, capacity, mount = line.split(/\s+/)
        names = [File.basename(device)]
        names << 'root' if mount == '/'
        names.each do |name|
          output["disk.#{name}.total_mb"] = total.to_f / 1024
          output["disk.#{name}.used_mb"] = used.to_f / 1024
          output["disk.#{name}.available_mb"] = available.to_f / 1024
          output["disk.#{name}.available_percent"] = available.to_f / total.to_f * 100
        end
      end
      output
    end

    def self.netstat(interface = 'en1')
      # mostly functional network io stats
      headers, *lines = `netstat -ibI #{interface}`.split("\n").map { |l| l.split(/\s+/) } # FIXME: vulnerability?
      headers = headers.map { |h| h.downcase }
      lines.each do |line|
        if !line[3].include?(':')
          return Hash[headers.zip(line)]
        end
      end
    end

    def self.load_filesystem
      {}
    end
  end

  module Linux
    def self.load_cpu
      output = { :gauges => {} }
      output[:gauges].merge!(cpu)
      output[:gauges].merge!(loadavg)
      output
    end

    def self.cpu
      categories = [:user, :nice, :system, :idle, :iowait]
      values = `cat /proc/stat | grep cpu[^0-9]`.chomp.split(/\s+/).slice(1, 5).map(&:to_f)
      SystemInspector.memory.store(:cpu_values, values.dup)
      if previous_values = SystemInspector.memory.retrieve(:cpu_values)
        values.map!.with_index { |value, index| (previous_values[index] - value).abs }
      end
      data = Hash[*categories.zip(values).flatten]
      total = values.inject { |memo, value| memo + value }

      output = {}
      if previous_values
        data.each do |category, value|
          output["cpu.#{category}"] = value / total * 100
        end
      end
      output["cpu.in_use"] = 100 - data['idle'] / total * 100
      output
    end

    def self.loadavg
      min_1, min_5, min_15 = `cat /proc/loadavg`.split(/\s+/)
      {
        'load.1min' => min_1.to_f,
        'load.5min' => min_5.to_f,
        'load.15min' => min_15.to_f,
      }
    end

    def self.load_memory
      output = { :gauges => {} }
      output[:gauges].merge!(memory)
      output[:gauges].merge!(swap)
      output
    end

    def self.memory
      _, total, used, free, shared, buffers, cached = `free -k -o | grep Mem`.chomp.split(/\s+/)
      {
        'memory.used_mb' => used.to_f / 1024,
        'memory.free_mb' => free.to_f / 1024,
        'memory.buffers_mb' => buffers.to_f / 1024,
        'memory.cached_mb' => cached.to_f / 1024,
        'memory.free_percent' => (free.to_f / total.to_f) * 100,
      }
    end

    def self.swap
      _, total, used, free = `free -k -o | grep Swap`.chomp.split(/\s+/)
      {
        'swap.used_mb' => used.to_f / 1024,
        'swap.free_mb' => free.to_f / 1024,
        'swap.free_percent' => (free.to_f / total.to_f) * 100,
      }
    end

    def self.load_disks
      { :gauges => disks.merge(iostat) }
    end

    def self.disks
      output = {}
      `df -Pk`.split("\n").grep(%r{^/dev/}).each do |line|
        device, total, used, available, capacity, mount = line.split(/\s+/)
        names = [File.basename(device)]
        names << 'root' if mount == '/'
        names.each do |name|
          output["disk.#{name}.total_mb"] = total.to_f / 1024
          output["disk.#{name}.used_mb"] = used.to_f / 1024
          output["disk.#{name}.available_mb"] = available.to_f / 1024
          output["disk.#{name}.available_percent"] = available.to_f / total.to_f * 100
        end
      end
      output
    end

    def self.iostat
      output = {}
      if system('iostat > /dev/null 2>&1')
        lines      = `iostat -dx`.split("\n")
        header     = lines.grep(/^Device/i).first
        dev_lines  = lines[(lines.index(header) + 1)..-1]
        util_index = header.split.index('%util')
        dev_lines.each do |dev_line|
          values = dev_line.split
          output["disk.#{values[0]}.percent_utilization"] = values[util_index].to_f
        end
      end
      output
    end

    def self.load_filesystem
      { :gauges => filesystem }
    end

    def self.filesystem
      allocated, unused, max = `sysctl fs.file-nr`.split[-3..-1].map(&:to_i)
      open_files = allocated - unused
      { 'filesystem.open_files' => open_files,
        'filesystem.max_open_files' => max,
        'filesystem.open_files_pct_max' => (open_files.to_f / max.to_f) * 100,
      }
    end
  end

  class Memory
    attr_reader :past_values, :current_values

    def initialize
      @past_values = {}
      @current_values = {}
    end

    def store(attribute, value)
      @current_values[attribute] = value
    end

    def retrieve(attribute)
      @past_values[attribute]
    end

    def cycle
      @past_values = @current_values
      @current_values = {}
    end
  end

end

class ServerController < Pidly::Control
  COMMANDS = [:start, :stop, :status, :restart, :clean, :kill]

  attr_accessor :run_options

  before_start do
    puts "Starting daemon process: #{@pid}"
  end

  start :run

  stop do
    puts "Attempting to kill daemon process: #{@pid}"
  end

  # after_stop :something

  error do
    puts 'Error encountered'
  end

  def self.run(options)
    agent = Instrumental::Agent.new(options[:api_key], :collector => [options[:collector], options[:port]].compact.join(':'))
    puts "Collecting stats under the hostname: #{options[:hostname]}"
    loop do
      sleep 10
      inspector = SystemInspector.new
      inspector.load_all
      inspector.gauges.each do |stat, value|
        agent.gauge("#{options[:hostname]}.#{stat}", value)
      end
      # agent.increment("#{host}.#{stat}", delta)
    end
  end

  def initialize(options={})
    @run_options = options.delete(:run_options) || {}
    super(options)
  end

  def run
    self.class.run(run_options)
  end

  alias_method :clean, :clean!
end

def terminal_messages
  if `which iostat`.chomp.empty?
    puts 'Install iostat (sysstat package) to collect disk I/O metrics.'
  end
end

def require_api_key(options, parser)
  unless options[:api_key] # present?
    print parser.help
    exit 1
  end
end

options = {
  :daemon => false,
  :collector => 'instrumentalapp.com',
  :port => '8000',
  :hostname => `hostname`.chomp,
}

option_parser = OptionParser.new do |opts|
  opts.banner = "Usage: instrument_server [-k API_KEY] [options] [-d #{ServerController::COMMANDS.join('|')}]"
  opts.on('-H', '--hostname HOSTNAME', 'Hostname to report as') do |hostname|
    options[:hostname] = hostname
  end
  opts.on('-k', '--api_key API_KEY', 'API key of your project') do |api_key|
    options[:api_key] = api_key
  end
  opts.on('-c', '--collector COLLECTOR[:PORT]', "Collector (default #{options[:collector]}:#{options[:port]})") do |collector|
    address, port = collector.split(':')
    options[:collector] = address
    options[:port] = port if port
  end
  opts.on('-d', '--daemonize', 'Run as daemon') do
    options[:daemon] = true
  end
  opts.on('-h', '--help', 'Display this screen') do
    puts opts
    terminal_messages
    exit
  end
end

option_parser.parse!
options[:api_key] ||= ENV["INSTRUMENTAL_TOKEN"]

if options.delete(:daemon)
  @server_controller = ServerController.spawn(
    :name => 'instrument_server',
    :path => Dir.tmpdir,
    :pid_file => File.join(Dir.tmpdir, 'pids', 'instrument_server.pid'),
    :verbose => true,
    # :signal => 'kill',
    :log_file => File.join(Dir.tmpdir, 'log', 'instrument_server.log'),
    :run_options => options
  )
  command = ARGV.first && ARGV.first.to_sym
  command ||= :start
  if [:start, :restart].include?(command)
    require_api_key(options, option_parser)
  end
  if ServerController::COMMANDS.include?(command)
    @server_controller.send command
  else
    puts "Command must be one of: #{ServerController::COMMANDS.join(', ')}"
  end
else
  require_api_key(options, option_parser)
  ServerController.run(options)
end