#!/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 unitfile(slot_start, slot_end, port, ip):
    myfile = """
####################################################################
[Unit]
Description=Redis Server Port %d
After=network.target

[Service]
WorkingDirectory=/db/data
Type=simple
LimitNOFILE=65535
User=redislabs
ExecStart=/usr/bin/taskset -c 2-32 /usr/local/bin/redis-server --port %d --protected-mode no --dbfilename %d-redis.rdb --loadmodule /usr/local/share/redisraft/redisraft.so addr=%s:%d raft-log-fsync=no follower-proxy=no raftize-all-commands=yes cluster-start-hslot=%d cluster-end-hslot=%d raft-log-filename=%d-raftlog.db cluster-mode=yes
StandardOutput=file:/db/errorlog/redis-%d.log
StandardError=file:/db/errorlog/redis-%d-error.log


[Install]
WantedBy=multi-user.target
Alias=redis-%d.service
    """
    return (myfile %(port, port, port, ip, port, slot_start, slot_end, port, port, port, port))

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)]
            #####
            uf = "redis-%d.service" %(19900+p)
            with open(os.path.join('/tmp/raft/', ip, uf), mode='w') as of:
                of.write(unitfile(x[0], x[-1], 19900+p, ip))
            
            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))