#include <iostream>
#include <string>
#include <filesystem>
#include <fstream>
#define STB_IMAGE_IMPLEMENTATION
#include "stb_image.h"
#include <vector>
#include <unordered_map>
#include <algorithm>
#include <iterator>
#include <cstdio>

using std::cout;
namespace fs = std::filesystem;

bool is_jpeg(const unsigned char* buf);
bool is_png(const unsigned char* buf);

struct Bucket {
    long r = 0;
    long g = 0;
    long b = 0;
    int count = 0;
};

int main(int argc, char *argv[]) {
    for (int i = 0; i <= argc - 1; i++) {
        std::string arg = argv[i];
        if (arg == "--help" || arg == "-h") {
            cout << 
                "ccolors - extract color palette from images\n\n"
                "Usage:\n"
                "  " << argv[0] << " <image>\n\n"
                "Options:\n"
                "  -h, --help     Show this help page\n\n"
                "Output:\n"
                "  Prints dominant colors as HEX and terminal swatches\n"
                "  Also writes palette.json\n";
            return 0;
        }
    }

    if (argc != 2) {
        cout << "Please specify a image file. " << argv[0] << " imagefile.jpg"
             << "\nTry \"" << argv[0] << " --help\" for more information.";
        return 1;
    }

    std::string imagefile = argv[1];

    if (!fs::is_regular_file(imagefile)) {
        std::cout << "File not found: " << imagefile << "\n";
        return 1;
    }

    std::ifstream file(imagefile, std::ios::binary);
    if (!file) {
        cout << "Cannot open file!\n";
        return 1;
    }

    // Read the file once; detect the format from the buffer and decode from
    // memory so the image is never read from disk a second time.
    std::vector<unsigned char> buffer((std::istreambuf_iterator<char>(file)),
                                      std::istreambuf_iterator<char>());
    if (buffer.size() < 8) {
        cout << "File is too small to be a valid image!\n";
        return 1;
    }
    const unsigned char* header = buffer.data();

    if (is_png(header)) {
        cout << "PNG file detected!\n";
    } else if (is_jpeg(header)) {
        cout << "JPEG file detected!\n";
    } else {
        cout << "Unknown image file format!\n";
        return 1;
    }

    int width, height, channels;

    unsigned char* data = stbi_load_from_memory(
        buffer.data(),
        static_cast<int>(buffer.size()),
        &width,
        &height,
        &channels,
        3 //forcing to RGB
        );

    if (!data) {
        cout << "Failed to load image!\n";
        return 1;
    }


    cout << "**********\n";
    cout << "Loaded image:\n";
    cout << "Width: " << width << "\n";
    cout << "Height: " << height << "\n";
    cout << "Channels: " << channels << "\n";
    cout << "**********\n";

    std::unordered_map<int, Bucket> hist;

    int step = 10;          //downsample rate
    long sampled = 0;
    for (int y = 0; y < height; y += step) {
        for (int x = 0; x < width; x += step) {
            int idx = (y * width + x) * 3;

            unsigned char r = data[idx];
            unsigned char g = data[idx + 1];
            unsigned char b = data[idx + 2];
            sampled++;

            int brightness = r + g + b;
            if (brightness < 60) continue;     //filtering out dark pixels

            int minc = std::min({r, g, b});
            int maxc = std::max({r, g, b});
            if (maxc - minc < 10) continue;    //filtering out low saturation
                                               //(gray) colors

            int rq = r / 16;
            int gq = g / 16;
            int bq = b / 16;

            int key = (rq << 8) | (gq << 4) | bq;

            auto& bucket = hist[key];
            bucket.r += r;
            bucket.g += g;
            bucket.b += b;
            bucket.count++;
        }
    }

    cout << "Sampled pixels: " << sampled << "\n";

    std::vector<Bucket> buckets;

    for (auto& kv : hist) {
        buckets.push_back(kv.second);
    }

    std::sort(buckets.begin(), buckets.end(),
            [](const Bucket& a, const Bucket& b) {
                return a.count > b.count;
            });

    int paletteSize = std::min(5, (int)buckets.size());

    std::ofstream json("palette.json");

    json << "{\n";
    json << "   \"image\": \"" << argv[1] << "\",\n";
    json << "   \"palette\": [\n";

    for (int i = 0; i < paletteSize; i++) {
        auto& b = buckets[i];
        int r = b.r / b.count;
        int g = b.g / b.count;
        int bcol = b.b / b.count;

        printf("\033[48;2;%d;%d;%dm  \033[0m #%02X%02X%02X\n",
        r, g, bcol,
        r, g, bcol);

        json << "   { "
             << "\"r\": " << r << ", "
             << "\"g\": " << g << ", "
             << "\"b\": " << bcol << ", "
             << "\"hex\": \"";

        char hexbuf[16];
        sprintf(hexbuf, "#%02X%02X%02X", r, g, bcol);

        json << hexbuf << "\" }";

        if (i != paletteSize - 1) json << ",";

        json << "\n";
    }

    json << "   ]\n";
    json << "}\n";
    json.close();

    stbi_image_free(data);

    return 0;
}

bool is_jpeg(const unsigned char* buf) {
    return buf[0] == 0xFF && buf[1] == 0xD8 && buf[2] == 0xFF;
}

bool is_png(const unsigned char* buf) {
    return buf[0] == 0x89 &&
           buf[1] == 0x50 &&
           buf[2] == 0x4E &&
           buf[3] == 0x47 &&
           buf[4] == 0x0D &&
           buf[5] == 0x0A &&
           buf[6] == 0x1A &&
           buf[7] == 0x0A;
}