#include "./packer.hpp"

#include <finders_interface.h>
#include <mijin/io/stream.hpp>
#include <mijin/util/iterators.hpp>

#define STBI_ASSERT(x) MIJIN_ASSERT(x, #x)
#define STB_IMAGE_IMPLEMENTATION
#include <stb_image.h>
#undef STB_IMAGE_IMPLEMENTATION
#undef STBI_ASSERT

#define STBIW_ASSERT(x) MIJIN_ASSERT(x, #x)
#define STB_IMAGE_WRITE_IMPLEMENTATION
#include <stb_image_write.h>
#include <mijin/util/string.hpp>
#undef STB_IMAGE_WRITE_IMPLEMENTATION
#undef STBIW_ASSERT

namespace sdl_gpu_test
{
void Packer::addImage(std::string name, const fs::path& path)
{
    int width = 0;
    int height = 0;
    int numComponents = 0;
    stbi_uc* pixels = stbi_load(
        /* filename         = */ path.generic_string().c_str(),
        /* x                = */ &width,
        /* y                = */ &height,
        /* channels_in_file = */ &numComponents,
        /* desired_channels = */ 4
    );
    if (pixels == nullptr)
    {
        throw std::runtime_error("Error loading image.");
    }

    mImages.push_back({
        .name = std::move(name),
        .data = std::unique_ptr<Pixel[], FreeDeleter>(reinterpret_cast<Pixel*>(pixels)),
        .width = static_cast<unsigned>(width),
        .height = static_cast<unsigned>(height)
    });
}

void Packer::addImageList(const fs::path& path)
{
    mijin::FileStream stream;
    mijin::throwOnError(stream.open(path.generic_string(), mijin::FileOpenMode::READ));
    while (!stream.isAtEnd())
    {
        std::string line;
        mijin::throwOnError(stream.readLine(line));
        if (line.empty())
        {
            continue;
        }
        const auto [name, file] = mijin::splitFixed<2>(line, " ");
        if (file.empty())
        {
            throw std::runtime_error("Invalid image list format, missing file name.");
        }
        addImage(std::string(name), path.parent_path() / file);
    }
}


void Packer::pack(const fs::path& outputPath)
{
    static constexpr bool allowFlip = false; // TODO: this could work?
    static constexpr rectpack2D::flipping_option flippingOption = rectpack2D::flipping_option::DISABLED;
    using spaces_type_t = rectpack2D::empty_spaces<allowFlip, rectpack2D::default_empty_spaces>;
    using rect_type_t = rectpack2D::output_rect_t<spaces_type_t>;

    auto reportSuccessful = [&](rect_type_t&)
    {
        return rectpack2D::callback_result::CONTINUE_PACKING;
    };

    auto reportUnsuccessful = [&](rect_type_t&)-> rectpack2D::callback_result
    {
        throw std::runtime_error("Could not fit images.");
    };

    static constexpr int maxSide = 4096;
    static constexpr int discardStep = -4;

    std::vector<rect_type_t> rectangles;
    rectangles.reserve(mImages.size());

    for (const ImageData& imageData : mImages)
    {
        rectangles.emplace_back(0, 0, imageData.width, imageData.height);
    }

    const rectpack2D::rect_wh resultSize = rectpack2D::find_best_packing<spaces_type_t>(
        rectangles,
        rectpack2D::make_finder_input(
            maxSide, discardStep, reportSuccessful, reportUnsuccessful, flippingOption
        )
    );

    std::vector<Pixel> resultPixels;
    resultPixels.resize(resultSize.area());

    for (const auto& [imageData, rect] : mijin::zip(mImages, rectangles))
    {
        for (unsigned row = 0; row < imageData.height; ++row)
        {
            const Pixel* src = &imageData.data[row * imageData.width];
            Pixel* dst = &resultPixels[(rect.y + row) * resultSize.w + rect.x];
            std::copy_n(src, imageData.width, dst);
        }
    }

    const int writeResult = stbi_write_png(
        /* filename     = */ outputPath.generic_string().c_str(),
        /* x            = */ resultSize.w,
        /* y            = */ resultSize.h,
        /* comp         = */ 4,
        /* data         = */ resultPixels.data(),
        /* stride_bytes = */ resultSize.w * sizeof(Pixel)
    );

    if (!writeResult)
    {
        throw std::runtime_error("Error writing result image.");
    }

    fs::path atlasPath(outputPath);
    atlasPath.replace_extension("txt");
    mijin::FileStream fileStream;
    mijin::throwOnError(fileStream.open(atlasPath.generic_string(), mijin::FileOpenMode::WRITE));
    for (const auto& [imageData, rect] : mijin::zip(mImages, rectangles))
    {
        mijin::throwOnError(fileStream.writeText(std::format("{} {} {} {} {}\n",
            imageData.name, rect.x, rect.y, rect.w, rect.h)));
    }
    fileStream.close();
}
}