#include "iwa/io/bitmap.hpp"

#include <cstring>
#include <filesystem>
#include <mijin/detect.hpp>

#if MIJIN_COMPILER == MIJIN_COMPILER_GCC || MIJIN_COMPILER == MIJIN_COMPILER_CLANG
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wunused-but-set-variable"
#pragma GCC diagnostic ignored "-Wmissing-field-initializers"
#endif // MIJIN_COMPILER == MIJIN_COMPILER_GCC || MIJIN_COMPILER == MIJIN_COMPILER_CLANG
#define STB_IMAGE_IMPLEMENTATION
#define STB_IMAGE_WRITE_IMPLEMENTATION
#define STBI_FAILURE_USERMSG
#include <stb_image.h>
#include <stb_image_write.h>
#undef STB_IMAGE_IMPLEMENTATION
#undef STB_IMAGE_WRITE_IMPLEMENTATION
#if MIJIN_COMPILER == MIJIN_COMPILER_GCC || MIJIN_COMPILER == MIJIN_COMPILER_CLANG
#pragma GCC diagnostic pop
#endif // MIJIN_COMPILER == MIJIN_COMPILER_GCC || MIJIN_COMPILER == MIJIN_COMPILER_CLANG

#include "iwa/log.hpp"

namespace iwa
{
namespace
{
vk::Format formatByComponents(int components, ImageFormatType formatType)
{
    switch (formatType)
    {
        case ImageFormatType::UNORM:
            switch (components)
            {
                case 1:
                    return vk::Format::eR8Unorm;
                case 2:
                    return vk::Format::eR8G8Unorm;
                case 3:
                    return vk::Format::eR8G8B8Unorm;
                case 4:
                    return vk::Format::eR8G8B8A8Unorm;
            }
            break;
        case ImageFormatType::UINT:
            switch (components)
            {
                case 1:
                    return vk::Format::eR8Uint;
                case 2:
                    return vk::Format::eR8G8Uint;
                case 3:
                    return vk::Format::eR8G8B8Uint;
                case 4:
                    return vk::Format::eR8G8B8A8Uint;
            }
            break;
        case ImageFormatType::SRGB:
            switch (components)
            {
                case 1:
                    return vk::Format::eR8Srgb;
                case 2:
                    return vk::Format::eR8G8Srgb;
                case 3:
                    return vk::Format::eR8G8B8Srgb;
                case 4:
                    return vk::Format::eR8G8B8A8Srgb;
            }
            break;
    }

    logAndDie("Could not detect image format.");
}

ObjectPtr<Bitmap> loadBitmapFromSTBI(stbi_uc* data, std::size_t width, std::size_t height, std::size_t /* components */,
                   const BitmapLoadOptions& options)
{
    if (data == nullptr)
    {
        logAndDie("Could not load texture: {}", stbi_failure_reason());
    }

    const BitmapCreationArgs args =
    {
        .format = formatByComponents(/*components */ 4, options.formatType),
        .size = {static_cast<unsigned>(width), static_cast<unsigned>(height)}
    };

    ObjectPtr <Bitmap> bitmap = Bitmap::create(args);
    std::memcpy(bitmap->getData().data(), data, bitmap->getData().size());
    stbi_image_free(data);
    return bitmap;
}

int stbiReadCallback(void* user, char* data, int size)
{
    mijin::Stream& stream = *static_cast<mijin::Stream*>(user);
    std::size_t bytesRead = 0;
    const mijin::StreamError error = stream.readRaw(data, size, /* partial = */ true, &bytesRead);
    if (error != mijin::StreamError::SUCCESS)
    {
        // TODO: return what?
        return 0;
    }
    return static_cast<int>(bytesRead);
}

void stbiSkipCallback(void* user, int bytes)
{
    mijin::Stream& stream = *static_cast<mijin::Stream*>(user);
    (void) stream.seek(bytes, mijin::SeekMode::RELATIVE);
}

int stbiEofCallback(void* user)
{
    mijin::Stream& stream = *static_cast<mijin::Stream*>(user);
    return stream.isAtEnd();
}
} // namespace

ObjectPtr<Bitmap> loadBitmap(const BitmapLoadOptions& options)
{
    int width, height, components; // NOLINT(cppcoreguidelines-init-variables)
    const stbi_io_callbacks callbacks{
        .read = &stbiReadCallback,
        .skip = &stbiSkipCallback,
        .eof = &stbiEofCallback
    };
    stbi_uc* data = stbi_load_from_callbacks(
        /* clbk             = */ &callbacks,
        /* user             = */ options.stream,
        /* x                = */ &width,
        /* y                = */ &height,
        /* channels_in_file = */ &components,
        /* desired_channels = */ 4);
    // NOLINTNEXTLINE(clang-analyzer-core.CallAndMessage) Yes, width isn't initialized if STB fails. No, that is not a problem
    return loadBitmapFromSTBI(data, width, height, components, options);
}

void saveBitmap(const Bitmap& bitmap, const BitmapSaveOptions& options)
{
    namespace fs = std::filesystem;

    BitmapCodec codec = options.codec;
    if (codec == BitmapCodec::NONE)
    {
        std::string extension = fs::path(options.fileName).extension().string();
        std::ranges::transform(extension, extension.begin(), [](char chr)
        { return static_cast<char>(std::tolower(chr)); });
        if (extension == ".png")
        {
            codec = BitmapCodec::PNG;
        } else if (extension == ".jpg" || extension == ".jpeg")
        {
            codec = BitmapCodec::JPEG;
        } else
        {
            logAndDie("Couldn't guess image file codec from file extension: {}.", extension);
        }
    }

    int comp = 0;
    switch (bitmap.getFormat())
    {
        case vk::Format::eR8G8B8A8Unorm:
        case vk::Format::eR8G8B8A8Srgb:
            comp = 4;
            break;
        case vk::Format::eR8G8B8Unorm:
        case vk::Format::eR8G8B8Srgb:
            comp = 3;
            break;
        default:
            logAndDie("Cannot write this image format (yet).");
            break;
    }

    switch (codec)
    {
        case BitmapCodec::NONE:
            assert(0);
            break;
        case BitmapCodec::PNG:
            stbi_write_png(
                /* filename        = */ options.fileName.c_str(),
                /* w               = */ static_cast<int>(bitmap.getSize().width),
                /* h               = */ static_cast<int>(bitmap.getSize().height),
                /* comp            = */ comp,
                /* data            = */ bitmap.getData().data(),
                /* stride_in_bytes = */ comp * static_cast<int>(bitmap.getSize().width));
            break;
        case BitmapCodec::JPEG:
            stbi_write_jpg(
                /* filename        = */ options.fileName.c_str(),
                /* x               = */ static_cast<int>(bitmap.getSize().width),
                /* y               = */ static_cast<int>(bitmap.getSize().height),
                /* comp            = */ comp,
                /* data            = */ bitmap.getData().data(),
                /* quality         = */ options.jpegQuality
            );
            break;
    }
}
} // namespace iwa