#include "Artwork.h"
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <stdexcept>
#include <string>
#include <unordered_set>
#include <zip.h>
//#include "LayerNodes/LayerNode.h"
#include "App.h"
#include "LayerNodes/PaintLayerNode.h"
#include "ManagementContext/LayerManagmentContext.h"
#include "Version.h"
#include "json.hpp"
using json = nlohmann::json;
#define LONESHA256_IMPLEMENTATION
#include "lonesha256.h"

char* AddNullTerminator(void* data, size_t& size) {
    char* buffer = (char*)malloc(size + 1);
    memcpy(buffer, data, size);
    buffer[size] = '\0';
    size++;
    return buffer;
}

void CollectJsonStrings(const json& node, std::unordered_set<std::string>& referenced) { 
    if (node.is_string()) {
        referenced.insert(node.get<std::string>());
    } else if (node.is_array()) {
        for (const auto& element : node) {
            CollectJsonStrings(element, referenced);
        }
    } else if (node.is_object()) {
        for (auto& [key, value] : node.items()) {
            CollectJsonStrings(value, referenced);
        }
    }
}

void* Artwork::readZipFile(std::string path, std::string fileName, size_t* outSize) {
    zip_t* zipFile = zip_open(path.c_str(), 0, NULL);
    if (zipFile == NULL) {
        throw std::runtime_error("Failed to open zip file");
    }
    struct zip_file* file = zip_fopen(zipFile, fileName.c_str(), 0);
    if (file == NULL) {
        throw std::runtime_error("Failed to open file in zip: " + fileName);
    }
    zip_stat_t stat;
    if (zip_stat(zipFile, fileName.c_str(), 0, &stat) != 0) {
        throw std::runtime_error("Failed to get file stats");
    }
    *outSize = stat.size;
    void* buffer = malloc(*outSize);
    if (buffer == NULL) {
        throw std::runtime_error("Failed to allocate memory for file");
    }
    if ((size_t)zip_fread(file, buffer, *outSize) != *outSize) {
        free(buffer);
        throw std::runtime_error("Failed to read file");
    }
    zip_fclose(file);
    zip_close(zipFile);
    return buffer;
}

Artwork::Artwork(std::string path) {
    this->path = path;
    size_t size;
    char* buffer = (char*)readZipFile(path, "details.json", &size);
    buffer = AddNullTerminator(buffer, size);
    json data = json::parse(buffer);
    width = data["width"];
    height = data["height"];
    free(buffer);
}

Artwork::~Artwork() {
    LayerManagementContext::Clear();
}

Artwork* Artwork::NewArtwork(int width, int height, std::string path) {
    // construct a minimial grdn project
    zip_t* zipFile = zip_open(path.c_str(), ZIP_CREATE | ZIP_TRUNCATE, NULL);
    if (zipFile == NULL) {
        throw std::runtime_error("Failed to create zip file");
    }
    // write details.json
    json details;
    details["width"] = width;
    details["height"] = height;
    details["name"] = "untitled";
    details["description"] = "";
    details["version"] = "1.0";
    // json.dump is valid, zip file is culprit here.
    std::string dump = details.dump();
    char* buf = static_cast<char*>(malloc(dump.size()));
    memcpy(buf, dump.c_str(), dump.size());
    zip_source_t *source = zip_source_buffer(zipFile, buf, dump.size(), 0);
    if (!source) {
        free(buf);
        zip_close(zipFile);
        throw std::runtime_error("Failed to create zip source");
    }
    zip_file_add(zipFile, "details.json", source, ZIP_FL_ENC_GUESS);
    zip_close(zipFile);
    free(buf);
    // write layers.json (minimum)
    zipFile = zip_open(path.c_str(), 0, NULL);
    std::string layersDump = json::array().dump();
    zip_source_t *layersSource = zip_source_buffer(zipFile, layersDump.c_str(), layersDump.size(), 0);
    if (layersSource) {
        zip_file_add(zipFile, "layers.json", layersSource, ZIP_FL_ENC_GUESS);
    }
    zip_close(zipFile);
    // create artwork from file
    return new Artwork(path);
}

void Artwork::LoadProject() {
    assert(width < 65535 && height < 65535 && "width or height unreasonable and likely unsuported");
    LayerManagementContext::InitCanvas(width, height);
    // for every layer in layers.json 
    size_t size;
    char* buffer = (char*)readZipFile(path, "layers.json", &size);
    buffer = AddNullTerminator(buffer, size);
    json data = json::parse(buffer);
    
    if (!data.empty()) {
        // Clear default canvas layers
        for (auto* l : LayerManagementContext::Layers) {
            delete l;
        }
        LayerManagementContext::Layers.clear();
        LayerManagementContext::Registry.clear();
        LayerManagementContext::SelectedLayer = nullptr;
        
        uint32_t maxId = 0;
        
        // reconstruct nodes and load textures
        for (auto& layer : data) {
            PaintLayerNode* paintLayer = new PaintLayerNode();
            
            uint32_t id = layer["id"];
            paintLayer->ID = id;
            if (id > maxId) {
                maxId = id;
            }
            
            paintLayer->Name = layer["name"];
            paintLayer->Opacity = layer["opacity"];
            paintLayer->BlendMode = (Blendmode)layer["blendingMode"];
            
            SDL_GPUTextureCreateInfo creationInfo = {};
            creationInfo.format = SDL_GPU_TEXTUREFORMAT_R8G8B8A8_UNORM;
            creationInfo.width = layer["width"];
            creationInfo.height = layer["height"];
            creationInfo.layer_count_or_depth = 1;
            creationInfo.num_levels = 1;
            creationInfo.sample_count = SDL_GPU_SAMPLECOUNT_1;
            creationInfo.usage = SDL_GPU_TEXTUREUSAGE_COLOR_TARGET | SDL_GPU_TEXTUREUSAGE_SAMPLER;
            paintLayer->Target = new Texture(creationInfo);
            
            std::string cid = layer["source"];
            if (!cid.empty()) {
                size_t qoiSize = 0;
                try {
                    void* qoiBuffer = readZipFile(path, "storage/" + cid, &qoiSize);
                    if (qoiBuffer) {
                        paintLayer->Target->UploadAsQoI(qoiBuffer, qoiSize);
                        free(qoiBuffer);
                    } else {
                        paintLayer->Target->Clear({ 0.0f, 0.0f, 0.0f, 0.0f });
                    }
                } catch (const std::exception& e) {
                    std::printf("Failed to load texture for layer %d: %s\n", id, e.what());
                    paintLayer->Target->Clear({ 0.0f, 0.0f, 0.0f, 0.0f });
                }
            } else {
                paintLayer->Target->Clear({ 0.0f, 0.0f, 0.0f, 0.0f });
            }
            
            LayerManagementContext::Layers.push_back(paintLayer);
            LayerManagementContext::Registry[id] = paintLayer;
        }
        
        // reconstruct parents
        int i = 0;
        for (auto& layer : data) {
            LayerNode* paintLayer = LayerManagementContext::Layers[i++];
            std::string parentIdStr = layer["parentId"];
            if (parentIdStr != "null" && !parentIdStr.empty()) {
                try {
                    uint32_t parentId = std::stoul(parentIdStr);
                    if (LayerManagementContext::Registry.find(parentId) != LayerManagementContext::Registry.end()) {
                        paintLayer->Parent = LayerManagementContext::Registry[parentId];
                    }
                } catch (...) {
                    // Ignore parsing errors
                }
            }
        }
        
        LayerNode::SetIdCounter(maxId);
        
        // Select the first layer as default selected layer
        if (!LayerManagementContext::Layers.empty()) {
            LayerManagementContext::SelectLayer(LayerManagementContext::Layers.front());
        }
    }
    
    free(buffer);
    lastLoadTime = std::chrono::system_clock::now();
}

void Artwork::SaveProject() { 
    std::printf("Saving artwork %s ...\n", path.c_str());
    // write layers.json
    zip_t* zipFile = zip_open(path.c_str(), 0, NULL);
    json layers = json::array();
    for (LayerNode *layer : LayerManagementContext::Layers) {
        json layerJson;
        layerJson["version"] = "1.0"; // version information is for each layer, considering other apps may support different versions, might need namespaces
        layerJson["id"] = layer->ID;
        layerJson["parentId"] = (layer->Parent) ? std::to_string(LayerManagementContext::ToId(layer->Parent)) : "null";
        layerJson["name"] = layer->Name;
        // for now gradien saves the entire canvas layer regardless of how much of it is used, TODO: optimize saved contents
        layerJson["width"] = LayerManagementContext::FinalWidth;
        layerJson["height"] = LayerManagementContext::FinalHeight;
        layerJson["horizontalOffset"] = 0;
        layerJson["verticalOffset"] = 0;
        // write source
        layerJson["source"] = "";
        // try cast to PaintLayerNode
        if (PaintLayerNode* paintLayer = dynamic_cast<PaintLayerNode*>(layer)) {
            size_t encodedSize;
            void* encodedData = paintLayer->Target->DownloadAsQoI(encodedSize);
            // cid
            auto cid = GenerateContentIdentifier(encodedData, encodedSize);
            layerJson["source"] = cid;
            
            zip_source_t *source = zip_source_buffer(zipFile, encodedData, encodedSize, 1);
            if (!source) {
                free(encodedData);
                zip_close(zipFile);
                throw std::runtime_error("Failed to create zip source");
            }
            zip_file_add(zipFile, ("storage/" + cid).c_str(), source, ZIP_FL_ENC_GUESS);
        }
        layerJson["opacity"] = layer->Opacity;
        layerJson["blendingMode"] = layer->BlendMode;
        layerJson["flags"] = {}; // empty array for now 
        layers.push_back(layerJson);
    }
    std::string layersEncoded = layers.dump();
    zip_source_t *layersSource = zip_source_buffer(zipFile, layersEncoded.c_str(), layersEncoded.size(), 0);
    zip_file_add(zipFile, "layers.json", layersSource, ZIP_FL_OVERWRITE);
    // write history.json
    std::printf("Appending history...\n");
    // get prev
    size_t prevHistorySize = 0;
    char* prevHistoryBuffer; 
    try {
        prevHistoryBuffer = (char*)readZipFile(path, "history.json", &prevHistorySize);
        prevHistoryBuffer = AddNullTerminator(prevHistoryBuffer, prevHistorySize);
    } catch (std::runtime_error& e) {
        // empty history
        prevHistoryBuffer = (char*)malloc(3);
        prevHistorySize = 3;
        memcpy(prevHistoryBuffer, "[]\0", prevHistorySize);
    }
    json history = json::parse(prevHistoryBuffer);
    free(prevHistoryBuffer);
    json session = json::object();
    session["events"] = App::History->GetEvents();
    session["version"] = "1.0";
    session["appName"] = APP_NAME;
    session["appVersion"] = SEMVER_STRING;
    session["beginTime"] =std::format("{:%Y-%m-%dT%H:%M:%SZ}", lastLoadTime);
    session["endTime"] =std::format("{:%Y-%m-%dT%H:%M:%SZ}", std::chrono::system_clock::now());
    history.push_back(session);
    std::string historyEncoded = history.dump();
    zip_source_t *historySource = zip_source_buffer(zipFile, historyEncoded.c_str(), historyEncoded.size(), 0);
    zip_file_add(zipFile, "history.json", historySource, ZIP_FL_OVERWRITE);
    zip_close(zipFile);
    std::printf("Clearing unused blobs...\n");
    RemoveUnusedBlobs();
    std::printf("Saved!\n");
}

void Artwork::RemoveUnusedBlobs() {
    // enumirate every blob and scans every json (including ones not in spec) for blob reference by cid (ie b....), then removes unused ones
    std::unordered_set<std::string> inStorage;
    std::unordered_set<std::string> referenced;
    std::unordered_set<std::string> jsonPaths;

    // add all files in storage
    zip_t* zipFile = zip_open(path.c_str(), 0, NULL);
    if (zipFile == NULL) {
        std::printf("Failed to open zip file: %s\n", path.c_str());
        return;
    }
    int n = zip_get_num_entries(zipFile, 0);
    for (int i = 0; i < n; i++) {
        const char* name = zip_get_name(zipFile, i, 0);
        if (name) {
            std::string nameStr(name);
            if (nameStr.starts_with("storage/")) {
                inStorage.insert(nameStr);
            }
            if (nameStr.ends_with(".json")) {
                jsonPaths.insert(nameStr);
            }
        }
    }

    // scan every json for blobs
    for (const auto& jsonPath : jsonPaths) {
        size_t len;
        void* data = readZipFile(path, jsonPath, &len);
        json j = json::parse(std::string((char*)data, len));
        CollectJsonStrings(j, referenced);
        free(data);
    }

    for (const auto& cid : referenced) {
        std::string fullPath = "storage/" + cid;
        if (inStorage.count(fullPath)) {
            inStorage.erase(fullPath);
        }
    }
    for (const auto& entry : inStorage) {
        zip_int64_t index = zip_name_locate(zipFile, entry.c_str(), 0);
        if (index >= 0) {
            zip_delete(zipFile, (zip_uint64_t)index);
        }
    }
    zip_close(zipFile);
}

std::string Artwork::Base32Encode(const void* data, size_t size) {
    // https://datatracker.ietf.org/doc/html/rfc4648#section-6
    const char* const alphabet = "abcdefghijklmnopqrstuvwxyz234567";
    std::string out;
    out.reserve(((size + 4) / 5) * 8);
    size_t i = 0;
    const uint8_t* bytes = (const uint8_t*)data;
    while (i < size) {
        uint64_t chunk = 0;
        int available = 0;
        for (int j = 0; j < 5 && i < size; ++j, ++i) {
            chunk = (chunk << 8) | bytes[i];
            available += 8;
        }
        chunk <<= (40 - available);
        for (int j = 35; j >= 0; j -= 5) {
            if (available > 0) {
                out += alphabet[(chunk >> j) & 0x1F];
                available -= 5;
            }
        }
    }
    return out;
}

std::string Artwork::GenerateContentIdentifier(void* data, size_t size) {
    // as per https://dasl.ing/cid.html
    unsigned char hash[32];
    lonesha256(hash, (const unsigned char*)data, size);

    // version 1, codec raw, hashtype sha256, hashsize 32 bytes
    const uint8_t header[] = { 0x01, 0x55, 0x12, 0x20 };
    uint8_t cidBytes[sizeof(header) + 32];
    memcpy(cidBytes, header, sizeof(header));
    memcpy(cidBytes + sizeof(header), hash, sizeof(hash));

    return "b" + Base32Encode(cidBytes, sizeof(cidBytes));
}