/*
 * 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]
 *
 * 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 <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>
#include <inttypes.h>
#include <sys/types.h>

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,
};

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

struct module {
  struct type_t *types;
  u_char *funcs[128];
  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;
  u_char *binary;
  struct stack stack;
  int scope;
};

enum INSTRUCTION {
  INSTR_CALL = 0x10,
  INSTR_ELSE = 0x05,
  INSTR_END = 0x0b,
  INSTR_F64_CONST = 0x44,
  INSTR_F64_LT = 0x63,
  INSTR_F64_MUL = 0xa2,
  INSTR_F64_SUB = 0xa1,
  INSTR_IF = 0x04,
  INSTR_LOCAL_GET = 0x20,
};  

enum TYPE {
  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
};

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;
};

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

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) {
    enum TYPE t = s->items[i];
    size_t type_size = TYPE_SIZE[t];
    void *value = &s->items[i - type_size];

    switch (t) {
      case TYPE_I32:
        printf("%d (I32)", *(int32_t*)value);
        break;
      case TYPE_I64:
        printf("%ld (I32)", *(int64_t*)value);
        break;
      case TYPE_F32:
        printf("%f (F32)", *(float*)value);
        break;
      case TYPE_F64:
        printf("%f (F64)", *(double*)value);
        break;
      case TYPE_V128:
        printf("%ld (V128)", *(__int128*)value);
        break;
      case TYPE_FUNCREF:
        printf("%ld (EREF)", *(int64_t*)value);
        break;
      case TYPE_EXTERNREF:
        printf("%ld (EREF)", *(int64_t*)value);
        break;
    }
    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_type(u_char *binary, int len) {
  int i = 0;
  enum TYPE param = binary[i];
  printf("type %x\n", param);
  incr(i, len);
  switch (param) {
    case TYPE_I32:
    case TYPE_I64:
    case TYPE_F32:
    case TYPE_F64:
    case TYPE_V128:
    case TYPE_FUNCREF:
    case TYPE_EXTERNREF:
      break;
    default:
      return -1;
  }
  return i;
}

int parse_function(struct module *module, u_char *binary, double param, int len);
int parse_instruction(struct module *module, u_char *binary, double param, int len) {
  int i = 0;
  enum INSTRUCTION instr = (u_char) binary[i];
  u_char *instr_addr = &binary[i];
  struct value_t a = {0};
  struct value_t b = {0};
  struct value_t result = {0};

  incr(i, len);

  switch (instr) {
    case INSTR_CALL: {
      int func_index = binary[i];
      incr(i, len);
      stack_pop(&module->stack, &a);
      parse_function(module, module->funcs[func_index], a.value.f64, len);
      break;
      }
    case INSTR_ELSE:
      printf("reached else instruction: impossible!\n");
    case INSTR_END:
      break;
    case INSTR_F64_CONST:
      result.type = TYPE_F64;
      result.value.f64 = *(double*)&binary[i];
      i += 8;
      stack_push(&module->stack, &result);
      break;
    case INSTR_F64_LT: {
      stack_pop(&module->stack, &a);
      stack_pop(&module->stack, &b);
      if (a.type != TYPE_F64 || b.type != TYPE_F64) 
        printf("Wrong types!\n");
      result.type = TYPE_F64;
      result.value.f64 = b.value.f64 < a.value.f64;
      stack_push(&module->stack, &result);
      break;
       }
    case INSTR_F64_MUL: {
      stack_pop(&module->stack, &a);
      stack_pop(&module->stack, &b);
      if (a.type != TYPE_F64 || b.type != TYPE_F64) 
        printf("Wrong types!\n");
      result.type = TYPE_F64;
      result.value.f64 = a.value.f64 * b.value.f64;
      stack_push(&module->stack, &result);
      break;
      }
    case INSTR_F64_SUB: {
      stack_pop(&module->stack, &a);
      stack_pop(&module->stack, &b);
      if (a.type != TYPE_F64 || b.type != TYPE_F64) 
        printf("Wrong types!\n");
      result.type = TYPE_F64;
      result.value.f64 = b.value.f64 - a.value.f64;
      stack_push(&module->stack, &result);
      break;
      }
    case INSTR_IF: {
      stack_pop(&module->stack, &a);
      enum TYPE condition_type = binary[i];
      incr(i, len);
      if (a.type != condition_type) 
        printf("Wrong types!\n");

      while (binary[i] != INSTR_ELSE) {
        // TODO test condition with correct type. 
        // This might not matter since all types are false with 0x0
        if (a.value.i64) {
          i += parse_instruction(module, &binary[i], param, len);
        } else {
          incr(i, len);
        }
      } 
      incr(i, len);
      while (binary[i] != INSTR_END) {
        if (a.value.i64) {
          incr(i, len);
        } else {
          i += parse_instruction(module, &binary[i], param, len);
        }
      }
      incr(i, len);
      break;
      }
    case INSTR_LOCAL_GET: {
      int local_index = binary[i];
      incr(i, len);
      stack_push(&module->stack, &(struct value_t) {.value.f64 = param, .type = TYPE_F64});
      break;
      }
    default:
      printf("unknown instruction! %x at %lx\n", instr, instr_addr - module->binary);
      exit(1);
  }
  return i;
}

int parse_function(struct module *module, u_char *binary, double param, int len) {
  int i = 0;
  int body_size = binary[i];
  incr(i, len);
  // int local_decl_cound = binary[i];
  incr(i, len);
  module->scope = 1;
  while (binary[i] != INSTR_END) {
    i += parse_instruction(module, &binary[i], param, len);
  }
  incr(i, len);
  return i;
}

int parse_section(struct module *module, u_char *binary, int len) {
  int 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");
        int num_types = binary[i];
        incr(i, len);
        for (int type_i = 0; type_i < num_types; type_i++) {
          if (binary[i] != 0x60) {
            printf("expected function type, found %x\n", binary[i]);
            return -1;
          }
          incr(i, len);
          int num_params = binary[i];
          incr(i, len);
          for (int params_i = 0; params_i < num_params; params_i++) {
            i += (parse_type(&binary[i], len));
          }
          int num_results = binary[i];
          incr(i, len);
          for (int results_i = 0; results_i < num_results; results_i++) {
            i += (parse_type(&binary[i], len));
          }
        }
        break;
    case Section_Import:
      break;
    case Section_Function:
      printf("section: function\n");
      int num_functions = binary[i];
      incr(i, len);
      for (int function_i = 0; function_i < num_functions; function_i++) {
        incr(i, len);
      }
      break;
    case Section_Table:
      break;
    case Section_Memory:
      break;
    case Section_Global:
      break;
    case Section_Export:
      printf("section: exports\n");
      int num_exports = binary[i];
      incr(i, len);
      for (int exports_i = 0; exports_i < num_exports; exports_i++) {
        int string_len = binary[i];
        incr(i, len);
        printf("export name: ");
        for (int si = 0; si < string_len; si++) {
          putchar(binary[i]);
          incr(i, len);
        }
        putchar('\n');
        // export kind
        incr(i, len);
        // export func index
        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->funcs[function_i] = &binary[i];
        i += parse_function(module, &binary[i], 4, len);
      }
      // 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(u_char *binary, size_t len) {
  int 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 %d\n", i);

  struct module module = {0};
  module.binary = binary;

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


int main(int argc, char **argv) {
  
  if (argc != 2) {
    fprintf(stderr, "Usage: %s [file]\n", argv[0]);
    return 1;
  };
  FILE *file = fopen(argv[1], "r");
  if (file == NULL) {
    fprintf(stderr, "Failed to open file\n");
    fclose(file);
    return 1;
  }
  struct stat st;
  stat(argv[1], &st);
  printf("size: %ld\n", st.st_size);
  
  u_char *binary = malloc(st.st_size);
  fread(binary, st.st_size, st.st_size, file);

  if (parse_module(binary, st.st_size) == -1) {
    printf("error :(\n");
  }
  
  free(binary);
  fclose(file);
  return 0;
}