#!/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
    TOKEN_CHANNEL: Token channel (True) / ID channel (False) routing
    TOKO: TOken channel Out
    TOKI: TOken channel In
    IDO: ID channel Out
    IDI: ID channel 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, TOKEN_CHANNEL=True, TOKO=1, TOKI=1, IDO=1, IDI=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 = TOKO
    if not TOKEN_CHANNEL:
        num_out_mult = IDO
    num_in_mult = TOKI
    if not TOKEN_CHANNEL:
        num_in_mult = IDI

    # Storage and Ingestion data stored per_client
    heap_size += clients_per_server * 100
    heap_size += (clients_per_server + M * M) * (B+8)

    # 2 Buffers of clients_per_server items of B size each, plus 1 of
    # size (clients_per_server + M * M) items, for token channel routing
    heap_size += (clients_per_server * B * 2) * num_out_mult
    heap_size += ((clients_per_server + M * M) * B) * num_in_mult

    # Additional buffers for ID channel routing

    # Round (M-1)^2 up to a multiple of M
    round1b_size = (M-1)*M
    wn_size = max(clients_per_server + M*M, 2*round1b_size)
    # Round up to a multiple of M
    colsort_size = int((wn_size + M - 1) / M) * M

    if not TOKEN_CHANNEL:
        heap_size += (colsort_size * B * 3) + (2 * round1b_size * B)

    # 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) * 9))

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

def generate_config(N, M, T, B, TOKEN_CHANNEL=True, TOKO=1, TOKI=1, IDO=1, IDI=1, num_WN_to_precompute=12):

    cf = open(CONFIG_FILE, 'w+')
    heap_size_page_aligned = get_heap_size(N, M, T, B, TOKEN_CHANNEL,
        TOKO, TOKI, IDO, IDI, 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)