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

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

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

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

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

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 (locals 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 (locals 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);

  assemble(vm, source);

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

bool init_vm(VM *vm) {
  vm->memory = (u8*)malloc(MEMORY_SIZE * sizeof(u8));
  vm->memory_size = MEMORY_SIZE;

  vm->code = (u8*)malloc(CODE_SIZE * sizeof(u8));
  vm->code_size = CODE_SIZE;

  vm->frames = (Frame*)malloc(FRAMES_SIZE * sizeof(Frame));
  vm->frames_size = FRAMES_SIZE;

  vm->stack = (u32*)malloc(STACK_SIZE * sizeof(u32))
  vm->stack_size = STACK_SIZE;

  vm->devices = (Device*)malloc(DEVICES_SIZE * sizeof(Device));
  vm->devices_size = DEVICES_SIZE;
  
  return true;
}

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

  // Parse command line arguments
  for (i32 i = 1; i < argc; i++) {
    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;
      }
      if (ext && (strcmp(ext, ".ir") == 0)) {
        is_ir = true;
      }
    } else if (output_file == nil && dump_rom) {
      // This is the output file for -o flag
      output_file = argv[i];
    }
  }

  VM vm = {0};
  if (!init_vm(&vm)) {
    printf("vm did not initialize for some reason.");
    return 1;
  }

  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 {
    printf("usage: undar <src.ul>...");
    return 1;
  }

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

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

  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: {

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

        mouse_data.x = (i32)x;
        mouse_data.y = (i32)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
    i32 cycles_this_frame = 0;
    i32 max_cycles_per_frame = 100; // Adjust this value
    while (cycles_this_frame < max_cycles_per_frame) {
      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
        f32 scale_x = (f32)output_rect.w / screen_data.width;
        f32 scale_y = (f32)output_rect.h / screen_data.height;
        f32 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
    }
  }

  return vm.flag;
}