#!/usr/bin/perl

use strict;
use warnings;

use FindBin;
use lib "$FindBin::Bin/../../lib";
use Wigo::Probe qw/:all/;

use IO::Socket::INET;

###
# DEFAULT CONFIG
###

my $conf = {
    'instances' => [
       # 'localhost:2181',
	],
    'details' => [
        'zk_version',
        'zk_server_state'
    ],
    'metrics' => [
        'zk_avg_latency',
        'zk_max_latency',
        'zk_min_latency',
        'zk_packets_received',
        'zk_packets_sent',
        'zk_num_alive_connections',
        'zk_outstanding_requests',
        'zk_znode_count',
        'zk_watch_count',
        'zk_ephemerals_count',
        'zk_approximate_data_size',
        'zk_open_file_descriptor_count',
        'zk_max_file_descriptor_count',
        'zk_followers',
        'zk_synced_followers',
        'zk_pending_syncs',
    ],
};

init( config => $conf );

if ( ! scalar @{config->{'instances'}} )
{
    message "Nothing to monitor";
    exit 13;
}

###
# GET STATUS
###

my @messages;
for my $instance ( @{config->{'instances'}} )
{
    my $socket;
    eval {
        local $SIG{ALRM} = sub { die "connect timeout\n" }; # NB: \n required
        alarm 1;
        
        $socket = new IO::Socket::INET(
            'PeerAddr'  => $instance,
            'Proto'     => 'tcp',
        ) or die("$!\n");

        alarm 0;
    };

    if ( $@ )
    {
        chomp $@;
        raise 300;
        push @messages, "$instance : $@";
        next;
    }
    
    my $response;
    eval {
        local $SIG{ALRM} = sub { die "ruok timeout\n" }; # NB: \n required
        alarm 1;
        
        $socket->send("ruok\n");
        $socket->recv($response, 1024);

        alarm 0;
    };
    
    if ( $@ )
    {
        chomp $@;
        raise 300;
        push @messages, "$instance : $@";
        next;
    }

    push @messages, "$instance : $response";

    if ( $response !~ /^imok/ )
    {
        raise 300;
    }

    $socket->close();
    
    eval {
        local $SIG{ALRM} = sub { die "mntr connect timeout\n" }; # NB: \n required
        alarm 1;

        $socket = new IO::Socket::INET(
            'PeerAddr'  => $instance,
            'Proto'     => 'tcp',
        ) or die("$!\n");

        alarm 0;
    };

    if ( $@ )
    {
        chomp $@;
        raise 300;
        push @messages, "$@";
        next;
    } 

    eval {
        local $SIG{ALRM} = sub { die "mntr timeout\n" }; # NB: \n required
        alarm 1;

        $socket->send("mntr\n");
        $socket->recv($response, 1024);

        alarm 0;
    };
    
    if ( $@ )
    {
        chomp $@;
        raise 300;
        push @messages, "$instance : $@";
        next;
    }
  
    $socket->close();

    for my $line ( split "\n", $response )
    {
        if ( my ( $key, $value ) = $line =~ /(\S+)\s+(.*)/ )
        {
            if ( grep { $key eq $_ } @{config->{'metrics'}} )
            {
                add_metric { 'Tags' => { 'instance' => "$instance", 'metric' => $key }, 'Value' => $value };
            }
            
            if ( grep { $key eq $_ } @{config->{'details'}} )
            {
                detail->{$instance}->{$key} = $value;
            }
        }
    }
}

message join( " , " , @messages ) if scalar @messages;

output 0;