import base64
import urllib2
import socket
socket.setdefaulttimeout(5)

try:
    from simplejson import dumps, loads
except ImportError:
    from json import dumps, loads

def perfdata(rbmq_queue_len):
    return "rabbitmq_queue_len=%s" % (rbmq_queue_len)

def __run__(params):
    try:
	if params.get("timeout"):
            timeout = params.get("timeout")
            socket.setdefaulttimeout(timeout)

        host      = params.get("host", "localhost")
        port      = params.get("port", "55672")
        vhost     = params.get("vhost")
        queue     = params.get("queue")
        user      = params.get("user", "guest")
        passwd    = params.get("passwd", "guest")
        check_workers = params.get("check_workers", "true")
        maxjobs_c = int(params.get("maxjobs", "100"))
        maxjobs_w = maxjobs_c * 0.7
        br = urllib2.Request("http://%s:%s/api/queues" % (host, port))
        br.add_header("Authorization", "Basic %s" % base64.encodestring("%s:%s" % (user, passwd)).strip())
        rbmq_json = loads(urllib2.urlopen(br).read().strip())
        if not queue and not vhost:
            rbmq_queue_len = sum([q["messages"] for q in rbmq_json if q['node'] == 'rabbit@%s' % socket.gethostname()])
            queue = "all"
        else:
            try:
                rbmq_queue = filter(lambda q: q["vhost"] == vhost and q["name"] == queue, rbmq_json)[0]
            except IndexError:
                return [2, "CRITICAL - Queue not found: %s" % queue]
        
            rbmq_queue_len = int(rbmq_queue["messages"])
        if rbmq_queue_len >= maxjobs_c:
            return [2, "CRITICAL - Queue '%s' has %s pending jobs" % (queue, rbmq_queue_len), perfdata(rbmq_queue_len)]
        if rbmq_queue_len >= maxjobs_w:
            return [1, "WARNING - Queue '%s' has %s pending jobs" % (queue, rbmq_queue_len), perfdata(rbmq_queue_len)]
        if check_workers.strip().lower() in ("true", "yes") and int(rbmq_queue["consumers"]) == 0:
            return [1, "WARNING - No workers connected"]

    except socket.timeout:
        return [2, "CRITICAL: Socket timeout"]
    except socket.gaierror:
        return [2, "CRITICAL: Unable to resolve domain"]
    except Exception, e:
        return [2,"CRITICAL: %s" % repr(e)]
    else:
        return [0, "OK - Queue '%s' has '%s' jobs" % (queue, rbmq_queue_len), perfdata(rbmq_queue_len)]