#include "../../tools/assembler.h"
#include "../../tools/parser.h"
#include "../../tools/lexer.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};
static ConsoleDeviceData console_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->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, vm->cp, file) != vm->cp) {
    fprintf(stderr, "Failed to write code section\n");
    fclose(file);
    return false;
  }

  // Write memory section
  if (fwrite(vm->memory, 1, vm->mp, file) != vm->mp) {
    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->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, vm->cp, file) != vm->cp) {
    fprintf(stderr, "Failed to read code section\n");
    fclose(file);
    return false;
  }

  // Read memory section
  if (fread(vm->memory, 1, vm->mp, file) != vm->mp) {
    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) {
  USED(vm);
  USED(output_file);
  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);

  initLexer(source);
  Token token;
  do {
    token = nextToken();
    if (token.type == TOKEN_ERROR) {
      printf("ERROR at line %d: %.*s\n", token.line, token.length, token.start);
      break; // Stop on error, or continue if you want to see more
    }
    if (token.type != TOKEN_EOF) {
      printf("Line %d [%s]: %.*s\n", 
            token.line, 
            tokenTypeToString(token.type), 
            token.length, 
            token.start);
    }
  } while (token.type != TOKEN_EOF);

  return true;
}

// Function to assemble and optionally save
bool assembleAndSave(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 {
    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);
}

#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",
    
    /* Immediate loads (only 32-bit variant needed) */
    [OP_LOAD_IMM]       = "load-immediate",
    
    /* Register-indirect loads */
    [OP_LOAD_IND_8]     = "load-indirect-8",
    [OP_LOAD_IND_16]    = "load-indirect-16",
    [OP_LOAD_IND_32]    = "load-indirect-32",
    
    /* Absolute address loads */
    [OP_LOAD_ABS_8]     = "load-absolute-8",
    [OP_LOAD_ABS_16]    = "load-absolute-16",
    [OP_LOAD_ABS_32]    = "load-absolute-32",
    
    /* Base+offset loads */
    [OP_LOAD_OFF_8]     = "load-offset-8",
    [OP_LOAD_OFF_16]    = "load-offset-16",
    [OP_LOAD_OFF_32]    = "load-offset-32",
    
    /* Absolute address stores */
    [OP_STORE_ABS_8]    = "store-absolute-8",
    [OP_STORE_ABS_16]   = "store-absolute-16",
    [OP_STORE_ABS_32]   = "store-absolute-32",
    
    /* Register-indirect stores */
    [OP_STORE_IND_8]    = "store-indirect-8",
    [OP_STORE_IND_16]   = "store-indirect-16",
    [OP_STORE_IND_32]   = "store-indirect-32",
    
    /* Base+offset stores */
    [OP_STORE_OFF_8]    = "store-offset-8",
    [OP_STORE_OFF_16]   = "store-offset-16",
    [OP_STORE_OFF_32]   = "store-offset-32",
    
    /* Memory operations */
    [OP_MALLOC]         = "malloc",
    [OP_MEMSET_8]       = "memset-8",
    [OP_MEMSET_16]      = "memset-16",
    [OP_MEMSET_32]      = "memset-32",
    
    /* Stack operations */
    [OP_PUSH]           = "push",
    [OP_POP]            = "pop",
    
    /* Register operations */
    [OP_REG_MOV]        = "register-move",
    [OP_SYSCALL]        = "syscall",
    
    /* Bit operations */
    [OP_SLL]            = "bit-shift-left",
    [OP_SRL]            = "bit-shift-right",
    [OP_SRE]            = "bit-shift-re",
    [OP_BAND]           = "bit-and",
    [OP_BOR]            = "bit-or",
    [OP_BXOR]           = "bit-xor",
    
    /* Integer arithmetic */
    [OP_ADD_INT]        = "add-int",
    [OP_SUB_INT]        = "sub-int",
    [OP_MUL_INT]        = "mul-int",
    [OP_DIV_INT]        = "div-int",
    
    /* Natural number arithmetic */
    [OP_ADD_NAT]        = "add-nat",
    [OP_SUB_NAT]        = "sub-nat",
    [OP_MUL_NAT]        = "mul-nat",
    [OP_DIV_NAT]        = "div-nat",
    
    /* Floating point operations */
    [OP_ADD_REAL]       = "add-real",
    [OP_SUB_REAL]       = "sub-real",
    [OP_MUL_REAL]       = "mul-real",
    [OP_DIV_REAL]       = "div-real",
    
    /* Type conversions */
    [OP_INT_TO_REAL]    = "int-to-real",
    [OP_NAT_TO_REAL]    = "nat-to-real",
    [OP_REAL_TO_INT]    = "real-to-int",
    [OP_REAL_TO_NAT]    = "real-to-nat",
    
    /* Integer comparisons */
    [OP_JEQ_INT]        = "jump-eq-int",
    [OP_JNEQ_INT]       = "jump-neq-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",
    
    /* Natural number comparisons */
    [OP_JEQ_NAT]        = "jump-eq-nat",
    [OP_JNEQ_NAT]       = "jump-neq-nat",
    [OP_JGT_NAT]        = "jump-gt-nat",
    [OP_JLT_NAT]        = "jump-lt-nat",
    [OP_JLE_NAT]        = "jump-le-nat",
    [OP_JGE_NAT]        = "jump-ge-nat",
    
    /* Floating point comparisons */
    [OP_JEQ_REAL]       = "jump-eq-real",
    [OP_JNEQ_REAL]      = "jump-neq-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",
    
    /* String operations */
    [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",
    
    /* String conversions */
    [OP_INT_TO_STRING]  = "int-to-string",
    [OP_NAT_TO_STRING]  = "nat-to-string",
    [OP_REAL_TO_STRING] = "real-to-string",
    [OP_STRING_TO_INT]  = "string-to-int",
    [OP_STRING_TO_NAT]  = "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, is_assembly = 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;
      }
      if (ext && (strcmp(ext, ".lisp") == 0)) {
        is_assembly = 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 if (is_assembly) {
      // Compile Lisp file
      if (dump_rom && output_file) {
        compilation_success = assembleAndSave(input_file, output_file, &vm);
      } else {
        compilation_success = assembleAndSave(input_file, nil, &vm);
      }
    } else {
      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;
  }

  vm_register_device(&vm, "/dev/term/0", "terminal", &console_data,
                     &console_device_ops, 4);

  if (gui_mode) {
    if (SDL_Init(SDL_INIT_VIDEO) < 0) {
      printf("SDL initialization failed: %s\n", SDL_GetError());
      return 1;
    }
    SDL_SetHint(SDL_HINT_TOUCH_MOUSE_EVENTS, "0");

    screen_data.width = 640;
    screen_data.height = 480;
    screen_data.buffer_size = screen_data.width * screen_data.height;

    vm_register_device(&vm, "/dev/screen/0", "screen", &screen_data,
                       &screen_ops, 16 + screen_data.buffer_size);

    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 = 16;

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

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

    SDL_Event event;
    bool running = true;
    SDL_PumpEvents();
    while (running) {
      while (SDL_PollEvent(&event)) {
        switch (event.type) {
        case SDL_QUIT:
          running = false;
          break;
        // Mouse events
        case SDL_MOUSEMOTION:
          mouse_data.x = event.motion.x;
          mouse_data.y = event.motion.y;
          break;

        case SDL_MOUSEBUTTONDOWN:
          if (event.button.button == SDL_BUTTON_LEFT)
            mouse_data.btn1 = 1;
          if (event.button.button == SDL_BUTTON_RIGHT)
            mouse_data.btn2 = 1;
          if (event.button.button == SDL_BUTTON_MIDDLE)
            mouse_data.btn3 = 1;
          if (event.button.button == SDL_BUTTON_X1)
            mouse_data.btn4 = 1;
          break;

        case SDL_MOUSEBUTTONUP:
          if (event.button.button == SDL_BUTTON_LEFT)
            mouse_data.btn1 = 0;
          if (event.button.button == SDL_BUTTON_RIGHT)
            mouse_data.btn2 = 0;
          if (event.button.button == SDL_BUTTON_MIDDLE)
            mouse_data.btn3 = 0;
          if (event.button.button == SDL_BUTTON_X1)
            mouse_data.btn4 = 0;
          break;

        // Touch events (map to mouse_data as left-click equivalent)
        case SDL_FINGERMOTION:
        case SDL_FINGERDOWN:
        case SDL_FINGERUP: {

          float x = event.tfinger.x * 640;
          float y = event.tfinger.y * 480;

          mouse_data.x = (int)x;
          mouse_data.y = (int)y;

          // Only treat the first finger as mouse input (ignore multi-touch
          // beyond 1 finger)
          if (event.tfinger.fingerId == 0) {
            if (event.type == SDL_FINGERDOWN ||
                event.type == SDL_FINGERMOTION) {
              mouse_data.btn1 = 1;
            } else if (event.type == SDL_FINGERUP) {
              mouse_data.btn1 = 0;
            }
          }
          break;
        }
        }
      }

      // Run VM for a fixed number of cycles or a time slice
      int cycles_this_frame = 0;
      int max_cycles_per_frame = 100; // Adjust this value
      while (cycles_this_frame < max_cycles_per_frame) {
        #ifdef ASM_DEBUG
        printf("| %s %d\n", opcode_to_string(vm.code[vm.pc]),vm.pc);
        #endif
        if (!step_vm(&vm)) {
          running = false;
          break;
        }
        cycles_this_frame++;
      }

      // Render only if the screen buffer was updated AND at a reasonable rate
      if (screen_data.update) {
        if (screen_data.renderer && screen_data.texture) {
          // Clear and render
          SDL_RenderClear(screen_data.renderer);

          SDL_Rect output_rect;
          SDL_RenderGetViewport(screen_data.renderer, &output_rect);

          // Calculate aspect ratio preserving scaling
          float scale_x = (float)output_rect.w / screen_data.width;
          float scale_y = (float)output_rect.h / screen_data.height;
          float scale = SDL_min(scale_x, scale_y);

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

          SDL_RenderCopy(screen_data.renderer, screen_data.texture, NULL,
                         &dstrect);
          SDL_RenderPresent(screen_data.renderer);
        }
        screen_data.update = false; // Reset flag after rendering
      }
    }
  } else {
    bool running = true;
    while (running) {
      running = step_vm(&vm);
    }
  }

  return EXIT_SUCCESS;
}