"""
Add devices

Add formatted drives to a swift ring.
--master-node: host that master api runs on.
--ring-type: ring to add devices to
--drive-list: file with newline separated devices
hosts: hosts of devices to add

Will automatically add all devices in the drive list file on all of the listed hosts
to the designated swift ring.  All requests get routed to the master swift node.
"""

import requests
import json
import sys
import argparse
import socket

DEFAULT_MASTER_NODE = 'swift-proxy-001.sfo01.justin.tv'
HEADERS = {'Content-type': 'application/json', 'Accept': 'text/plain'}

def get_sizes(host):
    """
    Return a dictionary of device names to sizes of those devices, for a given node.
    """

    resp = requests.get('http://{0}:8082/disks'.format(host))
    hw = resp.json()
    devices = hw.get('devices', {}).get('node', [])
    if type(devices) != list:
        devices = [devices]

    sizes = {}

    for device in devices:
        name = device.get('logicalname', device.get('@handle'))
        if type(name) == list:
            name = '/'.join(name)

        if 'node' in device:
            nodes = device['node']
            if type(nodes) is dict:
                nodes = [nodes]
            for node in nodes:
                subname = node.get('logicalname')
                if type(subname) == list:
                    name = '/'.join(subname)

                size = node.get('size')
                if size:
                    sizes[subname] = int(size[:-2])

    return sizes

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Add devices to a ring')
    parser.add_argument('--master-node', type=str, default=DEFAULT_MASTER_NODE, help='Node on which master api runs')
    parser.add_argument('--ring-type', type=str, required=True, help='Ring to add to: object, account, or container')
    parser.add_argument('--drive-list', type=str, required=True, help='Drive file with newline separated list of drives')
    parser.add_argument('--instant', action='store_true', help='When set, instantly insert devices instead of gradually')
    parser.add_argument('hosts', type=str, nargs='+', help='List of hosts to add drives to')
    args = parser.parse_args()

    headers = {'Content-type': 'application/json', 'Accept': 'text/plain'}

    print 'Starting add devices onto master node {0}...'.format(args.master_node)

    if args.ring_type == 'container':
        port = 6001
    elif args.ring_type == 'account':
        port = 6002
    elif args.ring_type == 'object':
        port = 6000
    else:
        print 'ring not valid'
        sys.exit(1)

    for host in args.hosts:
        host = host + '.justin.tv'
        devices = requests.get('http://' + host + ':8082/uuids').json()
        zone = requests.get('http://' + host + ':8082/zone').json()['zone']
        sizes = get_sizes(host)
        print sizes

        with open(args.drive_list) as f:
            for line in f:
                raw_device = '/dev/' + line[:-1]
                device = '/dev/' + line[:-1] + '1'

                print host
                print zone
                print device
                print '====='

                if device in devices:
                    uid = devices[device]['uuid']
                    host_ip = socket.gethostbyname(host)
                    if host_ip == '127.0.0.1':
                        print 'WARNING: resolved to loopback device, skipping!'
                        continue

                    print 'adding ' + uid
                    url = 'http://{0}:8077/swift/'.format(args.master_node) + args.ring_type + '?wait=true'
                    if args.instant:
                        url += '&instant=true'
                    requests.post(url, data=json.dumps({
                        'weight': sizes[raw_device],
                        'ip': host_ip,
                        'region': 1,
                        'zone': zone,
                        'port': port,
                        'device': uid,
                        'meta': host,
                    }), headers=headers)
