#include "iwa/io/font.hpp"

#include <mijin/util/iterators.hpp>

#define STB_RECT_PACK_IMPLEMENTATION
#define STB_TRUETYPE_IMPLEMENTATION
#include <stb_rect_pack.h>
#include <stb_truetype.h>
#undef STB_RECT_PACK_IMPLEMENTATION
#undef STB_TRUETYPE_IMPLEMENTATION

namespace iwa
{
namespace
{
std::vector<int> getRenderCodePoints()
{
    std::vector<int> result;
    for (int chr = 32; chr <= 255; ++chr) {
        result.push_back(chr);
    }
    return result;
}
}

ObjectPtr<Font> loadFont(const FontLoadOptions& options)
{
    mijin::TypelessBuffer fontData;
    if (const mijin::StreamError error = options.stream->readRest(fontData); error != mijin::StreamError::SUCCESS) {
        throw std::runtime_error("IO error reading font from stream.");
    }

    // init font data
    stbtt_fontinfo fontInfo;
    const int fontOffset = stbtt_GetFontOffsetForIndex(static_cast<const unsigned char*>(fontData.data()), 0);
    stbtt_InitFont(&fontInfo, static_cast<const unsigned char*>(fontData.data()), fontOffset);

    // prepare data for packing
    std::vector<int> codepoints = getRenderCodePoints();
    std::vector<stbtt_packedchar> chardata;
    std::vector<stbrp_rect> rects;

    chardata.resize(codepoints.size());
    rects.resize(codepoints.size());
    stbtt_pack_range range = {
        .font_size = options.size,
        .first_unicode_codepoint_in_range = 0,
        .array_of_unicode_codepoints = codepoints.data(),
        .num_chars = static_cast<int>(codepoints.size()),
        .chardata_for_range = chardata.data()
    };

    // TODO: this is just a guess, maybe there is a better way to detect this
    int oversample = 1;
    if (options.size < 30) {
        oversample = 4;
    }
    else if (options.size < 35)  {
        oversample = 3;
    }
    else if (options.size < 40) {
        oversample = 4;
    }

    int textureSize = 64;
    mijin::TypelessBuffer pixels;
    while (textureSize <= 4096)
    {
        pixels.resize(textureSize * textureSize * sizeof(unsigned char)); // NOLINT

        // reset char data
        for (stbtt_packedchar& packedChar : chardata) {
            packedChar.x0 = packedChar.y0 = packedChar.x1 = packedChar.y1 = 0;
        }

        // try packing with this size
        stbtt_pack_context packContext;
        stbtt_PackBegin(
            /* spc             = */ &packContext,
            /* pixels          = */ static_cast<unsigned char*>(pixels.data()),
            /* width           = */ textureSize,
            /* height          = */ textureSize,
            /* stride_in_bytes = */ 0,
            /* padding         = */ 1,
            /* alloc_context   = */ nullptr
        );
        stbtt_PackSetOversampling(&packContext, oversample, oversample); // TODO: make adjustable for better font quality

        const int nRects = stbtt_PackFontRangesGatherRects(&packContext, &fontInfo, &range, 1, rects.data());
        stbtt_PackFontRangesPackRects(&packContext, rects.data(), nRects);

        bool allPacked = true;
        for (int rectIdx = 0; rectIdx < nRects; ++rectIdx)
        {
            if (!rects[rectIdx].was_packed) {
                allPacked = false;
                break;
            }
        }
        if (!allPacked)
        {
            textureSize *= 2;
            continue;
        }

        stbtt_PackFontRangesRenderIntoRects(
            /* spc        = */ &packContext,
            /* info       = */ &fontInfo,
            /* ranges     = */ &range,
            /* num_ranges = */ 1,
            /* rects      = */ rects.data()
        );

        stbtt_PackEnd(&packContext);
        break;
    }

    // now create a bitmap from pixels
    ObjectPtr<Bitmap> bitmap = Bitmap::create(BitmapCreationArgs{
        .format = vk::Format::eR8Unorm,
        .size = {
            .width = static_cast<std::uint32_t>(textureSize),
            .height = static_cast<std::uint32_t>(textureSize)
        },
        .initialData = std::move(pixels)
    });
    std::unordered_map<char32_t, GlyphInfo> glyphMap;
    glyphMap.reserve(chardata.size());

    for (auto [codepoint, chrdata] : mijin::zip(codepoints, chardata))
    {
        glyphMap.emplace(static_cast<char32_t>(codepoint), GlyphInfo{
            .uvPos0 = {
                static_cast<float>(chrdata.x0) / static_cast<float>(textureSize),
                static_cast<float>(chrdata.y0) / static_cast<float>(textureSize)
            },
            .uvPos1 = {
                static_cast<float>(chrdata.x1) / static_cast<float>(textureSize),
                static_cast<float>(chrdata.y1) / static_cast<float>(textureSize)
            },
            .xOffsetBefore = chrdata.xoff,
            .xOffsetAfter = chrdata.xoff2,
            .yOffsetBefore = chrdata.yoff,
            .yOffsetAfter = chrdata.yoff2,
            .xAdvance = chrdata.xadvance
        });
    }

    const float scale = options.size > 0 ? stbtt_ScaleForPixelHeight(&fontInfo, options.size) : stbtt_ScaleForMappingEmToPixels(&fontInfo, -options.size);
    int ascent = 0, descent = 0, lineGap = 0;
    stbtt_GetFontVMetrics(&fontInfo, &ascent, &descent, &lineGap);

    return Font::create(FontCreationArgs{
        .bitmap = std::move(bitmap),
        .glyphMap = std::move(glyphMap),
        .metrics = {
            .ascent = scale * static_cast<float>(ascent),
            .descent = scale * static_cast<float>(descent),
            .lineGap = scale * static_cast<float>(lineGap),
            .sizeFactor = 1.f / static_cast<float>(oversample)
        }
    });
}
} // namespace iwa