/*
 * SPDX-FileCopyrightText: 2025 orangerot <me@orangerot.dev>
 *
 * SPDX-License-Identifier: GPL-3.0
 *
 * This program, named WAI, is a WebAssembly Interpreter. 
 * Compile using make. 
 * Usage: wai [file.wasm] [function name] [function arguments ...]
 *
 * Copyright (C) 2025  orangrot <me@orangerot.dev>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

#include <endian.h>
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>

#define STACK_CAPACITY 1024
#define MAX_FUNCTIONS 128
#define MAX_FUNCTION_PARAMETERS 32
#define MAX_FUNCTION_RESULTS 32
#define MAX_EXPORT_NAME_LENGTH 128

enum section {
  Section_Custom,
  Section_Type,
  Section_Import,
  Section_Function,
  Section_Table,
  Section_Memory,
  Section_Global,
  Section_Export,
  Section_Start,
  Section_Element,
  Section_Code,
  Section_Data,
  Section_Data_Count,
};

enum TYPE {
  TYPE_ANY = 0,
  TYPE_I32 = 0x7F,
  TYPE_I64 = 0x7E,
  TYPE_F32 = 0x7D,
  TYPE_F64 = 0x7C,
  TYPE_V128 = 0x7B,
  TYPE_FUNCREF = 0x70,
  TYPE_EXTERNREF = 0x6F
};

static const int TYPE_SIZE[] = {
  [TYPE_I32] = 4,
  [TYPE_I64] = 8,
  [TYPE_F32] = 4,
  [TYPE_F64] = 8,
  [TYPE_V128] = 16,
  [TYPE_FUNCREF] = 16,
  [TYPE_EXTERNREF] = 16
};

static const char *TYPE_NAME[] = {
  [TYPE_I32] = "I32",
  [TYPE_I64] = "I64",
  [TYPE_F32] = "F32",
  [TYPE_F64] = "F64",
  [TYPE_V128] = "V128",
  [TYPE_FUNCREF] = "FREF",
  [TYPE_EXTERNREF] = "EXTR",
  [TYPE_ANY] = "ANY",
};

struct stack {
  u_char items[STACK_CAPACITY];
  size_t bytes;
};

struct func_type_t {
  enum TYPE param[MAX_FUNCTION_PARAMETERS];
  enum TYPE result[MAX_FUNCTION_RESULTS];
  size_t num_params;
  size_t num_results;
};

struct func_t {
  size_t func_type_index;
  size_t num_local_vars;
  u_char *addr;
};

enum export_desc {
  Export_Func,
  Export_Table,
  Export_Mem,
  Export_Global,
};

struct export_t {
  char name[MAX_EXPORT_NAME_LENGTH];
  size_t name_length;
  size_t func_index;
  enum export_desc description;
};

struct module {
  struct func_type_t func_types[MAX_FUNCTIONS];
  struct func_t func[MAX_FUNCTIONS];
  struct table_t *tables;
  struct mem_t *mems;
  struct global_t *globals;
  struct elem_t *elems;
  struct data_t *datas;
  struct start_t *start;
  struct import_t *imports;
  struct export_t exports[MAX_FUNCTIONS];
  u_char *binary;
  struct stack stack;
  size_t num_exports;
};

struct value_t {
  enum TYPE type;
  union {
    int32_t i32; 
    int64_t i64; 
    float f32; 
    double f64; 
    __int128 v128; 
    int64_t funcref; 
    int64_t extref; 
  } value;
};

struct context {
  struct module *module;
  size_t func_i;
  size_t func_stack_begin;
};

#define incr(i, len) i++; if (i >= len) {return 0;}

void print_value(struct value_t *value) {
  void *number = &value->value;
  switch (value->type) {
    case TYPE_I32:
      printf("%d", *(int32_t*)number);
      break;
    case TYPE_I64:
    case TYPE_ANY:
      printf("%ld", *(int64_t*)number);
      break;
    case TYPE_F32:
      printf("%f", *(float*)number);
      break;
    case TYPE_F64:
      printf("%f", *(double*)number);
      break;
    case TYPE_V128:
      printf("%f", (double) *(__int128*)number);
      break;
    case TYPE_FUNCREF:
      printf("%ld", *(int64_t*)number);
      break;
    case TYPE_EXTERNREF:
      printf("%ld", *(int64_t*)number);
      break;
  }
  printf(" (%s)", TYPE_NAME[value->type]);
}

void stack_peak(struct stack *s, struct value_t *value, size_t nth, size_t from) {
  int byte_i = from - 1; 
  for (size_t element_i = 0; element_i < nth && byte_i > 0; element_i++) {
    byte_i -= TYPE_SIZE[s->items[byte_i]];
  }
  value->type = s->items[byte_i];
  memcpy(&value->value, &(s->items[byte_i - TYPE_SIZE[value->type]]), TYPE_SIZE[value->type]);
} 

void stack_push(struct stack *s, const struct value_t *value) {
  size_t type_size = TYPE_SIZE[value->type];
  memcpy(&(s->items[s->bytes]), &value->value, type_size);
  s->items[s->bytes + type_size] = value->type;
  s->bytes += type_size + 1;

  printf("stack: ");
  for (int i = s->bytes - 1; i > 0; i -= TYPE_SIZE[s->items[i]] + 1) {
    struct value_t stack_value = {0};
    stack_peak(s, &stack_value, 0, i + 1);
    print_value(&stack_value);
    printf(", ");
  }
  printf("\n");
} 

void stack_top(struct stack *s, struct value_t *value) {
  value->type = s->items[s->bytes-1];
  memcpy(&value->value, &(s->items[s->bytes - 1 - TYPE_SIZE[value->type]]), TYPE_SIZE[value->type]);
} 

void stack_pop(struct stack *s, struct value_t *value) {
  stack_top(s, value);
  s->bytes -= TYPE_SIZE[value->type] + 1;
}

int parse_function(struct module *module, size_t func_i, size_t len);
int parse_instruction(struct context context, u_char *binary, size_t len);

#define PARAMS(...) __VA_ARGS__

// https://webassembly.github.io/spec/core/appendix/index-instructions.html
// OP(NAME, CODE, PARAM, NUM_RESULTS, LEN_IMMIDIATE, BODY)
#define DEFINE_OPERATIONS(OP)                                                                                                         \
OP(INSTR_F64_MUL,   0xa2, PARAMS(TYPE_F64, TYPE_F64), 1, 0, result->type = TYPE_F64; result->value.f64 = a->value.f64 * b->value.f64) \
OP(INSTR_F64_SUB,   0xa1, PARAMS(TYPE_F64, TYPE_F64), 1, 0, result->type = TYPE_F64; result->value.f64 = b->value.f64 - a->value.f64) \
OP(INSTR_F64_LT,    0x63, PARAMS(TYPE_F64, TYPE_F64), 1, 0, result->type = TYPE_F64; result->value.f64 = b->value.f64 < a->value.f64) \
OP(INSTR_F64_CONST, 0x44, PARAMS(), 1, 8, result->type = TYPE_F64; result->value.f64 = *(double*)immidiate)                           \
OP(INSTR_LOCAL_GET, 0x20, PARAMS(), 1, 1,                                                                                             \
  size_t func_type_i = context.module->func[context.func_i].func_type_index;                                                          \
  struct func_type_t *func_type = &context.module->func_types[func_type_i];                                                           \
  int num_locals = func_type->num_params + context.module->func[context.func_i].num_local_vars;                                       \
                                                                                                                                      \
  printf("num locals %d, %d\n", num_locals, num_locals - 1 - *(u_char*)immidiate);                                                    \
  stack_peak(&context.module->stack, result, num_locals - 1 - *(u_char*)immidiate, context.func_stack_begin);                         \
)                                                                                                                                     \
OP(INSTR_CALL, 0x10, PARAMS(), 0, 1, parse_function(context.module, *(u_char*)immidiate, len))                                        \
OP(INSTR_IF, 0x04, PARAMS(TYPE_ANY), 0, 1,                                                                                            \
  size_t i = 0;                                                                                                                       \
  enum TYPE condition_type = *(u_char*) immidiate;                                                                                    \
  if (a->type != condition_type)                                                                                                      \
    printf("Wrong types!\n");                                                                                                         \
                                                                                                                                      \
  while (binary[i] != INSTR_ELSE) {                                                                                                   \
    if (a->value.i64) {                                                                                                               \
      i += parse_instruction(context, &binary[i], len);                                                                               \
    } else {                                                                                                                          \
      incr(i, len);                                                                                                                   \
    }                                                                                                                                 \
  }                                                                                                                                   \
  incr(i, len);                                                                                                                       \
  while (binary[i] != INSTR_END) {                                                                                                    \
    if (a->value.i64) {                                                                                                               \
      incr(i, len);                                                                                                                   \
    } else {                                                                                                                          \
      i += parse_instruction(context, &binary[i], len);                                                                               \
    }                                                                                                                                 \
  }                                                                                                                                   \
  incr(i, len);                                                                                                                       \
  return i;                                                                                                                           \
)

enum OP_CODES {
#define AS_ENUM(NAME, CODE, PARAM, NUM_RESULTS, LEN_IMMIDIATE, BODY) NAME = CODE,
DEFINE_OPERATIONS(AS_ENUM)
  INSTR_END = 0x0B,
  INSTR_ELSE = 0x05
};  


#define AS_FUNCTION(NAME, CODE, PARAM, NUM_RESULTS, LEN_IMMIDIATE, BODY)       \
  int exec_##NAME(struct context context, struct value_t *a,                   \
                  struct value_t *b, void *immidiate, struct value_t *result,  \
                  u_char *binary, size_t len) {                                \
    (void) context;                                                            \
    (void) a;                                                                  \
    (void) b;                                                                  \
    (void) immidiate;                                                          \
    (void) result;                                                             \
    (void) binary;                                                             \
    (void) len;                                                                \
    BODY;                                                                      \
    return 0;                                                                  \
  }
DEFINE_OPERATIONS(AS_FUNCTION)

struct instruction {
  size_t num_param;
  enum TYPE params[2];
  size_t len_immidiate;
  size_t num_results;
  int (*exec) (struct context context, struct value_t *a, struct value_t *b, void *immidiate, struct value_t *result, u_char *binary, size_t len);
};

struct instruction INSTRUCTIONS[] = {
#define AS_INSTRUCTION(NAME, CODE, PARAM, NUM_RESULTS, LEN_IMMIDIATE, BODY)    \
  [NAME] = {                                                                   \
    .num_param = sizeof((enum TYPE[]) {PARAM}) / sizeof(enum TYPE),            \
    .params = {PARAM},                                                         \
    .num_results = NUM_RESULTS,                                                \
    .len_immidiate = LEN_IMMIDIATE,                                            \
    .exec = &exec_##NAME                                                       \
  },                                                                                                    
DEFINE_OPERATIONS(AS_INSTRUCTION)
};

int parse_instruction(struct context context, u_char *binary, size_t len) {
  size_t i = 0;	
  enum OP_CODES op_code = binary[i];
  u_char *instr_addr = &binary[i];
  struct value_t result = {0};
  struct value_t arguments[2];

  incr(i, len);

  struct instruction *instr = &INSTRUCTIONS[op_code];
  if (instr->exec == NULL) {
    printf("not implemented/illegal instruction %x at %lx\n", op_code, instr_addr - context.module->binary);
    exit(1);
  };

  for (size_t param_i = 0; param_i < instr->num_param; param_i++) {
    stack_pop(&context.module->stack, &arguments[param_i]);
    if (instr->params[param_i] != TYPE_ANY && arguments[param_i].type != instr->params[param_i]) {
      printf("wrong type! %x\n", op_code);
    }
  }
  i += instr->exec(context, &arguments[0], &arguments[1], &binary[i], &result, &binary[i + instr->len_immidiate], len);
  i += instr->len_immidiate;
  if (instr->num_results) {
    stack_push(&context.module->stack, &result);
  }

  return i;
}

int parse_function(struct module *module, size_t func_i, size_t len) {
  size_t i = 0;	
  struct func_t *func = &module->func[func_i];
  u_char *binary = func->addr;
  size_t func_type_i = func->func_type_index;
  struct func_type_t *func_type = &module->func_types[func_type_i];
  // int body_size = binary[i];
  size_t func_stack_begin = module->stack.bytes;
  size_t func_stack_end;
  struct value_t result = {0};
  struct context context = {
    .module = module,
    .func_i = func_i,
    .func_stack_begin = func_stack_begin
  };

  incr(i, len);
  func->num_local_vars = binary[i];
  incr(i, len);
  for (size_t local_var_i = 0; local_var_i < func->num_local_vars; local_var_i++) {
    stack_push(&module->stack, &(struct value_t) {.type = binary[i], .value = {0}}); 
    incr(i, len);
  }
  while (binary[i] != INSTR_END) {
    i += parse_instruction(context, &binary[i], len);
  }
  incr(i, len);
  
  func_stack_end = module->stack.bytes;
  module->stack.bytes = func_stack_begin;
  for (size_t local_i = 0; local_i < func_type->num_params + func->num_local_vars; local_i++) {
    stack_pop(&module->stack, &result);
  }
  for (size_t result_i = 0; result_i < func_type->num_results; result_i++) {
    stack_peak(&module->stack, &result, func_type->num_results - 1 - result_i, func_stack_end);
    stack_push(&module->stack, &result);
  }
  return i;
}

int parse_section(struct module *module, u_char *binary, size_t len) {
  size_t i = 0;	
  enum section type = binary[i];
  incr(i, len);
  int size = binary[i];
  incr(i, len);
  printf("section %x with size %d\n", type, size);

  switch ((enum section) type) {
    case Section_Custom:
      break;
    case Section_Type: 
        printf("section: type\n");
        struct func_type_t *func_type;
        size_t num_types = binary[i];
        incr(i, len);
        for (size_t type_i = 0; type_i < num_types; type_i++) {
          if (binary[i] != 0x60) {
            printf("expected function type, found %x\n", binary[i]);
            return 0;
          }
          incr(i, len);
          func_type = &module->func_types[type_i];
          func_type->num_params = binary[i];
          incr(i, len);
          for (size_t param_i = 0; param_i < func_type->num_params; param_i++) {
            func_type->param[param_i] = binary[i];
            incr(i, len);
          }
          func_type->num_results = binary[i];
          incr(i, len);
          for (size_t result_i = 0; result_i < func_type->num_results; result_i++) {
            func_type->result[result_i] = binary[i];
            incr(i, len);
          }
        }
        break;
    case Section_Import:
      break;
    case Section_Function:
      printf("section: function\n");
      size_t num_functions = binary[i];
      incr(i, len);
      for (size_t function_i = 0; function_i < num_functions; function_i++) {
        module->func[function_i].func_type_index = binary[i];
        incr(i, len);
      }
      break;
    case Section_Table:
      break;
    case Section_Memory:
      break;
    case Section_Global:
      break;
    case Section_Export:
      printf("section: exports\n");
      module->num_exports = binary[i];
      
      if(module->num_exports > MAX_FUNCTIONS) {
        printf("Number of exports exceeds maximum number of functions in a module (%d)", MAX_FUNCTIONS);
        return 0;
      }
      
      incr(i, len);
      for (size_t exports_i = 0; exports_i < module->num_exports; exports_i++) {
        struct export_t *export = &module->exports[exports_i];

        export->name_length = binary[i];
        incr(i, len);
        
        for (size_t si = 0; si < export->name_length; si++) {
          export->name[si] = (char) binary[i];
          incr(i, len);
        }
        export->description = (enum export_desc) binary[i];
        incr(i, len);
        export->func_index = (size_t) binary[i];
        printf("export name: %s of type %d\n", export->name, export->description);
        if (export->description == Export_Func) {
          printf("exported function %s(", export->name);
          size_t func_type_index = module->func[export->func_index].func_type_index;
          struct func_type_t *func_type = &module->func_types[func_type_index];
          for (size_t param_i = 0; param_i < func_type->num_params; param_i++) {
            printf("%s", TYPE_NAME[func_type->param[param_i]]);
            if (param_i == func_type->num_params -1) {
              printf(") -> ");
            } else {
              printf(",");
            }
          }
          for (size_t result_i = 0; result_i < func_type->num_results; result_i++) {
            printf("%s", TYPE_NAME[func_type->result[result_i]]);
            if (result_i != func_type->num_params -1) {
              printf(",");
            }
          }
          printf("\n");
        }
        incr(i, len);
      }
      break;
    case Section_Start:
      break;
    case Section_Element:
      break;
    case Section_Code:
      printf("section: code\n");
      int num_functions2 = binary[i];
      incr(i, len);
      for (int function_i = 0; function_i < num_functions2; function_i++) {
        module->func[function_i].addr = &binary[i];
        stack_push(&module->stack, &(struct value_t) {.type = TYPE_F64, .value.f64 = 1});
        i += parse_function(module, function_i, len);
        stack_pop(&module->stack, &(struct value_t) {0});
      }
      // printf("result: %f\n", module->stack.items[0]);
      break;
    case Section_Data:
      break;
    case Section_Data_Count:
      break;
    default:
      fprintf(stderr, "expectet section\n");
      exit(1);
  }

  if (size == 0x0) {incr(i, len);}
  return i;
}

int parse_module(struct module *module, u_char *binary, size_t len) {
  size_t i = 0;	
  char *magic = "\0asm";
  while (i < 4) {
    if ((char) binary[i] != magic[i]) {
      fprintf(stderr, "no wasm magic\n");
      return 0;
    }
    incr(i, len);
  }
  printf("magic found\n");
  printf("wasm version: %d\n", le32toh(*(int*)&binary[i]));
  i += 4;
  printf("addr %zu\n", i);

  module->binary = binary;

  while (i < len) {
    i += parse_section(module, &binary[i], len);
  }
  return i;
}

int main(int argc, char **argv) {
  FILE *file;
  u_char *binary;
  struct stat st;
  struct module module = {0};

  if (argc < 3) {
    printf("Usage: %s [file] [function name] [function arguments ...]\n", argv[0]);
    exit(1);
  };
  file = fopen(argv[1], "r");
  if (file == NULL) {
    fprintf(stderr, "Failed to open file\n");
    fclose(file);
    return 1;
  }
  stat(argv[1], &st);
  printf("size: %ld\n", st.st_size);
  
  binary = malloc(st.st_size);
  fread(binary, st.st_size, 1, file);
  fclose(file);

  if (parse_module(&module, binary, st.st_size) == -1) {
    printf("error :(\n");
  }
  printf("%zu\n", module.num_exports);
  printf("%s\n", module.exports[0].name);
  size_t export_search_i = 0;
  while (export_search_i < module.num_exports && 
      (strcmp(module.exports[export_search_i].name, argv[2]) != 0)) {
    export_search_i++;
  }
  if (export_search_i == module.num_exports) {
    printf("Provided function name %s not recognised. \n", argv[2]);
    exit(1);
  }
  size_t function_search_i = module.exports[export_search_i].func_index;
  size_t function_search_type_index = module.func[function_search_i].func_type_index;
  struct func_type_t *func_type_search = &module.func_types[function_search_type_index];
  if (func_type_search->num_params != (size_t) (argc - 3)) {
    printf("Not enough function arguments provided. Got %d expected %zu. \n", argc - 3, func_type_search->num_params);
    exit(1);
  }

  for (size_t param_i = 0; param_i < func_type_search->num_params; param_i++) {
    enum TYPE param_type = func_type_search->param[param_i];
    struct value_t param = {
      .type = param_type,
      .value = {0}
    };
    char *param_str = argv[param_i + 3];
    switch (param_type) {
    case TYPE_I32:
      param.value.i32 = atoi(param_str);
      break;
    case TYPE_I64:
      param.value.i64 = atoll(param_str);
      break;
    case TYPE_F32:
      param.value.f32 = strtof(param_str, NULL);
      break;
    case TYPE_F64:
      param.value.f64 = strtod(param_str, NULL);
      break;
    case TYPE_V128:
    case TYPE_FUNCREF:
    case TYPE_EXTERNREF:
    default:
      printf("%s, %s, %s unsupported\n", TYPE_NAME[TYPE_V128], TYPE_NAME[TYPE_FUNCREF], TYPE_NAME[TYPE_EXTERNREF]);
      exit(1);
    }
    stack_push(&module.stack, &param);
  }
  parse_function(&module, function_search_i, 100);

  free(binary);
  return 0;
}