#include "wasm_runner.h"
#include "FreeRTOS.h"
#include "nova64_api.h"
#include "nova64_fs.h"
#include "nova64_internal.h"
#include "task.h"
#include "wasm_export.h"
#include <stdlib.h>

#define MAX_WASM_TASKS 4

typedef enum { SLOT_FREE = 0, SLOT_RUNNING, SLOT_STOPPING } slot_state_t;

typedef struct {
  int id;
  slot_state_t state;
  // TaskHandle_t freertos_handle;
  wasm_module_inst_t module_inst;
  wasm_exec_env_t exec_env;
} wasm_task_slot_t;

static wasm_task_slot_t g_wasm_slots[MAX_WASM_TASKS];

char error_buf[128];
uint32_t stack_size = 8092, heap_size = 8092;

bool wasm_runner_initialize(void) {

  for (int i = 0; i < MAX_WASM_TASKS; i++) {
    g_wasm_slots[i].id = i;
    g_wasm_slots[i].state = SLOT_FREE;
    // g_wasm_slots[i].freertos_handle = NULL;
    g_wasm_slots[i].module_inst = NULL;
  }

  wasm_runtime_init();
  nova64_register_imports();
  nova64_log("WASM runtime initialized successfully.\n");

  return true;
}

bool wasm_runner_destroy(void) {
  wasm_runtime_destroy();
  return true;
}

bool wasm_runner_start_program(
    const char *module_path,
    const wasm_runner_framebuffer_config_t *framebuffer_config) {

  nova64_log("Running WASM module: %s\n", module_path);

  wasm_task_slot_t *slot = NULL;
  for (int i = 0; i < MAX_WASM_TASKS; i++) {
    if (g_wasm_slots[i].state == SLOT_FREE) {
      slot = &g_wasm_slots[i];
      break;
    }
  }

  if (!slot) {
    nova64_log("Maximum number of WASM tasks (%d) reached!\n", MAX_WASM_TASKS);
    return false;
  }

  if (!wasm_runtime_init_thread_env()) {
    nova64_log("Failed to initialize WASM runtime thread environment.\n");
    return false;
  }

  int32_t file_size = nova64_get_file_size(module_path);
  if (file_size <= 0) {
    nova64_log("Failed to determine WASM module size.\n");
    return false;
  }

  unsigned char *buffer = malloc((size_t)file_size);
  if (!buffer) {
    nova64_log("Failed to allocate WASM module buffer.\n");
    return false;
  }

  if (!nova64_read_file_to_buffer(module_path, buffer, (size_t)file_size)) {
    nova64_log("Failed to read WASM module into buffer.\n");
    free(buffer);
    return false;
  }

  wasm_module_t module = wasm_runtime_load(buffer, (uint32_t)file_size,
                                           error_buf, sizeof(error_buf));

  slot->module_inst = wasm_runtime_instantiate(module, stack_size, heap_size,
                                               error_buf, sizeof(error_buf));

  wasm_function_inst_t start_func =
      wasm_runtime_lookup_function(slot->module_inst, "_start");

  if (start_func == NULL) {
    nova64_log("Failed to find '_start' function in WASM module: %s\n",
               wasm_runtime_get_exception(slot->module_inst));
    free(buffer);
    return false;
  }

  slot->exec_env = wasm_runtime_create_exec_env(slot->module_inst, stack_size);

  nova64_log("WASM module execution environment created successfully.\n");

  uint32_t argv[2] = {5, 3}; // Example arguments for the "add" function

  slot->state = SLOT_RUNNING;

  if (wasm_runtime_call_wasm(slot->exec_env, start_func, 0, NULL)) {
    nova64_log("WASM module executed successfully.\n");
  } else {
    nova64_log("Failed to execute WASM module: %s\n",
               wasm_runtime_get_exception(slot->module_inst));
  }

  slot->state = SLOT_STOPPING;

  wasm_runtime_destroy_exec_env(slot->exec_env);
  wasm_runtime_deinstantiate(slot->module_inst);
  wasm_runtime_unload(module);
  free(buffer);
  wasm_runtime_destroy_thread_env();

  slot->module_inst = NULL;
  slot->exec_env = NULL;
  slot->state = SLOT_FREE;
  vTaskDelete(NULL);

  return true;
}

void wasm_runner_terminate(void) {
  nova64_log("Stopping all WASM tasks...\n");

  for (int i = 0; i < MAX_WASM_TASKS; i++) {
    if (g_wasm_slots[i].state == SLOT_RUNNING && g_wasm_slots[i].module_inst) {
      g_wasm_slots[i].state = SLOT_STOPPING;
      wasm_runtime_terminate(g_wasm_slots[i].module_inst);
    }
  }

  bool any_running = true;
  while (any_running) {
    any_running = false;
    for (int i = 0; i < MAX_WASM_TASKS; i++) {
      if (g_wasm_slots[i].state != SLOT_FREE) {
        any_running = true;
        break;
      }
    }
    if (any_running) {
      vTaskDelay(pdMS_TO_TICKS(10));
    }
  }

  nova64_log("All WASM tasks successfully stopped.\n");
}