#include "../lib.h"
#include "../wasm_export.h"
#include "bh_platform.h"

#include "__WASM_ARRAY_FILE__"
#ifdef TARGET_LINUX
#include <stdio.h>
#include <string.h>
#endif

// FAIL* instrumentation symbols
void fail_start_trace(void) {}
void fail_stop_trace(void) {}
void fail_marker_positive(void) {}
void fail_marker_negative(void) {}
void fail_marker_detected(void) {}

// Those functions will be called from WASM
void host_fail_start_trace(wasm_exec_env_t exec_env) { fail_start_trace(); }
void host_fail_stop_trace(wasm_exec_env_t exec_env) { fail_stop_trace(); }
void host_fail_marker_positive(wasm_exec_env_t exec_env) {
  fail_marker_positive();
}
void host_fail_marker_negative(wasm_exec_env_t exec_env) {
  fail_marker_negative();
}
void host_fail_marker_detected(wasm_exec_env_t exec_env) {
  fail_marker_detected();
}
void host_print(wasm_exec_env_t exec_env, const char *msg) {
  PRINT("[WASM]    %s", msg);
}

#define STACK_SIZE (4 * 1024)
#define HEAP_SIZE STACK_SIZE

#if WASM_MEM_ALLOC_WITH_USAGE

// Another bump allocator
#define WASM_LINEAR_MEMORY_SIZE (64 * 1024) // Need to match --initial-memory?
#define RUNTIME_POOL_SIZE       (2 * 1024 * 1024)
#define LINEAR_POOL_SIZE        (WASM_LINEAR_MEMORY_SIZE + 512 * 1024)
#define ALIGNMENT 8

typedef struct {
  char  *buf;
  size_t size;
  size_t offset;
} BumpPool;

#if WAMR_RUNTIME_POOL_IN_TEXT
__attribute__((section(".text.wamr_runtime_pool"), aligned(4096)))
#endif
static char runtime_pool_buf[RUNTIME_POOL_SIZE];

#if WAMR_LINEAR_POOL_IN_TEXT
__attribute__((section(".text.wamr_linear_pool"), aligned(4096)))
#endif
static char linear_pool_buf[LINEAR_POOL_SIZE];

// mem_alloc_usage_t: 0 = Alloc_For_Runtime, 1 = Alloc_For_LinearMemory
static BumpPool pools[] = {
  { runtime_pool_buf, RUNTIME_POOL_SIZE, 0 },
  { linear_pool_buf,  LINEAR_POOL_SIZE,  0 },
};

static size_t align_up(size_t x, size_t a) {
  return (x + a - 1) & ~(a - 1);
}

static size_t alloc_size(void *ptr) {
  size_t header_size = align_up(sizeof(size_t), ALIGNMENT);
  return *(size_t *)((char *)ptr - header_size);
}

static void *bump_alloc(BumpPool *pool, unsigned int size) {
  size_t header_size = align_up(sizeof(size_t), ALIGNMENT);
  size_t start = align_up(pool->offset, ALIGNMENT);
  size_t end = start + header_size + align_up(size, ALIGNMENT);

  if (end > pool->size) {
    return NULL;
  }

  *(size_t *)&pool->buf[start] = size;
  void *ptr = &pool->buf[start + header_size];
  pool->offset = end;

  return ptr;
}

void *wamr_malloc(mem_alloc_usage_t usage, unsigned int size) {
  PRINT("wamr_malloc: usage=%d size=%u\n", usage, size);
  return bump_alloc(&pools[usage], size);
}

void *wamr_realloc(mem_alloc_usage_t usage, bool full_size_mmaped, void *ptr,
                   unsigned int new_size) {
  PRINT("wamr_realloc: usage=%d full_size_mmaped=%d ptr=%p new_size=%u\n",
        usage, full_size_mmaped, ptr, new_size);

  void *new_addr = bump_alloc(&pools[usage], new_size);
  if (!new_addr) {
    return NULL;
  }

  if (ptr) {
    size_t old_size = alloc_size(ptr);
    memcpy(new_addr, ptr, old_size < new_size ? old_size : new_size);
  }

  return new_addr;
}

void wamr_free(mem_alloc_usage_t usage, void *ptr) {
  // PRINT("wamr_free: usage=%d ptr=%p\n", usage, ptr);
}

#else

#define RUNTIME_POOL_SIZE (4 * STACK_SIZE)

#if WAMR_GLOBAL_HEAP_IN_TEXT
__attribute__((section(".text.wamr_global_heap"), aligned(4096)))
#endif
static char global_heap_buf[RUNTIME_POOL_SIZE];

#endif // WASM_MEM_ALLOC_WITH_USAGE

MAIN {
  char error_buf[128];

  // Step 1: Initialize WAMR Runtime
  static RuntimeInitArgs init_args;
  memset(&init_args, 0, sizeof(RuntimeInitArgs));

#if WASM_MEM_ALLOC_WITH_USAGE
  init_args.mem_alloc_type = Alloc_With_Allocator;
  init_args.mem_alloc_option.allocator.malloc_func = (void *)wamr_malloc;
  init_args.mem_alloc_option.allocator.realloc_func = (void *)wamr_realloc;
  init_args.mem_alloc_option.allocator.free_func = (void *)wamr_free;
#else
  init_args.mem_alloc_type = Alloc_With_Pool;
  init_args.mem_alloc_option.pool.heap_buf = global_heap_buf;
  init_args.mem_alloc_option.pool.heap_size = sizeof(global_heap_buf);
#endif

  init_args.max_thread_num = 1;
  if (!wasm_runtime_full_init(&init_args)) {
    PRINT_ERROR("wasm_runtime_full_init failed.\n");
    goto error_cleanup;
  }

  // Step 2: Export Native Symbols
  static NativeSymbol native_symbols[] = {
      {"fail_start_trace", (void *)host_fail_start_trace, "()", NULL},
      {"fail_stop_trace", (void *)host_fail_stop_trace, "()", NULL},
      {"fail_marker_positive", (void *)host_fail_marker_positive, "()", NULL},
      {"fail_marker_negative", (void *)host_fail_marker_negative, "()", NULL},
      {"fail_marker_detected", (void *)host_fail_marker_detected, "()", NULL},
      {"print", (void *)host_print, "(*)", NULL},
  };
  int count = sizeof(native_symbols) / sizeof(NativeSymbol);
  if (!wasm_runtime_register_natives("env", native_symbols, count)) {
    PRINT_ERROR("wasm_runtime_register_natives failed.\n");
    goto error_cleanup;
  }

  // Step 3: Parse and Validate Module
  wasm_module_t module = wasm_runtime_load(__WASM_ARRAY__, __WASM_ARRAY_LEN__,
                                           error_buf, sizeof(error_buf));
  if (!module) {
    PRINT_ERROR("wasm_runtime_load failed with \"%s\".\n", error_buf);
    goto error_cleanup;
  }

  // Step 4: Instantiate Module
  wasm_module_inst_t module_inst = wasm_runtime_instantiate(
      module, STACK_SIZE, HEAP_SIZE, error_buf, sizeof(error_buf));
  if (!module_inst) {
    PRINT_ERROR("wasm_runtime_instantiate failed with \"%s\".\n", error_buf);
    goto error_cleanup;
  }

  // Step 5: Create Execution Environment
  wasm_exec_env_t exec_env =
      wasm_runtime_create_exec_env(module_inst, STACK_SIZE);
  if (!exec_env) {
    PRINT_ERROR("wasm_runtime_create_exec_env failed.\n");
    goto error_cleanup;
  }

  // Step 6: Find and Call Exported Function
  wasm_function_inst_t func =
      wasm_runtime_lookup_function(module_inst, "wasm_module");
  if (!func) {
    PRINT_ERROR("wasm_runtime_lookup_function failed.\n");
    goto error_cleanup;
  }

  // In case wasm_module accepts arguments, set them here
  uint32_t args[1];

  if (!wasm_runtime_call_wasm(exec_env, func, 0, args)) {
    const char *exception = wasm_runtime_get_exception(module_inst);
    PRINT_ERROR("wasm_runtime_call_wasm failed with \"%s\".\n",
                exception ? exception : "unknown");
    goto error_cleanup;
  }

  PRINT_SUCCESS("wasm function execution finished.\n");

  // In case wasm_module returns a value we can do sth with it
  uint32_t retval = args[0];

success_cleanup:
  wasm_runtime_destroy_exec_env(exec_env);
  wasm_runtime_deinstantiate(module_inst);
  wasm_runtime_unload(module);
  wasm_runtime_destroy();

  RET(0);

error_cleanup:
  if (exec_env) {
    wasm_runtime_destroy_exec_env(exec_env);
  }
  if (module_inst) {
    wasm_runtime_deinstantiate(module_inst);
  }
  if (module) {
    wasm_runtime_unload(module);
  }
  wasm_runtime_destroy();

  RET(1);
}