//SPDX-License-Identifier: LGPL-3.0-only
//SPDX-FileCopyrightText: Copyright (C) 2026 Catcrafts®

module;
#define STB_IMAGE_IMPLEMENTATION
#include "../lib/stb_image.h"
#define STB_IMAGE_RESIZE_IMPLEMENTATION
#include "../lib/stb_image_resize2.h"
export module Crafter.Asset:Texture;
import :Compression;
import std;
import Crafter.Math;
namespace fs = std::filesystem;

export namespace Crafter {
    enum class OpaqueType : std::uint8_t {
        FullyOpaque,      // All pixels have A of 255
        SemiOpaque,       // All pixels have A of 0 or 255 (no blending needed)
        Transparent       // Color blending is used
    };

    struct TextureAssetInfo {
        std::uint16_t sizeX;
        std::uint16_t sizeY;
        OpaqueType opaque;
    };

    // GDeflate-compressed counterpart of TextureAsset<T>::Save output.
    // Single-region blob: the pixel array as one stream.
    struct CompressedTextureAsset {
        std::uint16_t sizeX = 0;
        std::uint16_t sizeY = 0;
        OpaqueType opaque = OpaqueType::FullyOpaque;
        std::uint32_t pixelStride = 0;
        Compression::CompressedBlob blob;
    };

    namespace TextureAssetFormat {
        inline constexpr char magic[4] = {'C', 'G', 'D', 'T'};
        inline constexpr std::uint32_t version = 1;
    }

    inline CompressedTextureAsset LoadCompressedTexture(fs::path path) {
        std::ifstream file(path, std::ios::binary);
        char magic[4];
        file.read(magic, 4);
        if (std::memcmp(magic, TextureAssetFormat::magic, 4) != 0) {
            Compression::Fatal("LoadCompressedTexture: bad magic on " + path.string());
        }
        std::uint32_t version = 0;
        file.read(reinterpret_cast<char*>(&version), sizeof(version));
        if (version != TextureAssetFormat::version) {
            Compression::Fatal("LoadCompressedTexture: unsupported version on " + path.string());
        }
        CompressedTextureAsset out;
        file.read(reinterpret_cast<char*>(&out.sizeX), sizeof(out.sizeX));
        file.read(reinterpret_cast<char*>(&out.sizeY), sizeof(out.sizeY));
        file.read(reinterpret_cast<char*>(&out.opaque), sizeof(out.opaque));
        file.read(reinterpret_cast<char*>(&out.pixelStride), sizeof(out.pixelStride));
        out.blob = Compression::ReadBlob(file);
        return out;
    }

    template <typename T>
    struct TextureAsset {
        std::uint16_t sizeX;
        std::uint16_t sizeY;
        OpaqueType opaque;
        std::vector<T> pixels;

        void Save(fs::path path) {
            std::ofstream file(path, std::ios::binary);

            file.write(reinterpret_cast<char*>(&sizeX), sizeof(sizeX));
            file.write(reinterpret_cast<char*>(&sizeY), sizeof(sizeY));
            file.write(reinterpret_cast<char*>(&opaque), sizeof(opaque));
            file.write(reinterpret_cast<char*>(pixels.data()), pixels.size() * sizeof(T));
        }

        void SaveCompressed(fs::path path) const {
            std::array<std::span<const std::byte>, 1> streams = {
                std::as_bytes(std::span(pixels)),
            };
            Compression::CompressedBlob blob = Compression::CompressStreams(streams);

            std::ofstream file(path, std::ios::binary);
            file.write(TextureAssetFormat::magic, 4);
            std::uint32_t version = TextureAssetFormat::version;
            std::uint32_t stride = static_cast<std::uint32_t>(sizeof(T));
            file.write(reinterpret_cast<const char*>(&version), sizeof(version));
            file.write(reinterpret_cast<const char*>(&sizeX), sizeof(sizeX));
            file.write(reinterpret_cast<const char*>(&sizeY), sizeof(sizeY));
            file.write(reinterpret_cast<const char*>(&opaque), sizeof(opaque));
            file.write(reinterpret_cast<const char*>(&stride), sizeof(stride));
            Compression::WriteBlob(file, blob);
        }

        static TextureAsset<T> Load(fs::path path) {
            TextureAsset<T> tex;

            std::ifstream file(path, std::ios::binary);

            file.read(reinterpret_cast<char*>(&tex.sizeX), sizeof(tex.sizeX));
            file.read(reinterpret_cast<char*>(&tex.sizeY), sizeof(tex.sizeY));
            file.read(reinterpret_cast<char*>(&tex.opaque), sizeof(tex.opaque));
            tex.pixels.resize(tex.sizeX * tex.sizeY);

            file.read(reinterpret_cast<char*>(tex.pixels.data()), tex.sizeX * tex.sizeY * sizeof(T));
            return tex;
        }

        // Bilinear-resize the pixel buffer in-place to `newW × newH`.
        // Used to normalize albedos to a uniform size before stacking them
        // into a WebGPU texture_2d_array (which requires identical layer
        // dimensions). stb_image_resize2 handles RGBA8 directly.
        void Resize(std::uint16_t newW, std::uint16_t newH) requires (sizeof(T) == 4) {
            if (sizeX == newW && sizeY == newH) return;
            std::vector<T> out(static_cast<std::size_t>(newW) * newH);
            stbir_resize_uint8_linear(
                reinterpret_cast<const std::uint8_t*>(pixels.data()), sizeX, sizeY, 0,
                reinterpret_cast<std::uint8_t*>(out.data()),          newW,  newH,  0,
                STBIR_RGBA);
            pixels = std::move(out);
            sizeX = newW;
            sizeY = newH;
        }

        static TextureAssetInfo LoadInfo(fs::path path) {
            TextureAssetInfo info;

            std::ifstream file(path, std::ios::binary);

            file.read(reinterpret_cast<char*>(&info.sizeX), sizeof(info.sizeX));
            file.read(reinterpret_cast<char*>(&info.sizeY), sizeof(info.sizeY));
            file.read(reinterpret_cast<char*>(&info.opaque), sizeof(info.opaque));

            return info;
        }

        static void Load(fs::path path, T* pixels, std::uint16_t sizeX, std::uint16_t sizeY) {
            std::ifstream file(path, std::ios::binary);
            file.seekg(sizeof(std::uint16_t) + sizeof(std::uint16_t) + sizeof(std::uint8_t), std::ios::cur);
            file.read(reinterpret_cast<char*>(pixels), sizeX * sizeY * sizeof(T));
        }

        // Decodes anything stb_image reads (PNG, JPEG, ...) to 8-bit RGBA, or
        // to normalized floats when TT is a floating-point component type.
        template <typename TT>
        static TextureAsset<T> LoadPNG(fs::path path) {
            std::filesystem::path abs = std::filesystem::absolute(path);

            int sizeX;
            int sizeY;

            unsigned char* data = stbi_load(abs.string().c_str(), &sizeX, &sizeY, nullptr, 4);
            if (data == nullptr) {
                Compression::Fatal(std::format("LoadPNG: {}: {}", abs.string(), stbi_failure_reason()));
            }
            if (!FitsSize(sizeX, sizeY)) {
                stbi_image_free(data);
                Compression::Fatal(std::format("LoadPNG: {}: {}x{} is more than a texture's 65535 a side", abs.string(), sizeX, sizeY));
            }
            return FromRGBA8<TT>(data, sizeX, sizeY);
        }

        // LoadPNG from an image that is already in memory, e.g. one embedded
        // in a GLB. An image that doesn't decode comes back as an error value
        // rather than a Fatal: the caller is reading a larger asset and can
        // say which part was bad.
        template <typename TT>
        static std::expected<TextureAsset<T>, std::string> LoadPNGFromMemory(std::span<const std::byte> encoded) {
            int sizeX;
            int sizeY;

            unsigned char* data = stbi_load_from_memory(reinterpret_cast<const stbi_uc*>(encoded.data()),
                                                        static_cast<int>(encoded.size()), &sizeX, &sizeY, nullptr, 4);
            if (data == nullptr) {
                return std::unexpected(std::string(stbi_failure_reason()));
            }
            if (!FitsSize(sizeX, sizeY)) {
                stbi_image_free(data);
                return std::unexpected(std::format("{}x{} is more than a texture's 65535 a side", sizeX, sizeY));
            }
            return FromRGBA8<TT>(data, sizeX, sizeY);
        }

    private:
        // sizeX and sizeY hold 16 bits.
        static bool FitsSize(int sizeX, int sizeY) {
            return sizeX <= std::numeric_limits<std::uint16_t>::max() && sizeY <= std::numeric_limits<std::uint16_t>::max();
        }

        // Converts stb's RGBA8 buffer, classifies its opacity and frees it.
        template <typename TT>
        static TextureAsset<T> FromRGBA8(unsigned char* data, int sizeX, int sizeY) {
            TextureAsset<T> tex;

            tex.sizeX = sizeX;
            tex.sizeY = sizeY;

            tex.pixels.resize(tex.sizeX*tex.sizeY);

            tex.opaque = OpaqueType::FullyOpaque;

            if constexpr(
#ifndef __wasm__
                std::same_as<TT, _Float16> ||
#endif
                std::same_as<TT, float> || std::same_as<TT, double>) {
                for(std::uint32_t i = 0; i < sizeX*sizeY; i++) {
                    tex.pixels[i].r = TT(data[i*4])/255;
                    tex.pixels[i].g = TT(data[i*4+1])/255;
                    tex.pixels[i].b = TT(data[i*4+2])/255;
                    tex.pixels[i].a = TT(data[i*4+3])/255;
                }

                for(std::uint32_t i = 0; i < tex.sizeX* tex.sizeY; i++) {
                    if(tex.pixels[i].a != 1) {
                        tex.opaque = OpaqueType::SemiOpaque;
                        for(std::uint32_t i2 = i; i2 < tex.sizeX* tex.sizeY; i2++) {
                            if(tex.pixels[i2].a != 0 && tex.pixels[i2].a != 1) {
                                tex.opaque = OpaqueType::Transparent;
                                stbi_image_free(data);
                                return tex;
                            }
                        }
                        stbi_image_free(data);
                        return tex;
                    }
                }
            } else {
                std::memcpy(tex.pixels.data(), data, tex.sizeX * tex.sizeY * 4);

                for(std::uint32_t i = 0; i < tex.sizeX* tex.sizeY; i++) {
                    if(tex.pixels[i].a != 255) {
                        tex.opaque = OpaqueType::SemiOpaque;
                        for(std::uint32_t i2 = i; i2 < tex.sizeX* tex.sizeY; i2++) {
                            if(tex.pixels[i2].a != 0 && tex.pixels[i2].a != 255) {
                                tex.opaque = OpaqueType::Transparent;
                                stbi_image_free(data);
                                return tex;
                            }
                        }
                        stbi_image_free(data);
                        return tex;
                    }
                }
            }

            stbi_image_free(data);
            return tex;
        }
    };
}