#!/usr/bin/python3

import math
import sys

###############################################################################

# CONFIGS TO SET:

CONFIG_FILE = "Enclave/Enclave.config.xml"

###############################################################################

'''
  Generate a manifest file with:
    N: number of clients
    M: number of servers
    T: number of threads for each server
    B: msg_size
    PRIVATE_ROUTE: Private (True) / Public (False) route
    PRO: PRivate Out
    PRI: PRivate In
    PUO: PUblic Out
    PUI: PUblic In
    num_WN_to_precompute: The default num_WN_to_precompute is 12 in App/start.cpp
'''

def get_heap_size(N, M, T, B, PRIVATE_ROUTE=True, PRO=1, PRI=1, PUO=1, PUI=1, num_WN_to_precompute=12):
    clients_per_server = math.ceil(N/M)

    # Base heap of 2 MB per thread
    heap_size = 2000000 * T

    num_out_mult = PRO
    if not PRIVATE_ROUTE:
        num_out_mult = PUO
    num_in_mult = PRI
    if not PRIVATE_ROUTE:
        num_in_mult = PUI

    # Storage and Ingestion data stored per_client = 52 bytes
    heap_size += clients_per_server * (B + 60)

    # 2 Buffers of clients_per_server items of B size each, plus 1 of
    # size (clients_per_server + M) items, for private routing
    heap_size += (clients_per_server * B * 2) * num_out_mult
    heap_size += ((clients_per_server + M) * B) * num_in_mult
    # Additional buffers for public routing
    wn_size = max(clients_per_server, 2*(M-1)**2)
    # Round up to a multiple of M
    colsort_size = int((wn_size + M - 1) / M) * M

    if not PRIVATE_ROUTE:
        heap_size += (colsort_size * B * 3) + (2 * (M-1)**2 * B * 2)

    # num_WN_to_precompute times size of each WN
    heap_size += (num_WN_to_precompute * num_out_mult * \
        (wn_size * math.ceil(math.log(wn_size,2) + 1) * 8))

    heap_size_page_aligned = math.ceil(heap_size/4096) * 4096
    return heap_size_page_aligned

def generate_config(N, M, T, B, PRIVATE_ROUTE=True, PRO=1, PRI=1, PUO=1, PUI=1, num_WN_to_precompute=12):

    cf = open(CONFIG_FILE, 'w+')
    heap_size_page_aligned = get_heap_size(N, M, T, B, PRIVATE_ROUTE,
        PRO, PRI, PUO, PUI, num_WN_to_precompute)
    hex_heap_size = hex(heap_size_page_aligned)

    enclave_config = '''<!-- Please refer to User's Guide for the explanation of each field -->
<EnclaveConfiguration>
  <ProdID>0</ProdID>
  <ISVSVN>0</ISVSVN>
  <StackMaxSize>0x40000</StackMaxSize>
  <HeapMaxSize>{H}</HeapMaxSize>
  <TCSNum>32</TCSNum>
  <TCSPolicy>1</TCSPolicy>
  <DisableDebug>0</DisableDebug>
  <MiscSelect>0</MiscSelect>
  <MiscMask>0xFFFFFFFF</MiscMask>
</EnclaveConfiguration>
    '''.format(H = hex_heap_size)

    #print (enclave_config)
    cf.write(enclave_config)

    cf.close()

if __name__ == "__main__":

    if len(sys.argv) != 5:
        print("Incorrect usage!\n")
        print("./gen_enclave_config.py expects 4 parameters.")
        print("Usage: ./gen_enclave_config.py <N = number of clients> <M = number of servers> <T = number of threads> <B = message_size>")
        exit()

    n = int(sys.argv[1])
    m = int(sys.argv[2])
    t = int(sys.argv[3])
    b = int(sys.argv[4])
    generate_config(n, m, t, b)