#include <stdio.h>
#include <stdarg.h>
#include <stdlib.h>
#include <stdint.h>
#include <signal.h>
#include <string.h>
#include <math.h>
#include <raylib.h>

#define MEMORY_SIZE 4096
#define PROGRAM_START 512
#define FONT_START 80
#define CPU_CLOCK_HZ 600
#define TIMER_HZ 60
#define SCREEN_WIDTH 64
#define SCREEN_HEIGHT 32

enum OpcodeGroups {
  OPCODE_GROUP_SYS_OR_RET = 0,
  OPCODE_GROUP_JUMP = 4096,
  OPCODE_GROUP_CALL = 8192,
  OPCODE_GROUP_SE_VX_BYTE = 12288,
  OPCODE_GROUP_SNE_VX_BYTE = 16384,
  OPCODE_GROUP_SE_VX_VY = 20480,
  OPCODE_GROUP_LD_VX_BYTE = 24576,
  OPCODE_GROUP_ADD_VX_BYTE = 28672,
  OPCODE_GROUP_ARITHMETIC = 32768,
  OPCODE_GROUP_SNE_VX_VY = 36864,
  OPCODE_GROUP_LD_I_ADDR = 40960,
  OPCODE_GROUP_JP_V0_ADDR = 45056,
  OPCODE_GROUP_RND_VX_BYTE = 49152,
  OPCODE_GROUP_DRW_VX_VY_N = 53248,
  OPCODE_GROUP_KEY_SKIP = 57344,
  OPCODE_GROUP_MISC = 61440
};

enum SysOpcodes {
  OPCODE_CLS = 224,
  OPCODE_RET = 238
};

enum ArithmeticSubOpcodes {
  SUB_OP_LD = 0,
  SUB_OP_OR = 1,
  SUB_OP_AND = 2,
  SUB_OP_XOR = 3,
  SUB_OP_ADD = 4,
  SUB_OP_SUB = 5,
  SUB_OP_SHR = 6,
  SUB_OP_SUBN = 7,
  SUB_OP_SHL = 14
};

enum KeySkipOpcodes {
  OPCODE_SKP = 158,
  OPCODE_SKNP = 161
};

enum MiscSubOpcodes {
  SUB_OP_LD_VX_DT = 7,
  SUB_OP_LD_VX_K = 10,
  SUB_OP_LD_DT_VX = 21,
  SUB_OP_LD_ST_VX = 24,
  SUB_OP_ADD_I_VX = 30,
  SUB_OP_LD_F_VX = 41,
  SUB_OP_LD_B_VX = 51,
  SUB_OP_LD_I_REG = 85,
  SUB_OP_LD_REG_I = 101
};

const uint8_t fontset[80] = {
  0xF0, 0x90, 0x90, 0x90, 0xF0, // 0
  0x20, 0x60, 0x20, 0x20, 0x70, // 1
  0xF0, 0x10, 0xF0, 0x80, 0xF0, // 2
  0xF0, 0x10, 0xF0, 0x10, 0xF0, // 3
  0x90, 0x90, 0xF0, 0x10, 0x10, // 4
  0xF0, 0x80, 0xF0, 0x10, 0xF0, // 5
  0xF0, 0x80, 0xF0, 0x90, 0xF0, // 6
  0xF0, 0x10, 0x20, 0x40, 0x40, // 7
  0xF0, 0x90, 0xF0, 0x90, 0xF0, // 8
  0xF0, 0x90, 0xF0, 0x10, 0xF0, // 9
  0xF0, 0x90, 0xF0, 0x90, 0x90, // A
  0xE0, 0x90, 0xE0, 0x90, 0xE0, // B
  0xF0, 0x80, 0x80, 0x80, 0xF0, // C
  0xE0, 0x90, 0x90, 0x90, 0xE0, // D
  0xF0, 0x80, 0xF0, 0x80, 0xF0, // E
  0xF0, 0x80, 0xF0, 0x80, 0x80  // F
};

const int CYCLES_PER_FRAME = CPU_CLOCK_HZ / TIMER_HZ;

volatile sig_atomic_t keep_running = 1;

void handle_sigint(int sig) {
  (void)sig;
  keep_running = 0;
}

enum ExitCodes {
  COMPLETED_SUCCESSFULLY,
  FILE_NOT_PROVIDED,
  FILE_MISSING,
  PROGRAM_TOO_LARGE,
  MEMORY_ALLOCATION_FAILURE,
};

// TODO: Add additional members
// i.e. registers
typedef struct ChipEight {
  uint8_t memory[MEMORY_SIZE];
  uint16_t pc;
  uint16_t I;
  uint8_t V[16];
  uint8_t dt;
  uint8_t st;
  uint16_t stack[16];
  uint8_t sc;
  uint32_t display[SCREEN_WIDTH * SCREEN_HEIGHT];
  uint8_t keys[16];
  uint8_t key_register;
  bool waiting_for_key;
} chip_eight_t;

void println(const char *format, ...) {
  va_list args;

  va_start(args, format);

  vprintf(format, args);

  va_end(args);

  printf("\n");
}


void setup_graceful_exit() {
  signal(SIGINT, handle_sigint);
  signal(SIGTERM, handle_sigint);
}

int load_rom(chip_eight_t *chip, char *file_path) {
  FILE *file_ptr = fopen(file_path, "rb");

  if (file_ptr == NULL) {
    return FILE_MISSING;
  }

  fseek(file_ptr, 0, SEEK_END);
  int num_of_bytes = ftell(file_ptr);

  int max_program_size = MEMORY_SIZE - PROGRAM_START;

  if (num_of_bytes > max_program_size) {
    fclose(file_ptr);
    return PROGRAM_TOO_LARGE;
  }

  fseek(file_ptr, 0, SEEK_SET);

  fread(&chip->memory[PROGRAM_START], sizeof(uint8_t), num_of_bytes, file_ptr);
  
  fclose(file_ptr);

  return COMPLETED_SUCCESSFULLY;
}

void init_chip(chip_eight_t *chip) {
  memset(chip, 0, sizeof(chip_eight_t));
  chip->pc = PROGRAM_START;
  chip->dt = 0;
  chip->st = 0;
  chip->waiting_for_key = false;

  for (int i = 0; i < 80; i++) {
    chip->memory[FONT_START + i] = fontset[i];
  }
}

void draw_sprite(chip_eight_t *chip, uint16_t n, uint16_t x, uint16_t y) {
  uint8_t x_coordinate = chip->V[x] % SCREEN_WIDTH;
  uint8_t y_coordinate = chip->V[y] % SCREEN_HEIGHT;
  chip->V[15] = 0;

  for (int row = 0; row < n; row++) {
    uint8_t sprite_byte = chip->memory[chip->I + row];

    if ((y_coordinate + row) >= SCREEN_HEIGHT) {
      break;
    }

    for (int col = 0; col < 8; col++) {
      if ((x_coordinate + col) >= SCREEN_WIDTH) {
        break;
      }
      
      // extracting individual bits (pixels) from byte in memory
      uint8_t sprite_pixel = sprite_byte & (128 >> col);

      // flatten 2-d coordinates to index
      int sprite_screen_index = ((y_coordinate + row) * SCREEN_WIDTH) + (x_coordinate + col);
      uint32_t *screen_pixel = &chip->display[sprite_screen_index];

      if (sprite_pixel != 0) {
        if (*screen_pixel == 1) {
          chip->V[15] = 1;
        }

        *screen_pixel ^= 1;
      }
    }
  }
}

void update_keys(chip_eight_t *chip) {
  chip->keys[1] = IsKeyDown(KEY_ONE);
  chip->keys[2] = IsKeyDown(KEY_TWO);
  chip->keys[3] = IsKeyDown(KEY_THREE);
  chip->keys[12] = IsKeyDown(KEY_FOUR);

  chip->keys[4] = IsKeyDown(KEY_Q);
  chip->keys[5] = IsKeyDown(KEY_W);
  chip->keys[6] = IsKeyDown(KEY_E);
  chip->keys[13] = IsKeyDown(KEY_R);

  chip->keys[7] = IsKeyDown(KEY_A);
  chip->keys[8] = IsKeyDown(KEY_S);
  chip->keys[9] = IsKeyDown(KEY_D);
  chip->keys[14] = IsKeyDown(KEY_F);

  chip->keys[10] = IsKeyDown(KEY_Z);
  chip->keys[0] = IsKeyDown(KEY_X);
  chip->keys[11] = IsKeyDown(KEY_C);
  chip->keys[15] = IsKeyDown(KEY_V);
}

int main(int argc, char *argv[]) {
  setup_graceful_exit();

  if (argc < 2) {
    println("You fucked up bro! We need a file to load into memory.");
    println("Usage: chip.exe <file-to-load>");
    return FILE_NOT_PROVIDED;
  }

  chip_eight_t *chip = malloc(sizeof(chip_eight_t));

  if (chip == NULL) {
    println("Failed to allocate memory for chip.");
    return MEMORY_ALLOCATION_FAILURE;
  }

  init_chip(chip);

  int load_rom_result = load_rom(chip, argv[1]);

  if (load_rom_result != COMPLETED_SUCCESSFULLY) {
    if (load_rom_result == FILE_MISSING) {
      println("Error: The file '%s' could not be found or opened.", argv[1]);
    } else if (load_rom_result == PROGRAM_TOO_LARGE) {
      println("Error: The program is too large to fit in CHIP-8 memory.");
    } else {
      println("Error: Failed to load ROM (code %d).", load_rom_result);
    }
    free(chip);
    return load_rom_result;
  }

  // Setup audio
  InitAudioDevice();

  // Generate 440Hz Sine wave beep in memory
  int sampleRate = 44100;
  float durationSeconds = 0.1f;
  int frameCount = sampleRate * durationSeconds;
  float *data = (float *)malloc(frameCount * sizeof(float));
  if (data != NULL) {
    for (int i = 0; i < frameCount; i++) {
      data[i] = sinf(2.0f * PI * 440.0f * ((float)i / sampleRate)) * 0.2f;
    }
  }
  Wave wave = { 0 };
  wave.frameCount = frameCount;
  wave.sampleRate = sampleRate;
  wave.sampleSize = 32;
  wave.channels = 1;
  wave.data = data;

  Sound beep_sound = LoadSoundFromWave(wave);
  UnloadWave(wave); // Frees the 'data' buffer in system RAM

  int scale_factor = 10;
  int display_width = SCREEN_WIDTH * scale_factor;
  int display_height = SCREEN_HEIGHT * scale_factor;
  InitWindow(display_width, display_height, "chip");
  SetTargetFPS(60);

  while(!WindowShouldClose() && keep_running)
  {
    update_keys(chip);

    if (chip->waiting_for_key) {
      bool key_pressed = false;

      for (int i = 0; i < 16; i++) {
        if (chip->keys[i]) {
          chip->V[chip->key_register] = i;
          chip->waiting_for_key = false;
          key_pressed = true;
          break;
        }
      }

      if (!key_pressed) {
        BeginDrawing();
        ClearBackground(BLACK);

        for (int i = 0; i < SCREEN_HEIGHT; i++) {
          for (int j = 0; j < SCREEN_WIDTH; j++) {
            uint32_t screen_pixel = chip->display[i * SCREEN_WIDTH + j];

            if (screen_pixel == 1) {
              DrawRectangle(j * scale_factor, i * scale_factor, scale_factor, scale_factor, RAYWHITE);
            }
          }
        }

        EndDrawing();
        continue;
      }

    }

    for(int i = 0; i < CYCLES_PER_FRAME; i++) {
      if (chip->waiting_for_key) {
        break;
      }
      // read 2 bytes at a time from memory
      // reading first 1 byte shifting left 1 byte ORing 2 byte
      uint16_t opcode = (chip->memory[chip->pc] << 8) | (chip->memory[chip->pc + 1]);
      chip->pc += 2;
      
      // extracting the lowest 12 bits of the instruction
      uint16_t nnn = opcode & 4095;

      //extracting the lowest 8 bits of the instruction
      uint8_t nn = opcode & 255;

      //extracting the lowest 4 bits of the instruction
      uint16_t n = opcode & 15;

      // extracting the lower 4 bits of the high byte of the instruction
      uint16_t x = (opcode & 3840) >> 8;

      // extracting the upper 4 bits of the low byte of the instruction
      uint16_t y = (opcode & 240) >> 4;
      
      // extracting the lowest 8 bits of the instruction
      uint16_t kk = opcode & 255;

      switch(opcode & 61440) {
        case OPCODE_GROUP_SYS_OR_RET: 
          if (opcode == OPCODE_CLS) {
            memset(chip->display, 0, sizeof(chip->display));
          } else if (opcode == OPCODE_RET) {
            chip->sc -= 1;
            chip->pc = chip->stack[chip->sc];
          }
          break;
        case OPCODE_GROUP_JUMP:
          chip->pc = nnn;
          break;
        case OPCODE_GROUP_CALL:
          chip->stack[chip->sc] = chip->pc;
          chip->sc += 1;
          chip->pc = nnn;
          break;
        case OPCODE_GROUP_SE_VX_BYTE:
          if (chip->V[x] == kk) {
            chip->pc += 2;
          }
          break;
        case OPCODE_GROUP_SNE_VX_BYTE:
          if (chip->V[x] != kk) {
            chip->pc += 2;
          }
          break;
        case OPCODE_GROUP_SE_VX_VY:
          if (chip->V[x] == chip->V[y]) {
            chip->pc += 2;
          }
          break;
        case OPCODE_GROUP_LD_VX_BYTE:
          chip->V[x] = kk;
          break;
        case OPCODE_GROUP_ADD_VX_BYTE:
          chip->V[x] += kk;
          break;
        case OPCODE_GROUP_ARITHMETIC:
          if (n == SUB_OP_LD) {
            chip->V[x] = chip->V[y];
          } else if (n == SUB_OP_OR) {
            chip->V[x] = chip->V[x] | chip->V[y];
          } else if (n == SUB_OP_AND) {
            chip->V[x] = chip->V[x] & chip->V[y];
          } else if (n == SUB_OP_XOR) {
            chip->V[x] = chip->V[x] ^ chip->V[y];
          } else if (n == SUB_OP_ADD) {
            uint16_t r = chip->V[x] + chip->V[y];
            chip->V[15] = (r > 255) ? 1 : 0;
            chip->V[x] = r & 255;
          } else if (n == SUB_OP_SUB) {
            chip->V[15] = (chip->V[x] >= chip->V[y]) ? 1 : 0;
            chip->V[x] = chip->V[x] - chip->V[y];
          } else if (n == SUB_OP_SHR) {
            uint8_t flag = chip->V[x] & 1;
            chip->V[15] = flag;
            chip->V[x] /= 2;
          } else if (n == SUB_OP_SUBN) {
            chip->V[15] = (chip->V[y] >= chip->V[x]) ? 1 : 0;
            chip->V[x] = chip->V[y] - chip->V[x];
          } else if (n == SUB_OP_SHL) {
            uint8_t flag = (chip->V[x] & 128) >> 7;
            chip->V[15] = flag;
            chip->V[x] *= 2;
          }
          break;
        case OPCODE_GROUP_SNE_VX_VY:
          if (chip->V[x] != chip->V[y]) {
            chip->pc += 2;
          }
          break;
        case OPCODE_GROUP_LD_I_ADDR:
          chip->I = nnn;
          break;
        case OPCODE_GROUP_JP_V0_ADDR:
          chip->pc = nnn + chip->V[0];
          break;
        case OPCODE_GROUP_RND_VX_BYTE:
          uint8_t num = rand() % 256;
          chip->V[x] = kk & num;
          break;
        case OPCODE_GROUP_DRW_VX_VY_N:
          draw_sprite(chip, n, x, y);
          break;
        case OPCODE_GROUP_KEY_SKIP:
          if (nn == OPCODE_SKP) {
            if (chip->keys[chip->V[x] & 0x0F] == 1) {
              chip->pc += 2;
            }
          } else if (nn == OPCODE_SKNP) {
            if (chip->keys[chip->V[x] & 0x0F] != 1) {
              chip->pc += 2;
            }
          }
          break;
        case OPCODE_GROUP_MISC:
          if (nn == SUB_OP_LD_VX_DT) {
            chip->V[x] = chip->dt;
          } else if (nn == SUB_OP_LD_VX_K) {
            chip->waiting_for_key = true;
            chip->key_register = x;
          } else if (nn == SUB_OP_LD_DT_VX) {
            chip->dt = chip->V[x];
          } else if (nn == SUB_OP_LD_ST_VX) {
            chip->st = chip->V[x];
          } else if (nn == SUB_OP_ADD_I_VX) {
            chip->I += chip->V[x];
          } else if (nn == SUB_OP_LD_F_VX) {
            chip->I = FONT_START + ((chip->V[x] & 0x0F) * 5);
          } else if (nn == SUB_OP_LD_B_VX) {
            chip->memory[chip->I] = chip->V[x] / 100;
            chip->memory[chip->I + 1] = (chip->V[x] / 10) % 10;
            chip->memory[chip->I + 2] = chip->V[x] % 10;
          } else if (nn == SUB_OP_LD_I_REG) {
            for (int i = 0; i <= x; i++) {
              chip->memory[chip->I + i] = chip->V[i];
            }
          } else if (nn == SUB_OP_LD_REG_I) {
            for (int i = 0; i <= x; i++) {
              chip->V[i] = chip->memory[chip->I + i];
            }
          }
          break;
        default:
          break;
      }
    }

    if (chip->dt > 0) {
      chip->dt--;
    }

    if (chip->st > 0) {
      chip->st--;
      if (!IsSoundPlaying(beep_sound)) {
        PlaySound(beep_sound);
      }
    } else {
      if (IsSoundPlaying(beep_sound)) {
        StopSound(beep_sound);
      }
    }

    BeginDrawing();
    ClearBackground(BLACK);

    for (int i = 0; i < SCREEN_HEIGHT; i++) {
      for (int j = 0; j < SCREEN_WIDTH; j++) {
        uint32_t screen_pixel = chip->display[i * SCREEN_WIDTH + j];

        if (screen_pixel == 1) {
          DrawRectangle(j * scale_factor, i * scale_factor, scale_factor, scale_factor, RAYWHITE);
        }
      }
    }

    EndDrawing();
  }

  free(chip);

  UnloadSound(beep_sound);
  CloseAudioDevice();

  CloseWindow();

  return COMPLETED_SUCCESSFULLY;
}