#include <cstdio>
#include <cstring>
#include <iostream>
#include <filesystem>
#include <fstream>
#include <sstream>
#include <string>

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

unsigned short RGB15(int r, int g, int b) {
  unsigned short rgb15 =
      ((r >> 3) << 10) |    // 5 bits for red
      ((g >> 3) << 5)  |    // 5 bits for green
      (b >> 3);             // 5 bits for blue

  // Set the alpha bit to 1
  rgb15 |= 1 << 15;

  return rgb15;
}

int main()
{
  std::string path_out = "../assets/out/";
  std::string path_palette = "../assets/palette.png";
  std::string path_resource = "../resource";
  std::string path_images = "../assets/out/out_";
  std::string path_test = "../assets/out/out_3667.bmp";

  int frame_total = -1;

  std::stringstream ss;
  std::ofstream out;

  // Find total frame
  for (const auto& entry : std::filesystem::directory_iterator(path_out)) {
    if (entry.is_regular_file()) {
      std::string fileName = entry.path().filename().string();

      if (fileName.substr(0, 4) == "out_" && fileName.substr(fileName.size() - 4) == ".bmp") {
        int number = std::stoi(fileName.substr(4, fileName.size() - 8));

        if (number > frame_total) {
          frame_total = number;
        }
      }
    }
  }

  // Generate palette
  uint16_t palette[256];
  int palette_w = -1;
  int palette_h = -1;
  int palette_n = -1;
  unsigned char *palette_data = stbi_load(path_palette.c_str(), &palette_w, &palette_h, &palette_n, 3);

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

    palette[i] = RGB15(r, g, b);
    // printf("%i> %i,%i,%i (%i)\n", i, r, g, b, palette[i]);
  }

  out.open(path_resource + "/palette.bin", std::ios::binary);
  out.write(reinterpret_cast<const char*>(palette), sizeof(palette));
  out.close();
  
  out.open(path_resource + "/image.bin", std::ios::binary);

  
  uint8_t batch_map[256*192*4];
  int batch_counter = 0;

  for (int frame_counter=1; frame_counter<=frame_total; frame_counter++)
  {
    ss.str("");
    ss << path_images << frame_counter << ".bmp";

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

    for (int i=0; i<256*192; i++)
    {

      int index = i*3;
      unsigned char r = frame_raw[index];
      unsigned char g = frame_raw[index + 1];
      unsigned char b = frame_raw[index + 2];

      uint16_t pixel_pal = RGB15(r, g, b);
      int pixel_index = -1;

      for (int i2=0; i2<256; i2++)
      {
        if (palette[i2] == pixel_pal) {
          pixel_index = i2;
          break;
        }
      }

      if (pixel_index == -1) printf("%i> %i palette not found!!!\n", i, pixel_index);
      else batch_map[i+ batch_counter*256*192] = pixel_index;
    }

    printf("%i) frame %i/%i (%1.f%s)> done\n", batch_counter, frame_counter, frame_total, ((float)frame_counter/frame_total)*100, "%");
    batch_counter++;

    if (batch_counter > 3)
    {
      uint8_t batch_map_compress[256*192*4*2];
      uint32_t size = fastlz_compress_level(1, batch_map, 256*192*4, batch_map_compress);

      out.write(reinterpret_cast<const char*>(&size), sizeof(size));
      out.write(reinterpret_cast<const char*>(&batch_map_compress), size);

      batch_counter = 0;
      printf("write batch %i\n", size);
    }
  }

  out.close();


  // for (int i=1; i<=frame_total/batch_size; i++)
  // {
  //   uint8_t batch[256*192*batch_size];
  //   uint8_t batch_compress[256*192*batch_size*2];
  //   uint32_t batch_compress_len = 0;
    

  //   for (batch_current=0; batch_current<batch_size; batch_current++)
  //   {
  //     if (frame_counter > frame_total) break;


  //     uint8_t current_map[256*192];
  //     ss.str("");
  //     ss << path_images << frame_counter << ".bmp";
  //     int current_w = -1;
  //     int current_h = -1;
  //     int current_n = -1;
  //     unsigned char *current_data = stbi_load(ss.str().c_str(), &current_w, &current_h, &current_n, 3);

  //     for (int i2=0; i2<256*192; i2++)
  //     {

  //       int index = i2*3;
  //       unsigned char r = current_data[index];
  //       unsigned char g = current_data[index + 1];
  //       unsigned char b = current_data[index + 2];


  //       uint16_t current_pal = RGB15(r, g, b);
  //       int current_index = -1;

  //       for (int i3=0; i3<256; i3++)
  //       {
  //         if (palette[i3] == current_pal) {
  //           current_index = i3;
  //           break;
  //         }
  //       }

  //       if (current_index == -1) printf("%i> %i palette not found!!!\n", i, current_index);
  //       else current_map[i2] = current_index;

  //     }

  //     memcpy(current_map, &batch[batch_current*256*192], 256*192);
  //     printf("frame %i/%i\n", frame_counter, frame_total);
  //     frame_counter++;
  //   }

  //   batch_compress_len = fastlz_compress_level(1, batch, 256*192*batch_size, batch_compress);
  //   out.write(reinterpret_cast<const char*>(&batch_compress_len), sizeof(batch_compress_len));
  //   out.write(reinterpret_cast<const char*>(&batch_compress), batch_compress_len);
  //   printf("batch done! (%i)\n", batch_compress_len);
  // }

  out.close();

  // uint8_t test[256*192];
  // uint8_t test_compress[(256*192)*2];
  // uint32_t test_compress_len = 0;
  // int test_w = -1;
  // int test_h = -1;
  // int test_n = -1;
  // unsigned char *test_data = stbi_load(path_test.c_str(), &test_w, &test_h, &test_n, 3);

  // for (int i=0; i<256*192; i++)
  // {
  //   int index = i*3;
  //   unsigned char r = test_data[index];
  //   unsigned char g = test_data[index + 1];
  //   unsigned char b = test_data[index + 2];

  //   uint16_t current_pal = RGB15(r, g, b);
  //   int current_index = -1;

  //   for (int i2=0; i2<256; i2++)
  //   {
  //     if (palette[i2] == current_pal) {
  //       current_index = i2;
  //       break;
  //     }
  //   }

  //   if (current_index == -1) printf("%i> %i\n", i, current_index);
  //   else test[i] = current_index;
  // }

  // test_compress_len = fastlz_compress_level(1, test, 256*192, test_compress);

  // printf("len: %i\n", test_compress_len);

  // out.open(path_resource + "/image.bin", std::ios::binary);
  // out.write(reinterpret_cast<const char*>(&test_compress_len), sizeof(test_compress_len));
  // out.write(reinterpret_cast<const char*>(test_compress), test_compress_len);
  // out.close();

  // Print
  printf("Total Frame: %i\n", frame_total);
}