//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"
// stb keeps the reason for its last failure until the next one, and a few
// of its failures give none: cleared before each decode, the reason after a
// failure is that failure's or null, never an older one's.
void CrafterAssetClearStbiFailure() {
    stbi__g_failure_reason = nullptr;
}
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, TGA, BMP, ...) to 8-bit
        // RGBA, or to normalized floats when TT is a floating-point component
        // type. A file that doesn't decode throws std::runtime_error, which
        // callers such as crafter-build catch to report the file; built
        // without exceptions (wasm), there is nothing to catch, and it is
        // Fatal.
        template <typename TT>
        static TextureAsset<T> LoadPNG(fs::path path) {
            std::filesystem::path abs = std::filesystem::absolute(path);

            // Read whole and decoded from memory, so that the size checked
            // before decoding is the size of the bytes that are decoded.
            std::error_code error;
            if (!fs::is_regular_file(abs, error)) {
                Fail(std::format("LoadPNG: {}: {}", abs.string(), error ? error.message() : "not a file"));
            }
            const std::uintmax_t size = fs::file_size(abs, error);
            if (error) Fail(std::format("LoadPNG: {}: {}", abs.string(), error.message()));
            if (size > kMaxEncoded) Fail(std::format("LoadPNG: {}: {} bytes is more than stb_image reads", abs.string(), size));
            std::vector<std::byte> encoded(static_cast<std::size_t>(size));
            std::ifstream file(abs, std::ios::binary);
            if (!file.read(reinterpret_cast<char*>(encoded.data()), static_cast<std::streamsize>(encoded.size()))) {
                Fail(std::format("LoadPNG: {}: can't read it", abs.string()));
            }
            std::expected<TextureAsset<T>, std::string> tex = LoadPNGFromMemory<TT>(encoded);
            if (!tex) Fail(std::format("LoadPNG: {}: {}", abs.string(), tex.error()));
            return std::move(*tex);
        }

        // 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 an exception or a Fatal: the caller is reading a larger
        // asset and can say which part was bad. One too large for a texture
        // is turned down from its header, before anything is decoded.
        template <typename TT>
        static std::expected<TextureAsset<T>, std::string> LoadPNGFromMemory(std::span<const std::byte> encoded) {
            const auto size = SizeFromMemory(encoded);
            if (!size) return std::unexpected(size.error());

            int sizeX;
            int sizeY;

            CrafterAssetClearStbiFailure();
            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(FailureReason());
            }
            if (!FitsSize(sizeX, sizeY)) {
                stbi_image_free(data);
                return std::unexpected(std::format("{}x{} decoded from a header of {}x{}", sizeX, sizeY, size->first, size->second));
            }
            return FromRGBA8<TT>(data, sizeX, sizeY);
        }

        // The width and height LoadPNGFromMemory would decode an image to,
        // read from its header alone, or why it wouldn't: an image it can't
        // read, or one too large for a texture.
        static std::expected<std::pair<std::uint32_t, std::uint32_t>, std::string> SizeFromMemory(std::span<const std::byte> encoded) {
            if (encoded.size() > kMaxEncoded) {
                return std::unexpected(std::format("{} bytes is more than stb_image reads", encoded.size()));
            }
            int sizeX;
            int sizeY;
            int components;
            CrafterAssetClearStbiFailure();
            if (!stbi_info_from_memory(reinterpret_cast<const stbi_uc*>(encoded.data()), static_cast<int>(encoded.size()),
                                       &sizeX, &sizeY, &components)) {
                return std::unexpected(FailureReason());
            }
            if (!FitsSize(sizeX, sizeY)) {
                return std::unexpected(std::format("{}x{} is more than a texture takes: 65535 a side, {} pixels in all", sizeX, sizeY, kMaxPixels));
            }
            return std::pair<std::uint32_t, std::uint32_t>(sizeX, sizeY);
        }

    private:
        // stb_image takes an encoded image's length as an int.
        static constexpr std::uintmax_t kMaxEncoded = std::numeric_limits<int>::max();
        // sizeX and sizeY hold 16 bits, and stb_image's arithmetic 32: its
        // 16-bit PNG decoding sizes a buffer of 8 bytes a pixel in an
        // unsigned int, which wraps round from 2^29 pixels. 2^28 is
        // 16384x16384, 1 GiB of RGBA8.
        static constexpr std::uint64_t kMaxPixels = std::uint64_t{1} << 28;

        static bool FitsSize(std::uint64_t sizeX, std::uint64_t sizeY) {
            return sizeX <= std::numeric_limits<std::uint16_t>::max() && sizeY <= std::numeric_limits<std::uint16_t>::max()
                && sizeX * sizeY <= kMaxPixels;
        }

        // stb_image's reason for the failure just now. CrafterAssetClearStbiFailure
        // before the call makes it this failure's, but some failures give
        // none.
        static std::string FailureReason() {
            const char* reason = stbi_failure_reason();
            return reason != nullptr ? reason : "corrupt image";
        }

        [[noreturn]] static void Fail(const std::string& message) {
#if __cpp_exceptions
            throw std::runtime_error(message);
#else
            Compression::Fatal(message);
#endif
        }

        // 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;
        }
    };
}