#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <iostream>
#include <filesystem>
#include <fstream>
#include <sstream>
#include <stdint.h>
#include <string>
#include <vector>

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

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";

  std::vector<int> usingRLE = {177, 178, 179, 220, 231, 232, 237, 238, 239, 240, 241, 263, 264, 409, 681, 682, 735, 753, 754, 755, 756, 757, 758, 759, 760, 761, 762, 763, 764, 770, 771, 772, 773, 774, 775, 776, 777, 779, 780, 783, 784, 785, 786, 787, 788, 789, 790, 791, 798, 799, 800, 801, 802, 803, 804, 805, 807, 808, 809, 811, 812, 813, 814, 815, 816, 817, 818, 819, 821, 822, 823, 824, 825, 826, 827, 828, 830, 831, 832, 833, 834, 835, 836, 837, 838, 839, 840, 841, 842, 843, 844, 845, 846, 847, 848, 849, 850, 851, 852, 853, 854, 855, 856, 857, 858, 859, 860, 861, 862, 863, 864, 897, 898, 899, 900, 908, 909, 914, 915, 916, 917, 918, 919, 920, 921, 922, 923, 925, 926, 927, 928, 929, 930, 931, 934, 935, 937, 938, 939, 940, 941, 942, 943, 944, 945, 946, 947, 948, 949, 950, 951, 952, 953, 954, 1119, 1120, 1121, 1122, 1123, 1124, 1125, 1126, 1154, 1155, 1156, 1157, 1158, 1245, 1247, 1248, 1249, 1250, 1251, 1355, 1356, 1357, 1358, 1359, 1360, 1361, 1362, 1363, 1364, 1365, 1366, 1367, 1368, 1390, 1391, 1392, 1393, 1394, 1395, 1396, 1397, 1398, 1399, 1400, 1401, 1402, 1403, 1404, 1405, 1406, 1407, 1408, 1409, 1410, 1411, 1412, 1413, 1414, 1415, 1427, 1428, 1429, 1430, 1431, 1432, 1433, 1434, 1459, 1460, 1461, 1462, 1463, 1465, 1466, 1467, 1468, 1469, 1470, 1472, 1473, 1474, 1475, 1476, 1477, 1478, 1479, 1480, 1481, 1482, 1483, 1484, 1493, 1494, 1495, 1496, 1501, 1504, 1507, 1510, 1513, 1515, 1516, 1517, 1520, 1521, 1522, 1524, 1525, 1526, 1528, 1529, 1530, 1531, 1532, 1533, 1535, 1536, 1537, 1538, 1539, 1541, 1542, 1547, 1549, 1551, 1552, 1553, 1554, 1555, 1557, 1559, 1560, 1561, 1562, 1563, 1571, 1572, 1573, 1574, 1575, 1576, 1577, 1578, 1579, 1580, 1582, 1634, 1635, 1636, 1637, 1638, 1639, 1640, 1641, 1643, 1644, 1645, 1646, 1647, 1648, 1649, 1650, 1651, 1652, 1653, 1654, 1655, 1656, 1657, 1699, 1701, 1702, 1704, 1705, 1706, 1707, 1708, 1709, 1710, 1711, 1712, 1713, 1714, 1715, 1716, 1717, 1718, 1719, 1720, 1721, 1722, 1723, 1724, 1725, 1726, 1727, 1728, 1729, 1730, 1731, 1732, 1733, 1734, 1735, 1736, 1750, 1751, 1752, 1753, 1754, 1755, 1756, 1757, 1758, 1759, 1760, 1763, 1764, 1765, 1766, 1787, 1789, 1808, 1809, 1810, 1823, 1834, 1835, 1836, 1837, 1843, 1876, 1892, 1893, 1894, 1897, 1901, 1904, 1905, 1906, 1907, 1908, 1910, 1911, 1969, 2090, 2091, 2092, 2128, 2129, 2180, 2181, 2182, 2195, 2196, 2197, 2199, 2201, 2281, 2282, 2286, 2287, 2295, 2300, 2309, 2333, 2334, 2335, 2336, 2337, 2338, 2339, 2340, 2343, 2345, 2346, 2347, 2349, 2670, 2727, 2729, 2730, 2812, 2831, 2833, 2850, 2851, 3140, 3141, 3142, 3143, 3144, 3146, 3147, 3148, 3152, 3153, 3154, 3155, 3156, 3157, 3158, 3159, 3160, 3161, 3162, 3163, 3164, 3165, 3166, 3167, 3168, 3169, 3170, 3171, 3172, 3173, 3174, 3175, 3176, 3177, 3178, 3179, 3182, 3183, 3184, 3185, 3186, 3187, 3189, 3190, 3255};

  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;
  int batch_biggest = 0;
  int batch_num = 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++;

    // int size;
    // char *batch_map_compress = LZS_Fast(batch_map, 256*192, &size);

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

    if (batch_counter >= 4)
    {

      bool isRLE = false;

      for (auto y : usingRLE) 
        if (batch_num == y)
          isRLE = true;
      

      uint32_t size, size2;
      uint8_t *batch_map_compress = (uint8_t*)malloc(256*192*4*2);
      
      if (isRLE)
      {
        char *temp = RLE_Code(batch_map, 256*192*4, (int*)&size);
        size2 = (size << 1) | 0x01;
        memcpy(batch_map_compress, temp, size);
        free(temp);

        printf("rle    > ");
      }
      else
      {
        size = fastlz_compress_level(1, batch_map, 256*192*4, batch_map_compress);
        size2 = (size << 1);

        printf("fastlz > ");
      }

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

      batch_counter = 0;
      if ((size) > batch_biggest) batch_biggest = (size);
      printf("write batch %i %i/%i\n", size, batch_num, frame_total/4);

      free(batch_map_compress);
      batch_num++;
    }
  }

  out.close();

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