#!/usr/bin/python3

import csv
import os

NODE_IPS=["192.168.0.1", "192.168.0.2", "192.168.0.3", "192.168.0.4"]
NUM_CPUS=6
RAFT_GROUP=3
TOTAL_SLOTS=16384

cluster_script = {}


def bucket(seq, num):
    avg = len(seq) / float(num)
    out = []
    last = 0.0

    while last < len(seq):
        out.append(seq[int(last):int(last + avg)])
        last += avg

    return out


try:
    os.mkdir("/tmp/raft")
except: 
    dir = 1

for ip in NODE_IPS:
    try:
        os.mkdir(os.path.join("/tmp/raft", ip))
    except: 
        dir = 1


with open('/tmp/raft/slots.csv', mode='w') as slots_file:
    slots_writer = csv.writer(
        slots_file,
        delimiter=',',
        quotechar='"',
        quoting=csv.QUOTE_MINIMAL)
    p = 1
    slot_row = []
    for x in bucket(range(TOTAL_SLOTS), (len(NODE_IPS)*NUM_CPUS)/RAFT_GROUP):
        r = "Slots: %s-%s Port: %d" %(x[0], x[-1], 19900 +p)
        for i in range(0,RAFT_GROUP):
            ip = NODE_IPS[len(slot_row)%len(NODE_IPS)]
            if 19900+p in cluster_script:
                cluster_script[19900+p].append(ip)
            else:
                cluster_script[19900+p] = [ip]
            if len(slot_row) == len(NODE_IPS):
                slots_writer.writerow(slot_row)
                slot_row = []
            slot_row.append(r)

        p += 1

    slots_writer.writerow(slot_row)

with open('/tmp/raft/setup_cluster.sh', mode='w') as outf:
    outf.write("#!/bin/bash -ex\n\n\n")
    for port in cluster_script:
        outf.write("\necho \"Setting up Raft Groups - port %d\"\n\n" %port)
        outf.write("/usr/local/bin/redis-cli -h %s -p %d raft.cluster init\n" %(cluster_script[port][0], port))
        for x in cluster_script[port][1:]:
            outf.write("/usr/local/bin/redis-cli -h %s -p %d raft.cluster join %s:%d\n" %(x, port, cluster_script[port][0], port))
    outf.write("\n\n\n# Sleep until raft cluster is created\necho \"Waiting for Raft to stabilize\"\nsleep 10\n\n# Linking Shards\n\n")
    l = list(cluster_script.keys())
    for p in l[1:-1]:
        outf.write("/usr/local/bin/redis-cli -h %s -p %d RAFT.SHARDGROUP LINK %s:%d\n" %(cluster_script[l[0]][0], l[0], cluster_script[p][0], p))