#!/usr/bin/env python
"""
usage: onos-gen-partitions [-h] [-s PARTITION_SIZE] [-n NUM_PARTITIONS]
                           [filename] [node_ip [node_ip ...]]

Generate the partitions json file given a list of IPs or from the $OC*
environment variables.

positional arguments:
  filename              File to write output to. If none is provided, output
                        is written to stdout.
  node_ip               IP Address(es) of the node(s) in the cluster. If no
                        IPs are given, will use the $OC* environment
                        variables. NOTE: these arguemnts are only processed
                        after the filename argument.

optional arguments:
  -h, --help            show this help message and exit
  -s PARTITION_SIZE, --partition-size PARTITION_SIZE
                        Number of nodes per partition. Note that partition
                        sizes smaller than 3 are not fault tolerant. Defaults
                        to 3.
  -n NUM_PARTITIONS, --num-partitions NUM_PARTITIONS
                        Number of partitions. Defaults to the number of nodes
                        in the cluster.
"""

from os import environ
from collections import deque
import argparse
import re
import json
import hashlib

convert = lambda text: int(text) if text.isdigit() else text.lower()
alphanum_key = lambda key: [convert(c) for c in re.split('([0-9]+)', key)]

def get_OC_vars():
  vars = []
  for var in environ:
    if re.match(r"OC[0-9]+", var):
      vars.append(var)
  return sorted(vars, key=alphanum_key)

def get_nodes(ips=None, default_port=9876):
    node = lambda id, ip, port : { 'id': id, 'ip': ip, 'port': port }
    result = []
    if not ips:
        ips = [ environ[v] for v in get_OC_vars() ]
    for ip_string in ips:
        address_tuple = ip_string.split(":")
        if len(address_tuple) == 3:
            id=address_tuple[0]
            ip=address_tuple[1]
            port=int(address_tuple[2])
        else:
            id=ip_string
            ip=ip_string
            port=default_port
        result.append(node(id, ip, port))
    return result

def generate_partitions(nodes, k, n):
  l = deque(nodes)
  perms = []
  for i in range(1, n+1):
    part = {
             'id': i,
             'members': list(l)[:k]
           }
    perms.append(part)
    l.rotate(-1)
  return perms

if __name__ == '__main__':
  parser = argparse.ArgumentParser(
      description="Generate the partitions json file given a list of IPs or from the $OC* environment variables.")
  parser.add_argument(
      '-s', '--partition-size', type=int, default=3,
      help="Number of nodes per partition. Note that partition sizes smaller than 3 are not fault tolerant. Defaults to 3." )
  parser.add_argument(
      '-n', '--num-partitions', type=int,
      help="Number of partitions. Defaults to the number of nodes in the cluster." )
 # TODO: make filename and nodes independent. This will break backwards compatibility with existing usage.
  parser.add_argument(
     'filename', metavar='filename', type=str, nargs='?',
     help='File to write output to. If none is provided, output is written to stdout.')
  parser.add_argument(
      'nodes', metavar='node_ip', type=str, nargs='*',
      help='IP Address(es) of the node(s) in the cluster. If no IPs are given, ' +
           'will use the $OC* environment variables. NOTE: these arguemnts' +
           ' are only processed after the filename argument.')

  args = parser.parse_args()
  filename = args.filename
  partition_size = args.partition_size
  nodes = get_nodes(args.nodes)
  num_partitions = args.num_partitions
  if not num_partitions:
    num_partitions = len(nodes)

  partitions = generate_partitions([v.get('id') for v in nodes], partition_size, num_partitions)
  m = hashlib.sha256()
  for node in nodes:
    m.update(node['ip'])
  name = int(m.hexdigest()[:8], base=16) # 32-bit int based on SHA256 digest
  data = {
           'name': name,
           'nodes': nodes,
           'partitions': partitions
         }
  output = json.dumps(data, indent=4)

  if filename:
    with open(filename, 'w') as f:
      f.write(output)
  else:
    print output