#!/usr/bin/env ruby

require 'socket'

INTERVAL = ENV['COLLECTD_INTERVAL'].to_i
HOSTNAME = ENV['COLLECTD_HOSTNAME']


def recv sock
  results = []
  loop do
    line = sock.gets
    if line.start_with? "| OK:"
      break
    else
      results << line.chomp
    end
  end
  results
end

=begin
def parse_table lines
  fields = lines.shift.split(/\|/).collect { |s| s.strip }
  p :fields => fields
  lines.collect { |s|
    data = s.split(/\|/).collect { |s| s.strip }
    obj = {}
    fields.each_with_index do |field,i|
      obj[field] = data[i]
    end
    obj
  }
end
=end

sock = TCPSocket.new "localhost", 5582
loop do
  stats = {}

  sock.puts "c2s:show()"
  recv(sock).each do |line|
    encrypted = line.index(" secure ") or line.index(" TLSv")
    ipv6 = line.index " IPv6 "
    k = "c2s:#{encrypted ? 'encrypted' : 'unencrypted'}-#{ipv6 ? 'ipv6' : 'ipv4'}"
    stats[k] = 0 unless stats[k]
    stats[k] += 1
  end

  sock.puts "s2s:show()"
  recv(sock).each do |line|
    out = if line.index "<-"
            false
          elsif line.index "->"
            true
          else
            next
          end
    encrypted = line.index " TLSv"
    ipv6 = line.index " IPv6 "
    k = "s2s-#{out ? 'out' : 'in'}:#{encrypted ? 'encrypted' : 'unencrypted'}-#{ipv6 ? 'ipv6' : 'ipv4'}"
    stats[k] = 0 unless stats[k]
    stats[k] += 1
  end

  stats.each do |k,v|
    k1, k2 = k.split(":", 2)
    puts "PUTVAL \"#{HOSTNAME}/exec-prosody-#{k1}/current_connections-#{k2}\" interval=#{INTERVAL} N:#{v}"
  end
  STDOUT.flush

ensure
  sleep INTERVAL

end