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

#include <stdint.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 = nil};

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

void compileFile(const char *path, VM *vm) {
  FILE *f = fopen(path, "rb");
  if (!f) {
    perror("fopen");
    exit(1);
  }

  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) {
    perror("source is larget than buffer");
    exit(1);
  }
  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");
  } else {
    assemble(vm, ast);
    expr_free(ast);
  }
}

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
    int paren_balance = 0;
    for (int 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 (int 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)
      int has_content = 0;
      for (int 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);
}

enum FlagType {
  FLAG_NONE = 0,
  FLAG_DUMP_ROM = 1,
  FLAG_GUI_MODE = 2,
};

#define MAX_INPUT_FILES 16 /* Adjust based on your system's constraints */

struct CompilerConfig {
  u32 flags;
  char *input_files[MAX_INPUT_FILES];
  i32 input_file_count;
};

i32 parse_arguments(i32 argc, char *argv[], struct CompilerConfig *config) {
  i32 i;

  /* Initialize config */
  config->flags = 0;
  config->input_file_count = 0;

  /* Zero out input files array for safety */
  for (i = 0; i < MAX_INPUT_FILES; i++) {
    config->input_files[i] = NULL;
  }

  for (i = 1; i < argc; i++) {
    if (argv[i][0] == '-') {
      /* Long and short flag handling */
      if (strcmp(argv[i], "-g") == 0 || strcmp(argv[i], "--gui") == 0) {
        config->flags |= FLAG_GUI_MODE;
      } else if (strcmp(argv[i], "-o") == 0 ||
                 strcmp(argv[i], "--dump-rom") == 0) {
        config->flags |= FLAG_DUMP_ROM;
      } else {
        fprintf(stderr, "Unknown flag: %s\n", argv[i]);
        return -1;
      }
    } else if (strstr(argv[i], ".ul") != NULL) {
      /* Collect input files */
      if (config->input_file_count >= MAX_INPUT_FILES) {
        fprintf(stderr, "Too many input files. Maximum is %d\n",
                MAX_INPUT_FILES);
        return -1;
      }
      config->input_files[config->input_file_count++] = argv[i];
    } else if (strstr(argv[i], ".lisp") != NULL) {
      /* Collect input files */
      if (config->input_file_count >= MAX_INPUT_FILES) {
        fprintf(stderr, "Too many input files. Maximum is %d\n",
                MAX_INPUT_FILES);
        return -1;
      }
      config->input_files[config->input_file_count++] = argv[i];
    }
  }

  return 0;
}

/*
 * This needs to be done dynamically eventually
 */
void register_sdl_devices(VM *vm) {
  screen_data.vm = vm;
  screen_data.width = 640;
  screen_data.height = 480;
  screen_data.framebuffer_size = 640 * 480;
  screen_data.framebuffer_pos = vm->mp;
  vm->mp += screen_data.framebuffer_size / 4; /* advance memory pointer */

  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;

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

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_IMM]     = "load-immediate",
        [OP_STORE]        = "store",
        [OP_PUSH]         = "push",
        [OP_POP]          = "pop",
        [OP_REG_MOV]      = "register-move",
        [OP_SYSCALL]      = "syscall",
        [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>";
}

i32 main(i32 argc, char *argv[]) {
  struct CompilerConfig config = {0};

  if (parse_arguments(argc, argv, &config) != 0) {
    fprintf(stderr, "Usage: %s [-d] [-g] [-o] <file1.ul> [file2.ul] ...\n",
            argv[0]);
    return 64;
  }

  VM vm = {0};
  if (config.input_file_count == 0) {
    repl(&vm);
  } else {
    i32 j;
    for (j = 0; j < config.input_file_count; j++) {
      compileFile(config.input_files[j], &vm);
    }

    if (config.flags & FLAG_DUMP_ROM) {
      FILE *file = fopen("memory_dump.bin", "wb");
      if (!file) {
        perror("Failed to open file");
        return EXIT_FAILURE;
      }

      size_t code_written = fwrite(vm.code, 1, CODE_SIZE, file);
      if (code_written != CODE_SIZE) {
        fprintf(stderr, "Incomplete code write: %zu bytes written out of %u\n",
                code_written, CODE_SIZE);
        fclose(file);
        return EXIT_FAILURE;
      }

      size_t memory_written = fwrite(vm.memory, 1, MEMORY_SIZE, file);
      if (memory_written != MEMORY_SIZE) {
        fprintf(stderr,
                "Incomplete memory write: %zu bytes written out of %u\n",
                memory_written, MEMORY_SIZE);
        fclose(file);
        return EXIT_FAILURE;
      }

      fclose(file);
    }
  }
  bool running = true;
  vm_register_device(&vm, "/dev/term/0", "terminal", nil, &console_device_ops);

  if (config.flags & FLAG_GUI_MODE) {
    u32 i;
    register_sdl_devices(&vm);
    while (running) {
      for (i = 0; i < vm.dc; i++) {
        Device *dev = &vm.devices[i];
        if (strcmp(dev->type, "mouse") == 0) {
          MouseDeviceData *mouse = (MouseDeviceData *)dev->data;
          SDL_PumpEvents();
          SDL_Event event;
          while (SDL_PollEvent(&event)) {
            switch (event.type) {
            case SDL_QUIT:
              running = false;
              break;
            case SDL_MOUSEMOTION:
              mouse->x = event.motion.x;
              mouse->y = event.motion.y;
              break;
            case SDL_MOUSEBUTTONDOWN:
              if (event.button.button == SDL_BUTTON_LEFT)
                mouse->btn1 = 1;
              if (event.button.button == SDL_BUTTON_RIGHT)
                mouse->btn2 = 1;
              if (event.button.button == SDL_BUTTON_MIDDLE)
                mouse->btn3 = 1;
              if (event.button.button == SDL_BUTTON_X1)
                mouse->btn4 = 1;
              break;
            case SDL_MOUSEBUTTONUP:
              if (event.button.button == SDL_BUTTON_LEFT)
                mouse->btn1 = 0;
              if (event.button.button == SDL_BUTTON_RIGHT)
                mouse->btn2 = 0;
              if (event.button.button == SDL_BUTTON_MIDDLE)
                mouse->btn3 = 0;
              if (event.button.button == SDL_BUTTON_X1)
                mouse->btn4 = 0;
              break;
            }
          }
        }
      }

      step_vm(&vm);

      for (i = 0; i < vm.dc; i++) {
        Device *dev = &vm.devices[i];
        if (strcmp(dev->type, "screen") == 0) {
          ScreenDeviceData *screen = (ScreenDeviceData *)dev->data;
          if (screen->texture && screen->renderer) {
            SDL_UpdateTexture(screen->texture, NULL,
                              &vm.memory[screen->framebuffer_pos],
                              screen->width * sizeof(u32));

            SDL_RenderClear(screen->renderer);

            SDL_Rect output_rect;
            SDL_RenderGetViewport(screen->renderer, &output_rect);
            float scale_x = (float)output_rect.w / screen->width;
            float scale_y = (float)output_rect.h / screen->height;
            float scale = SDL_min(scale_x, scale_y);

            SDL_Rect dstrect = {
                (i32)((output_rect.w - screen->width * scale) / 2),
                (i32)((output_rect.h - screen->height * scale) / 2),
                (i32)(screen->width * scale), (i32)(screen->height * scale)};

            SDL_RenderCopy(screen->renderer, screen->texture, NULL, &dstrect);
            SDL_RenderPresent(screen->renderer);
          }
          break;
        }
      }
    }
  } else {
    while (running) {
      //printf("| %d %s %d\n", vm.code[vm.pc], opcode_to_string(vm.code[vm.pc]), vm.pc);
      running = step_vm(&vm);
    }
  }

  return 0;
}