#!/usr/bin/env python

import sys
from random import choice
from argparse import ArgumentParser
from pprint import pprint
from elasticsearch import Elasticsearch


def auto_relocate(host='localhost'):
    es = Elasticsearch(host)

    data = es.cluster.state('routing_table,nodes')
    commands = []

    for index in data['routing_table']['indices'].values():
        shards = [s for i in index['shards'].values() for s in i]

        unassigned = [shard for shard in shards
                      if shard['state'] == 'UNASSIGNED' and
                      not shard['relocating_node']]

        if len(unassigned) == 1:
            unassigned = unassigned[0]
            dest_node = [shard['node'] for shard in shards
                         if shard['shard'] == unassigned['shard'] and
                         shard['state'] != 'UNASSIGNED'][0]
            to_move = choice([shard for shard in shards
                              if shard['shard'] != unassigned['shard'] and
                              shard['primary'] != unassigned['primary']])
            print('')
            print('Move the below to: {0}'.format(
                data['nodes'][dest_node]['name']
            ))
            pprint(to_move)

            commands.append({'move': {
                'index': to_move['index'],
                'shard': to_move['shard'],
                'from_node': data['nodes'][to_move['node']]['name'],
                'to_node': data['nodes'][dest_node]['name']
            }})

    if raw_input('Perform the above actions? [y/N]: ').lower() == 'y':
        es.cluster.reroute({'commands': commands})


def parse_arguments():
    parser = ArgumentParser()
    parser.add_argument('--elasticsearch', '-e', type=str, default='localhost')
    return parser.parse_args()


def main():
    args = parse_arguments()
    auto_relocate(host=args.elasticsearch)
    sys.exit(0)


if __name__ == '__main__':
    main()