#include "../../tools/assembler/assembler.h"
#include "../../vm/vm.h"
#include "devices.h"

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

#define MAX_SRC_SIZE 16384

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

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);
  USED(source_file);

  return true;
}

#ifdef STATIC
#define SCOPES_COUNT 2048
SymbolTable scopes[SCOPES_COUNT];
#endif

void symbol_table_init(ScopeTable *t) {
#ifdef STATIC
  memset(scopes, 0, SCOPES_COUNT * sizeof(SymbolTable));
  t->scopes = scopes;
  t->count = 0;
  t->capacity = SCOPES_COUNT;
#else
  t->scopes = calloc(16, sizeof(SymbolTable));
  t->count = 0;
  t->capacity = 16;
#endif

  // Make sure that all the parents are the 'global' namespace.
  for (u32 i = 0; i < t->capacity; i++) {
    t->scopes[i].parent = -1;
  }
}

bool table_realloc(ScopeTable *table) {
#ifdef STATIC
  if (table->count >= table->capacity) {
    return false;
  }
#else
  if (table->count >= table->capacity) {
    table->capacity *= 2;
    table->scopes =
        realloc(table->scopes, table->capacity * sizeof(SymbolTable));

    // Make sure that all the parents are the 'global' namespace.
    for (u32 i = table->count; i < table->capacity; i++) {
      table->scopes[i].parent = -1;
    }
  }
#endif
  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);

  ScopeTable table = {0};
  symbol_table_init(&table);
  assemble(vm, &table, source);
#ifndef STATIC
  free(table.scopes);
#endif

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

i32 main(i32 argc, char *argv[]) {
  bool dump_rom = false;
  char *input_file = nil;
  char *output_file = nil;
  bool is_rom = 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, ".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};
  bool compilation_success = true;
  if (input_file) {
    if (is_rom) {
      // Load ROM file directly
      compilation_success = loadVM(input_file, &vm);
    } else if (is_ir) {
      // 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);

  bool running = true;
  while (running) {
    if (!step_vm(&vm)) {
      running = false;
      break;
    }
  }

  return vm.flag;
}