import os

def beginSection(file, sectionComment):
    file.writelines("\n")
    file.writelines(sectionComment)
    file.writelines("\n")

def make_executable(path):
    mode = os.stat(path).st_mode
    mode |= (mode & 0o444) >> 2    # copy R bits to X
    os.chmod(path, mode)

def writeServerScript(id, nodeFolder, hosts, numServers, numThreads, dataset, ruleSet, isWindows):
    scriptName = "run" + str(numThreads)

    shellFilePath = nodeFolder + "/test" + str(numThreads)
    shellFile = open(shellFilePath, "w")
    shellFile.write("./DMAT -shell . " + scriptName)
    shellFile.close()
    make_executable(shellFilePath)

    scriptFile = scriptName.lower() + ".rdfox" 
    file = open(nodeFolder + "/scripts/" + scriptFile, "w")
    cmd = "server " + str(id)

    file.writelines("# Generated by Testgen for DMAT v0.3\n")
     
    #set variables
    beginSection(file, "# Variables")
    
    file.writelines("set node-server.init-timeout 9000000000\n")
    file.writelines("set node-server.stats true\n")

    # Initialise a data store and server
    beginSection(file, "# Initialisation")
    
    file.writelines("init par-complex-nn\n")
    file.writelines(cmd + " init " + str(numThreads) + " " + str(numServers) + "\n")
    if isWindows:
        file.writelines(cmd + " add .\n")
    else:
        if numServers == 1 or id < numServers - 1:
            file.writelines(cmd + " add " + dataset + ".p" + str(id) + ".etl.gz\n")

    #set network locations
    beginSection(file, "# Network Locations")
    
    for x in range(0, numServers):
        if isWindows:
            port = hosts[x].split()[2]
        else:
            port = "12000"
        if x == id:
            file.writelines(cmd + ' set ' + str(x) + ' localhost ' + port + '\n')
        else:
            file.writelines(cmd + ' set ' + str(x) + ' ' + hosts[x].split()[1] + ' ' + port + '\n')

    #load occurrences
    beginSection(file, "# Occurrences")
    
    if id == numServers - 1:
        file.writelines(cmd + " occurrences all " + dataset + ".ocr.gz\n");
    else:
        file.writelines(cmd + " occurrences " + dataset + ".ocr.gz\n"); 

    #start
    beginSection(file, "# Start")
    file.writelines(cmd + " start\n")

    # coordinator
    
    if id == numServers - 1:
        beginSection(file, "# Tasks")
        file.writelines([cmd + " waitready\n",
                         cmd + " import\n",
                         cmd + " count tuples\n",
                         cmd + " rules " + ruleSet + ".dlog\n",
                         cmd + " mat\n",
                         cmd + " count tuples\n",
                         cmd + " count DER\n",
                         cmd + " count FCT\n",
                         cmd + " count PAR"])

    file.close()
    return scriptName


def writeScripts(testPrefix, dataset, ruleSet, hostsFile, numServers, numThreads, isWindows):
    hosts = []
    for l in open(hostsFile, "r"):
        hosts.append(l)

    transferScriptPath = "transfer.sh"
    transferDMATPath = "transferDMAT.sh"

    transferScript = open(transferScriptPath, "w")
    transferDMAT = open(transferDMATPath, "w")

    if numServers > 1:
        numServers = numServers + 1

    for x in range(0, numServers):
        folder = testPrefix + "-" + str(x)
        writeServerScript(x, folder, hosts, numServers, numThreads, dataset, ruleSet, isWindows)
        transferScript.write("scp -o StrictHostKeyChecking=no -i key.pem -r " + folder + " ec2-user@" + hosts[x].split()[1] + ":" + " & " )
        transferDMAT.write("scp -o StrictHostKeyChecking=no -i key.pem DMAT ec2-user@" + hosts[x].split()[1] + ":" + folder + "/ & " )

    transferScript.write("wait")
    transferDMAT.write("wait")
    transferScript.close()
    transferDMAT.close()

    make_executable(transferScriptPath)
    make_executable(transferDMATPath)

from sys import argv
if len(argv) <= 1:
    settingsFile = open("settings.txt", "r")
    prefix = settingsFile.readline().strip()
    dataset = settingsFile.readline().strip()
    ruleset = settingsFile.readline().strip()
    hostsFile = "hosts.txt"
    numServers = int(settingsFile.readline().strip())
    numThreads = int(settingsFile.readline().strip())

    writeScripts(prefix, dataset, ruleset, hostsFile, numServers, numThreads, False)
else:
    writeScripts(argv[1], argv[2], argv[3], argv[4], int(argv[5]), int(argv[6]), argv[7])
#writeScripts("lh2", "LUBM-100", "LUBM_LE", "hosts.txt", 2, 1, False)

