#!/usr/bin/env python

# Copyright (c) 2014 Anchor Systems

# https://github.com/anchor/elasticsearch-scripts

"""
NAME

    es-index-locations - print the current state of /_status


SYNOPSIS

    es-index-locations [HOST[:PORT]]


DESCRIPTION

    This command will always output JSON to standard output.

"""
DEFAULT_HOST = "localhost"
DEFAULT_PORT = 9200

from esadmin import Connection
import json
import logging
import os
import sys

logger = logging.getLogger(__name__)

def usage(argv):
    print >>sys.stderr, ("Usage:\n\t%s <index> [HOST[:PORT]]\n" %
                         os.path.basename(argv[0]))

def main(argv=None):
    if argv is None:
        argv = sys.argv

    host = DEFAULT_HOST
    port = DEFAULT_PORT
    index = None
    try:
        index = sys.argv[1]
    except IndexError:
        usage(sys.argv)
        sys.exit(1)
    try:
        address = sys.argv[2]
    except IndexError:
        pass
    else:
        if address.find(":") == -1:
            host = address
        else:
            host, port = address.split(":")
            try:
                port = int(port)
            except ValueError:
                usage(argv)
                return 2

    shard_locations = []
    with Connection(host, port) as conn:
        resp = conn.get('/_cluster/health')
        cluster_name = resp["cluster_name"]
        resp = conn.get('/_status')
        for shard in resp["indices"][index]["shards"]:
            for replica in resp["indices"][index]["shards"][shard]:
                location = { 
                    "node_id" : replica["routing"]["node"],
                    "shard"   : replica["routing"]["shard"],
                    "primary" : replica["routing"]["primary"],
                }

                node_info = conn.get('/_cluster/nodes/%s' % location["node_id"])
                location["node_name"] = node_info["nodes"][location["node_id"]]["name"]

                location["path"] = "/var/lib/elasticsearch/%s/nodes/0/indices/%s/%s" % (
                                    cluster_name,
                                    index,
                                    location["shard"])
                    
                shard_locations.append(location)
   
    for location in shard_locations:
        print("%s:%s (shard %d, %s)" % (
            location["node_name"],
            location["path"], 
            location["shard"], 
            ("primary" if location["primary"] 
                else "replica")
            )
        )
 
    return 0

if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO,
                        format='%(levelname)s:\t%(message)s')

    sys.exit(main(sys.argv))
