#!/usr/bin/env python

"""
Based on Matt Weber es2graphite
https://github.com/mattweber/es2graphite
"""

import re
import sys
import json
import time
import pickle
import struct
import socket
import thread
import urllib2
import argparse
from pprint import pprint
from datetime import datetime

NODES = {}
CLUSTER_NAME = ''
STATUS = {'red': 0, 'yellow': 1, 'green': 2}
SHARD_STATE = {'CREATED': 0, 'RECOVERING': 1, 'STARTED': 2, 'RELOCATED': 3, 'CLOSED': 4}
HOST_IDX = -1


def log(what, force=False):
    if args.verbose or force:
        pprint(what)


def get_es_host():
    global HOST_IDX
    HOST_IDX = (HOST_IDX + 1) % len(args.es) # round-robin
    return args.es[HOST_IDX]


def normalize(what):
    if not isinstance(what, (list, tuple)):
        return re.sub('\W+', '_',  what.strip().lower()).encode('utf-8')
    elif len(what) == 1:
        return normalize(what[0])
    else:
        return '%s.%s' % (normalize(what[0]), normalize(what[1:]))


def add_metric(metrics, prefix, stat, val, timestamp):
    if isinstance(val, bool):
        val = int(val)

    if prefix[-1] == 'translog' and stat == 'id':
        return
    elif isinstance(val, (int, long, float)) and stat != 'timestamp':
        metrics.append((normalize((prefix, stat)), (timestamp, val)))
    elif stat == 'status' and val in STATUS:
        metrics.append((normalize((prefix, stat)), (timestamp, STATUS[val])))
    elif stat == 'state' and val in SHARD_STATE:
        metrics.append((normalize((prefix, stat)), (timestamp, SHARD_STATE[val])))


def process_node_stats(prefix, stats):
    metrics = []
    global CLUSTER_NAME
    CLUSTER_NAME = stats['cluster_name']
    for node_id in stats['nodes']:
        node_stats = stats['nodes'][node_id]
        NODES[node_id] = node_stats['name']
        process_section(int(time.time()), metrics, (prefix, CLUSTER_NAME, NODES[node_id]), node_stats)
    return metrics


def process_cluster_health(prefix, health):
    metrics = []
    process_section(int(time.time()), metrics, (prefix, CLUSTER_NAME), health)
    return metrics


def process_indices_status(prefix, status):
    metrics = []
    process_section(int(time.time()), metrics, (prefix, CLUSTER_NAME, 'indices'), status['indices'])
    return metrics


def process_indices_stats(prefix, stats):
    metrics = []