#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);

#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 10

uint16_t palette_buffer[PALETTE_SIZE];

size_t chunk_counter = 0;
size_t chunk_current = 0;

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

uint32_t compress_size;
uint32_t compress_size_total;
size_t compress_size_biggest;
uint8_t compress_buffer[FRAME_SIZE*CHUNK_SIZE*2];

std::ofstream file_out;

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

int main()
{
  printf("Finding total frame...");
  // Find total frame
  for (const auto& entry : std::filesystem::directory_iterator(ASSETS_PATH "/out/")) {
    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(ASSETS_PATH "/palette.png", &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(ASSETS_PATH "/palette.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(ASSETS_PATH "/image.bin", std::ios::binary);

  // 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 << ASSETS_PATH "/out/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]);
    }
    
    compress_size = fastlz_compress_level(1, 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);

    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 << ASSETS_PATH "/out/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]);
      i_chunk++;
    }

    compress_size = fastlz_compress_level(1, 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);

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

void Frame_RAWtoRGB15(unsigned char *frame, unsigned char *dest)
{
  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;
  }
}