import yaml
import socket

def format(number):
    for x in ["bytes", "KB", "MB", "GB"]:
        if number < 1024.0 and number > -1024.0:
            return "%3.2f %s" % (number, x)
        number /= 1024.0
    return "%3.1f%s" % (number, "TB")

def __run__(params):
    try:
        host = params.get("host", "localhost")
        port = int(params.get("port", "11211"))
        setget = params.get("setget", False)
        thresold = yaml.load(params.get("thresholds"))
    except Exception, e:
        return [2, "CRITICAL - Error: %s" % repr(e)]

    try:
        socket.setdefaulttimeout(5)
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.connect((host, port))
        sock.send("stats\n")
    except socket.error, e:
        return [2, "CRITICAL - Error: %s" % repr(e)]
    except Exception, e:
        return [2, "CRITICAL - Error: %s" % repr(e)]

    try:
        line = ""
        while True:
            line += sock.recv(4096)
            if line[-1] == "\n":
                sock.close()
                break
    except Exception, e:
        return [2, "CRITICAL - Error: %s" % repr(e)]

    if setget:
        try:
            socket.setdefaulttimeout(5)
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.connect((host, port))
            sock.send("set MEMCACHETEST 0 10 2\nOK\n")
            value = sock.send("get MEMCACHETEST\n")
        except:
            return [2, "CRITICAL - Error: Unable to set and get value on memcache" ]

    stats = {}
    [ stats.update({data.split()[1]: data.split()[2]}) for data in line.split("\r\n") if "STAT" in data ]

    used_memory = ((int(stats["bytes"]) * 100) / int(stats["limit_maxbytes"]))
    total_memory = format(float(stats["limit_maxbytes"]))

    if used_memory > thresold["critical"]:
        return [2, "CRITICAL - %d%% used of %s RAM" % (used_memory, total_memory)]
    elif used_memory > thresold["warning"]:
        return [1, "WARNING - %d%% used of %s RAM" % (used_memory, total_memory)]
    else:
        return [0, "OK - %d%% used of %s RAM" % (used_memory, total_memory), "used_memory=%d%%" % used_memory]