#include <destrum/Graphics/imageLoader.h>

#define STB_IMAGE_IMPLEMENTATION
#include <stb_image.h>

#define TINYEXR_IMPLEMENTATION
#include "tinyexr.h"

#include <cstdio>   // fprintf
#include <cstdlib>  // free

#include "spdlog/spdlog.h"

#include <destrum/Assets/TextureFile.h>

ImageData::~ImageData()
{
    freePixels();
}

void ImageData::freePixels()
{
    if (shouldSTBFree) {
        if (pixels)    stbi_image_free(pixels);
        if (hdrPixels) stbi_image_free(hdrPixels);
        spdlog::debug("Freed STB image memory");
    }
    if (shouldEXRFree) {
        if (hdrPixels) free(hdrPixels);
        spdlog::debug("Freed TinyEXR image memory");
    }

    cookedPixels.clear();
    pixels = nullptr;
    hdrPixels = nullptr;
    shouldSTBFree = false;
    shouldEXRFree = false;
}

namespace util
{
    static bool isExrExt(const std::filesystem::path& p)
    {
        auto ext = p.extension().string();
        for (auto& c : ext) c = char(::tolower(c));
        return ext == ".exr";
    }

    static bool isCookedTextureExt(const std::filesystem::path& p)
    {
        auto ext = p.extension().string();
        for (auto& c : ext) c = char(::tolower(c));
        return ext == TextureFile::Extension;
    }

    static VkFormat toVkFormat(TextureFile::PixelFormat format, TextureIntent intent)
    {
        switch (format) {
            case TextureFile::PixelFormat::RGBA16Float: return VK_FORMAT_R16G16B16A16_SFLOAT;
            case TextureFile::PixelFormat::RGBA8:
                return intent == TextureIntent::ColorSrgb ? VK_FORMAT_R8G8B8A8_SRGB : VK_FORMAT_R8G8B8A8_UNORM;
        }
        return VK_FORMAT_UNDEFINED;
    }

    // TheChef writes <image>.dtex next to the copied source, so every caller that has the source path gets it
    static std::filesystem::path findCookedSibling(const std::filesystem::path& p)
    {
        std::filesystem::path cooked = p;
        cooked += TextureFile::Extension;

        std::error_code error;
        if (!std::filesystem::is_regular_file(cooked, error)) {
            return {};
        }
        const auto sourceTime = std::filesystem::last_write_time(p, error);
        if (!error && sourceTime > std::filesystem::last_write_time(cooked, error)) {
            return {};
        }
        return cooked;
    }

    ImageData loadImage(const std::filesystem::path& p, TextureIntent intent)
    {
        ImageData data;

        if (!isCookedTextureExt(p)) {
            if (const auto cooked = findCookedSibling(p); !cooked.empty()) {
                if (ImageData cookedData = loadImage(cooked, intent); cookedData.byteSize != 0) {
                    return cookedData;
                }
            }
        }

        if (isCookedTextureExt(p)) {
            try {
                TextureFile::Texture texture = TextureFile::ReadFromFile(p);

                data.cookedPixels = std::move(texture.pixels);
                data.pixels = data.cookedPixels.data();
                data.width = static_cast<int>(texture.width);
                data.height = static_cast<int>(texture.height);
                data.channels = 4;

                data.vkFormat = toVkFormat(texture.format, intent);
                data.byteSize = data.cookedPixels.size();
            } catch (const std::exception& exception) {
                spdlog::error("Failed to read cooked texture '{}': {}", p.string(), exception.what());
                return {};
            }
            return data;
        }

        if (isExrExt(p)) {
            float* out = nullptr;
            int w = 0, h = 0;
            const char* err = nullptr;

            const int ret = LoadEXR(&out, &w, &h, p.string().c_str(), &err);
            if (ret != TINYEXR_SUCCESS || !out || w <= 0 || h <= 0) {
                if (err) {
                    spdlog::error("TinyEXR error: {}", err);
                    FreeEXRErrorMessage(err);
                }
                return {};
            }

            data.hdr = true;
            data.hdrPixels = out;
            data.width = w;
            data.height = h;
            data.channels = 4;
            data.shouldEXRFree = true;

            data.vkFormat = VK_FORMAT_R32G32B32A32_SFLOAT;
            data.byteSize = std::size_t(w) * std::size_t(h) * 4 * sizeof(float);
            return data;
        }

        // ---------- STB HDR (.hdr etc) ----------
        data.shouldSTBFree = true;

        if (stbi_is_hdr(p.string().c_str())) {
            data.hdr = true;
            data.hdrPixels = stbi_loadf(
                p.string().c_str(),
                &data.width,
                &data.height,
                &data.comp,
                4
            );

            if (!data.hdrPixels || data.width <= 0 || data.height <= 0) {
                return {};
            }

            data.channels = 4;
            data.vkFormat = VK_FORMAT_R32G32B32A32_SFLOAT;
            data.byteSize = std::size_t(data.width) * std::size_t(data.height) * 4 * sizeof(float);
            return data;
        }

        // ---------- STB LDR (png/jpg/...) ----------
        data.pixels = stbi_load(
            p.string().c_str(),
            &data.width,
            &data.height,
            &data.channels,
            4
        );

        if (!data.pixels || data.width <= 0 || data.height <= 0) {
            return {};
        }

        data.channels = 4;

        // Choose color space for LDR based on intent
        switch (intent) {
            case TextureIntent::ColorSrgb:
                data.vkFormat = VK_FORMAT_R8G8B8A8_SRGB;
                break;
            case TextureIntent::DataLinear:
                data.vkFormat = VK_FORMAT_R8G8B8A8_UNORM;
                break;
        }

        data.byteSize = std::size_t(data.width) * std::size_t(data.height) * 4 * sizeof(unsigned char);
        return data;
    }

} // namespace util