#include <cstddef>
#include <cstdint>
#include <cstdio>
#include <cstring>
#include <filesystem>
#include <fstream>

#define STB_IMAGE_IMPLEMENTATION
#include "stb_image.h"
#include "fastlz.h"

extern "C" char *RLE_Code(unsigned char *raw_buffer, int raw_len, int *new_len);
extern "C" unsigned char *LZS_Fast(unsigned char *raw_buffer, size_t raw_len, size_t *new_len);
extern "C" unsigned char *LZS_Code(unsigned char *raw_buffer, size_t raw_len, size_t *new_len, size_t best);

#define LZS_WRAM      0x00       // VRAM not compatible (LZS_WRAM | LZS_NORMAL)
#define LZS_VRAM      0x01       // VRAM compatible (LZS_VRAM | LZS_NORMAL)

#define RGBtoRGB15(r,g,b) (((r >> 3) & 0x1F) | (((g >> 3) & 0x1F) << 5) | (((b >> 3) & 0x1F) << 10))

#define ASSETS_PATH "../assets/"

#define PALETTE_SIZE 256
#define FRAME_SIZE 256*192
#define CHUNK_SIZE 8

void Frame_RAWtoRGB15(unsigned char *frame, unsigned char *dest, uint16_t *palette_buffer);

int Compress(uint8_t type, uint8_t *frame_buffer, uint32_t frame_buffer_len, uint8_t *compress_buffer)
{
  switch (type) {
    case 1: // fastlz
      return fastlz_compress_level(1, frame_buffer, frame_buffer_len, compress_buffer);
      break;

    case 2: // rle
      {
        int compress_len;
        char *compress = RLE_Code(frame_buffer, frame_buffer_len, &compress_len);
        memcpy(compress_buffer, compress, compress_len);
        return compress_len;
        break;
      }
    
    case 3: // lzss
      {
        size_t compress_len;
        unsigned char *compress = LZS_Fast(frame_buffer, frame_buffer_len, &compress_len);
        // int mode = LZS_VRAM;
        // unsigned char *compress = LZS_Code(frame_buffer, frame_buffer_len, &compress_len, mode);
        memcpy(compress_buffer, compress, compress_len);
        return compress_len;
        break;
      }
  }

  return 0;
}

void Convert(std::string frame_path, std::string palette_path, std::string output_path, std::string output_basename, uint8_t type)
{
  uint16_t palette_buffer[PALETTE_SIZE];

  size_t frame_total = 0;
  uint8_t frame_buffer[FRAME_SIZE*CHUNK_SIZE];

  uint32_t compress_size = 0;
  uint32_t compress_size_total = 0;
  size_t compress_size_biggest = 0;
  uint8_t compress_buffer[162000];

  std::ofstream file_out;
  
  printf("Finding total frame...");
  // Find total frame
  for (const auto& entry : std::filesystem::directory_iterator(frame_path)) {
    if (entry.is_regular_file()) {
      std::string filename = entry.path().filename().string();
      int extpos = filename.find(".bmp", 4);
      if (extpos != std::string::npos)
      {
        int number = std::stoi(filename.substr(4, extpos-4));
        if (number > frame_total) frame_total = number;
      }
    }
  }

  // Generate palette
  printf("Generating palette map...\n");
  int pal_w, pal_h, pal_n;
  unsigned char *palette_raw = stbi_load(palette_path.c_str(), &pal_w, &pal_h, &pal_n, 3);

  for (int i=0; i<PALETTE_SIZE; i++)
  {
    int index = i*3;
    unsigned char r = palette_raw[index];
    unsigned char g = palette_raw[index + 1];
    unsigned char b = palette_raw[index + 2];

    palette_buffer[i] = RGBtoRGB15(r, g, b);
  }

  printf("Writing palette map...\n");
  file_out.open(output_path + output_basename + "_pal.bin", std::ios::binary);
  file_out.write(reinterpret_cast<const char*>(palette_buffer), sizeof(palette_buffer));
  file_out.close();

  printf("Generating image map...\n");
  std::stringstream ss;
  file_out.open(output_path + output_basename + "_img.bin", std::ios::binary);

  // ########## PER FRAME ENCODER ##########
  // for (int i=0; i<frame_total; i++)
  // {
  //   ss.str("");
  //   ss << frame_path << "out_" << i+1 << ".bmp";

  //   int frame_w, frame_h, frame_n;
  //   unsigned char *frame_raw = stbi_load(ss.str().c_str(), &frame_w, &frame_h, &frame_n, 3);

  //   Frame_RAWtoRGB15(frame_raw, frame_buffer, palette_buffer);

  //   compress_size = Compress(type, frame_buffer, FRAME_SIZE, compress_buffer);
  //   if ((compress_size) > compress_size_biggest) compress_size_biggest = compress_size;
  //   compress_size_total += compress_size;

  //   file_out.write(reinterpret_cast<const char*>(&compress_size), sizeof(uint32_t));
  //   file_out.write(reinterpret_cast<const char*>(compress_buffer), compress_size);

  //   printf("write %i bytes (%i/%i)\n", compress_size, i+1, frame_total);
  // }
  // ########################################

  // Calculate the number of complete iterations needed
  int numIterations = frame_total / CHUNK_SIZE;

  // Calculate the number of remaining items
  int remainingItems = frame_total % CHUNK_SIZE;

  for (int i=0; i<numIterations; i++)
  {

    for (int i_chunk=0; i_chunk<CHUNK_SIZE; i_chunk++) {
      // Calculate the index of the current item within the total set
      int itemIndex = i * CHUNK_SIZE + i_chunk;

      ss.str("");
      ss << frame_path << "out_" << itemIndex+1 << ".bmp";

      // printf("loading %s\n", ss.str().c_str());

      int frame_w, frame_h, frame_n;
      unsigned char *frame_raw = stbi_load(ss.str().c_str(), &frame_w, &frame_h, &frame_n, 3);

      Frame_RAWtoRGB15(frame_raw, &frame_buffer[i_chunk*FRAME_SIZE], palette_buffer);
    }
    
    compress_size = Compress(type, frame_buffer, FRAME_SIZE*CHUNK_SIZE, compress_buffer);
    if ((compress_size) > compress_size_biggest) compress_size_biggest = compress_size;
    compress_size_total += compress_size;

    file_out.write(reinterpret_cast<const char*>(&compress_size), sizeof(uint32_t));
    file_out.write(reinterpret_cast<const char*>(compress_buffer), compress_size);
    // file_out.put('\0');

    printf("write chunk %i bytes %i/%i\n", compress_size, i, numIterations-1);
  }

  // Handle remaining frame
  if (remainingItems > 0) {
    memset(frame_buffer, 0, FRAME_SIZE*CHUNK_SIZE);

    int i_chunk = 0;
    for (int k = frame_total-remainingItems; k<frame_total; k++) {
      ss.str("");
      ss << frame_path << "out_" << k+1 << ".bmp";

      // printf("loading %s\n", ss.str().c_str());

      int frame_w, frame_h, frame_n;
      unsigned char *frame_raw = stbi_load(ss.str().c_str(), &frame_w, &frame_h, &frame_n, 3);

      Frame_RAWtoRGB15(frame_raw, &frame_buffer[i_chunk*FRAME_SIZE], palette_buffer);
      i_chunk++;
    }

    compress_size = Compress(type, frame_buffer, FRAME_SIZE*CHUNK_SIZE, compress_buffer);
    if ((compress_size) > compress_size_biggest) compress_size_biggest = compress_size;
    compress_size_total += compress_size;

    file_out.write(reinterpret_cast<const char*>(&compress_size), sizeof(uint32_t));
    file_out.write(reinterpret_cast<const char*>(compress_buffer), compress_size);
    // file_out.put('\0');

    printf("write remaining chunk %i bytes\n", compress_size);
  }

  file_out.close();

  // Print
  printf("Total Frame: %i\n", frame_total);
  printf("Biggest    : %i\n", compress_size_biggest);
  printf("Total Size : %1.fmb\n", static_cast<double>(compress_size_total) / (1024*1024));
}

int main()
{
  // Convert(ASSETS_PATH "out1/", ASSETS_PATH "palette1.png", "../nitrofs/", "data_fastlz", 1);
  // Convert(ASSETS_PATH "out1/", ASSETS_PATH "palette1.png", "../nitrofs/", "data_rle", 2);
  Convert(ASSETS_PATH "out/", ASSETS_PATH "palette1.png", "../nitrofs/", "main", 3);
  // Convert(ASSETS_PATH "out2/", ASSETS_PATH "palette2.png", "../nitrofs/", "sub");
}

void Frame_RAWtoRGB15(unsigned char *frame, unsigned char *dest, uint16_t *palette_buffer)
{
  for (int rgb=0; rgb<FRAME_SIZE; rgb++)
  {
    int index = rgb*3;
    unsigned char r = frame[index];
    unsigned char g = frame[index + 1];
    unsigned char b = frame[index + 2];

    uint16_t palette_value = RGBtoRGB15(r, g, b);
    int palette_index = -1;

    for (int i=0; i<256; i++)
    {
      if (palette_buffer[i] == palette_value) {
        palette_index = i;
        break;
      }
    }

    if (palette_index == -1) dest[rgb] = 0;
    else dest[rgb] = palette_index;
  }
}