#include "../../tools/assembler.h"
#include "../../tools/parser.h"
#include "../../vm/vm.h"
#include "devices.h"
#include <SDL2/SDL.h>

#include <stdint.h>
#include <stdlib.h>
#include <string.h>

#define MAX_SRC_SIZE 16384

static DeviceOps screen_ops = {.open = screen_open,
                               .read = screen_read,
                               .write = screen_write,
                               .close = screen_close,
                               .ioctl = screen_ioctl};

static DeviceOps mouse_ops = {.open = mouse_open,
                              .read = mouse_read,
                              .write = mouse_write,
                              .close = mouse_close,
                              .ioctl = nil};

static DeviceOps keyboard_ops = {.open = keyboard_open,
                                 .read = keyboard_read,
                                 .write = keyboard_write,
                                 .close = keyboard_close,
                                 .ioctl = nil};

static DeviceOps console_device_ops = {
    .open = console_open,
    .read = console_read,
    .write = console_write,
    .close = console_close,
    .ioctl = console_ioctl,
};

static ScreenDeviceData screen_data = {0};
static MouseDeviceData mouse_data = {0};
static KeyboardDeviceData keyboard_data = {0};

// Function to save VM state to ROM file
bool saveVM(const char *filename, VM *vm) {
  FILE *file = fopen(filename, "wb");
  if (!file) {
    perror("Failed to open file for writing");
    return false;
  }

  // Write VM state (registers and pointers)
  if (fwrite(&vm->pc, sizeof(u32), 1, file) != 1 ||
      fwrite(&vm->cp, sizeof(u32), 1, file) != 1 ||
      fwrite(&vm->fp, sizeof(u32), 1, file) != 1 ||
      fwrite(&vm->sp, sizeof(u32), 1, file) != 1 ||
      fwrite(&vm->rp, sizeof(u32), 1, file) != 1 ||
      fwrite(&vm->mp, sizeof(u32), 1, file) != 1 ||
      fwrite(&vm->dc, sizeof(u32), 1, file) != 1 ||
      fwrite(&vm->flag, sizeof(i32), 1, file) != 1) {
    fprintf(stderr, "Failed to write VM state\n");
    fclose(file);
    return false;
  }
  // Write code section
  if (fwrite(vm->code, 1, CODE_SIZE, file) != CODE_SIZE) {
    fprintf(stderr, "Failed to write code section\n");
    fclose(file);
    return false;
  }

  // Write memory section
  if (fwrite(vm->memory, 1, MEMORY_SIZE, file) != MEMORY_SIZE) {
    fprintf(stderr, "Failed to write memory section\n");
    fclose(file);
    return false;
  }

  fclose(file);
  return true;
}

// Function to load VM state from ROM file
bool loadVM(const char *filename, VM *vm) {
  FILE *file = fopen(filename, "rb");
  if (!file) {
    perror("Failed to open ROM file");
    return false;
  }

  // Read VM state (registers and pointers)
  if (fread(&vm->pc, sizeof(u32), 1, file) != 1 ||
      fread(&vm->cp, sizeof(u32), 1, file) != 1 ||
      fread(&vm->fp, sizeof(u32), 1, file) != 1 ||
      fread(&vm->sp, sizeof(u32), 1, file) != 1 ||
      fread(&vm->rp, sizeof(u32), 1, file) != 1 ||
      fread(&vm->mp, sizeof(u32), 1, file) != 1 ||
      fread(&vm->dc, sizeof(u32), 1, file) != 1 ||
      fread(&vm->flag, sizeof(i32), 1, file) != 1) {
    fprintf(stderr, "Failed to read VM state\n");
    fclose(file);
    return false;
  }

  // Read code section
  if (fread(vm->code, 1, CODE_SIZE, file) != CODE_SIZE) {
    fprintf(stderr, "Failed to read code section\n");
    fclose(file);
    return false;
  }

  // Read memory section
  if (fread(vm->memory, 1, MEMORY_SIZE, file) != MEMORY_SIZE) {
    fprintf(stderr, "Failed to read memory section\n");
    fclose(file);
    return false;
  }

  fclose(file);
  return true;
}

// Function to compile and optionally save
bool compileAndSave(const char *source_file, const char *output_file, VM *vm) {
  FILE *f = fopen(source_file, "rb");
  if (!f) {
    perror("fopen");
    return false;
  }

  static char source[MAX_SRC_SIZE + 1];

  fseek(f, 0, SEEK_END);
  long len = ftell(f);
  fseek(f, 0, SEEK_SET);
  if (len >= MAX_SRC_SIZE) {
    fprintf(stderr, "Source is larger than buffer\n");
    fclose(f);
    return false;
  }
  size_t read = fread(source, 1, len, f);
  source[read] = '\0';
  fclose(f);

  ExprNode *ast = expr_parse(source, strlen(source));
  if (!ast) {
    printf("Parse failed.\n");
    return false;
  } else {
#ifdef ASM_DEBUG
    expr_print(ast, 0);
#endif
    assemble(vm, ast);
    expr_free(ast);

    // If output file specified, save the VM
    if (output_file) {
      if (!saveVM(output_file, vm)) {
        printf("Failed to save VM to %s\n", output_file);
        return false;
      }
      printf("VM saved to %s\n", output_file);
    }
    return true;
  }
}

void repl(VM *vm) {
  USED(vm);

  char buffer[1024 * 10] = {0}; // Larger buffer for multi-line input
  char line[1024];

  for (;;) {
    // Count current parentheses balance
    i32 paren_balance = 0;
    for (i32 i = 0; buffer[i]; i++) {
      if (buffer[i] == '(')
        paren_balance++;
      else if (buffer[i] == ')')
        paren_balance--;
    }

    // Show appropriate prompt
    if (paren_balance > 0) {
      printf(".. "); // Continuation prompt when unbalanced
    } else {
      printf("> "); // Normal prompt when balanced
    }
    fflush(stdout);

    if (!fgets(line, sizeof(line), stdin)) {
      printf("\n");
      break;
    }

    // Append the new line to buffer
    strncat(buffer, line, sizeof(buffer) - strlen(buffer) - 1);

    // Recalculate balance after adding new line
    paren_balance = 0;
    for (i32 i = 0; buffer[i]; i++) {
      if (buffer[i] == '(')
        paren_balance++;
      else if (buffer[i] == ')')
        paren_balance--;
    }

    // Only parse when parentheses are balanced
    if (paren_balance == 0) {
      // Check if buffer has actual content (not just whitespace)
      i32 has_content = 0;
      for (i32 i = 0; buffer[i]; i++) {
        if (!isspace(buffer[i])) {
          has_content = 1;
          break;
        }
      }

      if (has_content) {
        ExprNode *ast = expr_parse(buffer, strlen(buffer));
        if (!ast) {
          printf("Parse failed.\n");
        } else {
          assemble(vm, ast);
          while (step_vm(vm)) {
          }
          expr_free(ast);
        }
      }

      // Reset buffer for next input
      buffer[0] = '\0';
    }
    // If unbalanced, continue reading more lines
  }
  exit(0);
}

/*
 * This needs to be done dynamically eventually
 */
void register_sdl_devices(VM *vm) {
  screen_data.width = 640;
  screen_data.height = 480;
  screen_data.size = 640 * 480;

  vm_register_device(vm, "/dev/screen/0", "screen", &screen_data, &screen_ops);

  mouse_data.x = 0;
  mouse_data.y = 0;
  mouse_data.btn1 = 0;
  mouse_data.btn2 = 0;
  mouse_data.btn3 = 0;
  mouse_data.btn4 = 0;
  mouse_data.size = 12;

  vm_register_device(vm, "/dev/mouse/0", "mouse", &mouse_data, &mouse_ops);

  keyboard_data.keys = SDL_GetKeyboardState(&keyboard_data.key_count);
  vm_register_device(vm, "/dev/keyboard/0", "keyboard", &keyboard_data,
                     &keyboard_ops);
}

#ifdef ASM_DEBUG
const char *opcode_to_string(Opcode op) {
  static const char *names[] = {[OP_HALT] = "halt",
                                [OP_JMP] = "jump",
                                [OP_JMPF] = "jump-if-flag",
                                [OP_CALL] = "call",
                                [OP_RETURN] = "return",
                                [OP_LOAD] = "load",
                                [OP_LOAD_REG] = "load-r",
                                [OP_LOAD_REG8] = "load-r8",
                                [OP_LOADI8] = "load-i8",
                                [OP_LOADU8] = "load-u8",
                                [OP_LOADI16] = "load-i16",
                                [OP_LOADU16] = "load-u16",
                                [OP_LOAD_IMM] = "load-immediate",
                                [OP_MALLOC] = "malloc",
                                [OP_STORE] = "store",
                                [OP_STORE8] = "store-8",
                                [OP_STORE16] = "store-16",
                                [OP_PUSH] = "push",
                                [OP_POP] = "pop",
                                [OP_REG_MOV] = "register-move",
                                [OP_SYSCALL] = "syscall",
                                [OP_SLL] = "bit-shift-left",
                                [OP_SRL] = "bit-shift-right",
                                [OP_SRE] = "bit-shift-right-extend",
                                [OP_BAND] = "bit-and",
                                [OP_BOR] = "bit-or",
                                [OP_BXOR] = "bit-xor",
                                [OP_ADD_INT] = "add-int",
                                [OP_SUB_INT] = "sub-int",
                                [OP_MUL_INT] = "mul-int",
                                [OP_DIV_INT] = "div-int",
                                [OP_ADD_UINT] = "add-nat",
                                [OP_SUB_UINT] = "sub-nat",
                                [OP_MUL_UINT] = "mul-nat",
                                [OP_DIV_UINT] = "div-nat",
                                [OP_ADD_REAL] = "add-real",
                                [OP_SUB_REAL] = "sub-real",
                                [OP_MUL_REAL] = "mul-real",
                                [OP_DIV_REAL] = "div-real",
                                [OP_INT_TO_REAL] = "int-to-real",
                                [OP_UINT_TO_REAL] = "nat-to-real",
                                [OP_REAL_TO_INT] = "real-to-int",
                                [OP_REAL_TO_UINT] = "real-to-nat",
                                [OP_JEQ_INT] = "jump-eq-int",
                                [OP_JGT_INT] = "jump-gt-int",
                                [OP_JLT_INT] = "jump-lt-int",
                                [OP_JLE_INT] = "jump-le-int",
                                [OP_JGE_INT] = "jump-ge-int",
                                [OP_JEQ_UINT] = "jump-eq-nat",
                                [OP_JGT_UINT] = "jump-gt-nat",
                                [OP_JLT_UINT] = "jump-lt-nat",
                                [OP_JLE_UINT] = "jump-le-nat",
                                [OP_JGE_UINT] = "jump-ge-nat",
                                [OP_JEQ_REAL] = "jump-eq-real",
                                [OP_JGE_REAL] = "jump-ge-real",
                                [OP_JGT_REAL] = "jump-gt-real",
                                [OP_JLT_REAL] = "jump-lt-real",
                                [OP_JLE_REAL] = "jump-le-real",
                                [OP_STRLEN] = "string-length",
                                [OP_STREQ] = "string-eq",
                                [OP_STRCAT] = "string-concat",
                                [OP_STR_GET_CHAR] = "string-get-char",
                                [OP_STR_FIND_CHAR] = "string-find-char",
                                [OP_STR_SLICE] = "string-slice",
                                [OP_INT_TO_STRING] = "int-to-string",
                                [OP_UINT_TO_STRING] = "nat-to-string",
                                [OP_REAL_TO_STRING] = "real-to-string",
                                [OP_STRING_TO_INT] = "string-to-int",
                                [OP_STRING_TO_UINT] = "string-to-nat",
                                [OP_STRING_TO_REAL] = "string-to-real"};

  if (op < 0 || op >= (int)(sizeof(names) / sizeof(names[0]))) {
    return "<invalid-opcode>";
  }

  const char *name = names[op];
  return name ? name : "<unknown-opcode>";
}
#endif

i32 main(i32 argc, char *argv[]) {
  bool gui_mode = false;
  bool dump_rom = false;
  char *input_file = nil;
  char *output_file = nil;
  bool is_rom = false;

  // Parse command line arguments
  for (i32 i = 1; i < argc; i++) {
    if (strcmp(argv[i], "-g") == 0 || strcmp(argv[i], "--gui") == 0) {
      gui_mode = true;
    } else if (strcmp(argv[i], "-o") == 0 ||
               strcmp(argv[i], "--dump-rom") == 0) {
      dump_rom = true;
    } else if (input_file == nil) {
      // This is the input file
      input_file = argv[i];

      // Check if it's a ROM file
      const char *ext = strrchr(argv[i], '.');
      if (ext && (strcmp(ext, ".rom") == 0)) {
        is_rom = true;
      }
    } else if (output_file == nil && dump_rom) {
      // This is the output file for -o flag
      output_file = argv[i];
    }
  }

  VM vm = {0};

  bool compilation_success = true;
  if (input_file) {
    if (is_rom) {
      // Load ROM file directly
      compilation_success = loadVM(input_file, &vm);
    } else {
      // Compile Lisp file
      if (dump_rom && output_file) {
        compilation_success = compileAndSave(input_file, output_file, &vm);
      } else {
        compilation_success = compileAndSave(input_file, nil, &vm);
      }
    }
  } else {
    // No input file - enter REPL mode
    repl(&vm);
    return 0;
  }

  if (dump_rom) {
    return (compilation_success) ? EXIT_SUCCESS : EXIT_FAILURE;
  }

  // If dump_rom flag was set without specifying output file, use default
  if (dump_rom && !is_rom && !output_file) {
    if (!saveVM("memory_dump.bin", &vm)) {
      printf("Failed to save VM to memory_dump.bin\n");
      return EXIT_FAILURE;
    }
    printf("VM saved to memory_dump.bin\n");
    return EXIT_SUCCESS;
  }

  bool running = true;

  vm_register_device(&vm, "/dev/term/0", "terminal", nil, &console_device_ops);
  if (gui_mode) {
    register_sdl_devices(&vm);

    while (running) {
#ifdef ASM_DEBUG
      printf("| %s %d\n", opcode_to_string(vm.code[vm.pc]),vm.pc);
#endif
      running = step_vm(&vm);
    }
  } else {
    while (running) {
#ifdef ASM_DEBUG
      printf("| %s %d\n", opcode_to_string(vm.code[vm.pc]),vm.pc);
#endif
      running = step_vm(&vm);
    }
  }

  return EXIT_SUCCESS;
}