#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <cstdio>
#include <cstring>
#include <filesystem>
#include <fstream>
#include <sys/types.h>
#include <vector>

#define STB_IMAGE_IMPLEMENTATION
#include "stb_image.h"

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_PER_CHUNK 2
#define FRAMERATE 30

size_t GetTotalFrame(std::string path)
{
  size_t total = 0;

  for (const auto& entry : std::filesystem::directory_iterator(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 > total) total = number;
      }
    }
  }

  return total;
}

void GetFrameDimension(std::string path, int *w, int *h)
{
  int n;
  unsigned char *raw = stbi_load(path.c_str(), w, h, &n, 3);
  stbi_image_free(raw);
}

void ConvertPalette(uint16_t *buffer_palette, std::string path)
{
  int pal_w, pal_h, pal_n;
  unsigned char *palette_raw = stbi_load(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];

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

  stbi_image_free(palette_raw);
}

void ConvertFrametoPalIndex(unsigned char *buffer_frame, size_t buffer_size, uint16_t *buffer_palette, unsigned char *dest)
{
  for (int rgb=0; rgb<buffer_size; rgb++)
  {
    int index = rgb*3;
    unsigned char r = buffer_frame[index];
    unsigned char g = buffer_frame[index + 1];
    unsigned char b = buffer_frame[index + 2];

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

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

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

int Compress(uint8_t type, uint8_t *frame_buffer, uint32_t frame_buffer_len, uint8_t *compress_buffer)
{
  // if (time_list.at(chunknum) > time_avg) type = 4;

  switch (type) {
    case 1: // 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;
      }

    case 2: // raw
      {
        memcpy(compress_buffer, frame_buffer, frame_buffer_len);
        return frame_buffer_len;
        break;
      }
  }

  return 0;
}

void Convert(
  std::string path_video1,
  std::string path_palette1,
  std::string path_video2,
  std::string path_palette2,
  std::string path_output)
{
  uint16_t metadata_width1;
  uint16_t metadata_width2;
  uint16_t metadata_height1;
  uint16_t metadata_height2;
  size_t metadata_total_frame1;
  size_t metadata_total_frame2;

  uint16_t *buffer_palette1 = new uint16_t[PALETTE_SIZE];
  uint16_t *buffer_palette2 = new uint16_t[PALETTE_SIZE];

  uint8_t *buffer_frame;
  uint8_t *buffer_compress;

  size_t size_video1_frame;
  size_t size_video2_frame;
  size_t size_buffer_frame;
  size_t size_compress = 0;
  size_t size_compress_total = 0;
  size_t size_compress_max = 0;
  size_t size_compress_min = 99999999;

  std::ofstream file_out;
  file_out.open(path_output, std::ios::binary);

  if (!file_out) {
    printf("fail to open output %s\n", path_output.c_str());
    return;
  }
  
  // ---
  {
    printf("Finding total frame...\n");
    metadata_total_frame1 = GetTotalFrame(path_video1);
    metadata_total_frame2 = GetTotalFrame(path_video2);
  }

  // ---
  {
    printf("Finding video dimension...\n");
    int w, h;
    GetFrameDimension(path_video1 + "out_1.bmp", &w, &h);
    metadata_width1 = w;
    metadata_height1 = h;

    GetFrameDimension(path_video2 + "out_1.bmp", &w, &h);
    metadata_width2 = w;
    metadata_height2 = h;

    size_video1_frame = metadata_width1*metadata_height1;
    size_video2_frame = metadata_width2*metadata_height2;
  }

  printf("video1 dimension: %ix%i\n", metadata_width1, metadata_height1);
  printf("video2 dimension: %ix%i\n", metadata_width2, metadata_height2);
  
  // ---
  {
    size_buffer_frame = size_video1_frame + size_video2_frame;
    buffer_frame = new uint8_t[size_buffer_frame*FRAME_PER_CHUNK];
    buffer_compress = new uint8_t[size_buffer_frame*FRAME_PER_CHUNK];
  }

  // ---
  {
    printf("Generating palette map...\n");
    ConvertPalette(buffer_palette1, path_palette1);
    ConvertPalette(buffer_palette2, path_palette2);
  }

  // ---
  {
    printf("Writing metadata...\n");
  
    uint8_t framerate = FRAMERATE;
    uint8_t frame_per_chunk = FRAME_PER_CHUNK;

    file_out.write(reinterpret_cast<const char*>(&framerate), sizeof(uint8_t));
    file_out.write(reinterpret_cast<const char*>(&frame_per_chunk), sizeof(uint8_t));
    file_out.write(reinterpret_cast<const char*>(&metadata_total_frame1), sizeof(uint32_t));
    file_out.write(reinterpret_cast<const char*>(&metadata_total_frame2), sizeof(uint32_t));
    file_out.write(reinterpret_cast<const char*>(&metadata_width1), sizeof(uint16_t));
    file_out.write(reinterpret_cast<const char*>(&metadata_height1), sizeof(uint16_t));
    file_out.write(reinterpret_cast<const char*>(&metadata_width2), sizeof(uint16_t));
    file_out.write(reinterpret_cast<const char*>(&metadata_height2), sizeof(uint16_t));
  }

  // ---
  {
    printf("Writing palette map...\n");
    file_out.write(reinterpret_cast<const char*>(buffer_palette1), PALETTE_SIZE*sizeof(uint16_t));
    file_out.write(reinterpret_cast<const char*>(buffer_palette2), PALETTE_SIZE*sizeof(uint16_t));
  }

  // ---

  printf("Generating image map...\n");
  std::stringstream ss;

  // Calculate the number of complete iterations needed
  int numIterations = std::max(metadata_total_frame1, metadata_total_frame2) / FRAME_PER_CHUNK;

  // Calculate the number of remaining items
  int remainingItems = std::max(metadata_total_frame1, metadata_total_frame2) % FRAME_PER_CHUNK;

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

    memset(buffer_frame, 0, size_buffer_frame*FRAME_PER_CHUNK);

    for (int i_chunk=0; i_chunk<FRAME_PER_CHUNK; i_chunk++) {
      int itemIndex = i * FRAME_PER_CHUNK + i_chunk + 1;

      int frame_w, frame_h, frame_n;

      if (itemIndex <= metadata_total_frame1)
      {
        ss.str("");
        ss << path_video1 << "out_" << itemIndex << ".bmp";
        unsigned char *frame_raw = stbi_load(ss.str().c_str(), &frame_w, &frame_h, &frame_n, 3);
        ConvertFrametoPalIndex(frame_raw, frame_w*frame_h, buffer_palette1, &buffer_frame[i_chunk*size_buffer_frame]);
        stbi_image_free(frame_raw);
      }

      if (itemIndex <= metadata_total_frame2)
      {
        ss.str("");
        ss << path_video2 << "out_" << itemIndex << ".bmp";
        unsigned char *frame_raw = stbi_load(ss.str().c_str(), &frame_w, &frame_h, &frame_n, 3);
        ConvertFrametoPalIndex(frame_raw, frame_w*frame_h, buffer_palette2, &buffer_frame[i_chunk*size_buffer_frame+size_video1_frame]);
        stbi_image_free(frame_raw);
      }
    }
    
    size_compress = Compress(1, buffer_frame, size_buffer_frame*FRAME_PER_CHUNK, buffer_compress);
    size_compress_max = std::max(size_compress, size_compress_max);
    size_compress_min = std::min(size_compress, size_compress_min);
    size_compress_total += size_compress;

    file_out.write(reinterpret_cast<const char*>(&size_compress), sizeof(uint32_t));
    file_out.write(reinterpret_cast<const char*>(buffer_compress), size_compress);

    printf("write %i bytes (%i/%i)\n", size_compress, i+1, numIterations);
  }

  // Handle remaining frame
  if (remainingItems > 0) {
    memset(buffer_frame, 0, size_buffer_frame*FRAME_PER_CHUNK);

    int i_chunk = 0;
    int frame_total = std::max(metadata_total_frame1, metadata_total_frame2);
    for (int k = frame_total-remainingItems+1; k<=frame_total; k++) {
      int frame_w, frame_h, frame_n;

      if (k <= metadata_total_frame1)
      {
        ss.str("");
        ss << path_video1 << "out_" << k << ".bmp";
        unsigned char *frame_raw = stbi_load(ss.str().c_str(), &frame_w, &frame_h, &frame_n, 3);
        ConvertFrametoPalIndex(frame_raw, frame_w*frame_h, buffer_palette1, &buffer_frame[i_chunk*size_buffer_frame]);
        stbi_image_free(frame_raw);
      }

      if (k <= metadata_total_frame2)
      {
        ss.str("");
        ss << path_video2 << "out_" << k << ".bmp";
        unsigned char *frame_raw = stbi_load(ss.str().c_str(), &frame_w, &frame_h, &frame_n, 3);
        ConvertFrametoPalIndex(frame_raw, frame_w*frame_h, buffer_palette2, &buffer_frame[i_chunk*size_buffer_frame+size_video1_frame]);
        stbi_image_free(frame_raw);
      }

      i_chunk++;
    }

    size_compress = Compress(1, buffer_frame, size_buffer_frame*FRAME_PER_CHUNK, buffer_compress);
    size_compress_max = std::max(size_compress, size_compress_max);
    size_compress_min = std::min(size_compress, size_compress_min);
    size_compress_total += size_compress;

    file_out.write(reinterpret_cast<const char*>(&size_compress), sizeof(uint32_t));
    file_out.write(reinterpret_cast<const char*>(buffer_compress), size_compress);

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

  file_out.close();

  // Print
  printf("\n");
  printf("fps: %i\n", FRAMERATE);
  printf("frame per chunk: %i\n", FRAME_PER_CHUNK);
  printf("video1 total frame: %i\n", metadata_total_frame1);
  printf("video2 total frame: %i\n", metadata_total_frame2);
  printf("video1 dimension: %ix%i\n", metadata_width1, metadata_height1);
  printf("video2 dimension: %ix%i\n", metadata_width2, metadata_height2);
  printf("compress size max: %i\n", size_compress_max);
  printf("compress size min: %i\n", size_compress_min);
  printf("size total : %1.fmb\n", static_cast<double>(size_compress_total) / (1024*1024));

  delete[] buffer_palette1;
  delete[] buffer_palette2;
  delete[] buffer_frame;
  delete[] buffer_compress;
}

int main()
{
  Convert(
    ASSETS_PATH "out/",
    ASSETS_PATH "palette1.png",
    ASSETS_PATH "out2/",
    ASSETS_PATH "palette2.png",
    "../nitrofs/video.sillyvideo");
}