#!/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

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

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

  sock.puts "s2s:show()"
  recv(sock).each do |line|
    if line.start_with? "|  "
      out = if line.index "<-"
              false
            elsif line.index "->"
              true
            else
              next
            end
      encrypted = line.index " (encrypted)"
      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
  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

ensure
  sleep INTERVAL

end