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

// Text drawn by the UI kernels has to be antialiased: a glyph's edge ramps
// from transparent to opaque over one screen pixel, at any size. FontAtlas
// keeps each glyph as an SDF in which stb_truetype writes kOnEdgeValue +
// kPixelDistScale * distance (in atlas pixels) per byte; the kernels decode a
// sample back into a distance and divide it by the atlas pixels one screen
// pixel spans along the glyph. ui-text and ui-fused used to multiply by
// kPixelDistScale instead of 255 / kPixelDistScale, which made the distance
// 1024/255 (about 4) times too large and the ramp a quarter of a pixel wide:
// nearly every edge pixel came out fully on or fully off, and text looked
// pixelated at every size.
//
// The test draws one scene through DispatchText and again through
// DispatchFused, captures the frame and checks:
//
//   - one string at 14 and 64 px, below and above the atlas's 32 px base
//     size, and at 14 px stretched to three times its width, each against
//     stb_truetype's own antialiased rasterizer at the same size, aspect and
//     sub-pixel position: about as many pixels are partly covered, and where
//     the reference is partly covered the coverage is close to it. The
//     quarter-pixel ramp gives a quarter as many; a ramp one atlas pixel
//     wide, which ignores the glyph's size, half as many at 14 px and twice
//     as many at 64 px; one ramp width for both axes of the stretched row,
//     the width across its horizontal edges, makes it three pixels wide
//     across the stems;
//   - a glyph turned by swapping its UVs, the one way a GlyphItem flips, has
//     to draw its mirror image. A signed UV span gave it hard edges;
//   - atlas texels no glyph reaches hold 0, like every glyph's padding, and
//     drawn shrunk far below 4 px they have to draw nothing. A ramp wider
//     than the distances a byte can hold left the whole quad translucent;
//   - the fused text phase draws what DispatchText draws.
//
// Needs a live Wayland compositor and a Vulkan device, like UIFallbackParity.

#include "vulkan/vulkan.h"
#include "vulkan/vk_enum_string_helper.h"
#include <cstdlib>

// PNG decode for the captured frame. The library only compiles the writer
// (Window::CaptureNextFrame); STB_IMAGE_STATIC keeps the reader private to this file.
#define STB_IMAGE_IMPLEMENTATION
#define STB_IMAGE_STATIC
#define STBI_ONLY_PNG
#include "../../lib/stb_image.h"
// The reference rasterizer. Declarations only: the definitions are the
// library's own (Font.cpp), and Font::font is the stbtt_fontinfo they take.
#include "../../lib/stb_truetype.h"

import Crafter.Graphics;
import Crafter.Event;
import std;

using namespace Crafter;
namespace fs = std::filesystem;

namespace {

int failures = 0;

void Check(bool ok, std::string_view what) {
    std::println("{}  {}", ok ? "PASS" : "FAIL", what);
    if (!ok) ++failures;
}

// The test binary's own directory: the build places the compiled shaders
// and font.ttf next to it, and the test runner's working directory is not
// that directory.
fs::path ExeDir() {
    return fs::read_symlink("/proc/self/exe").parent_path();
}

constexpr std::uint32_t kWidth  = 640;
constexpr std::uint32_t kHeight = 560;

// Straight stems, curves and diagonals.
constexpr std::string_view kText = "Hamburgefonstiv";
constexpr float kTextX = 20.0f;
constexpr std::uint32_t kGlyphCapacity = 64;

struct Row {
    float size;
    float baseline;
    float stretch;   // horizontal scale applied to ShapeText's glyphs
};
constexpr std::array<Row, 3> kRows{{ {14.0f, 40.0f, 1.0f}, {64.0f, 120.0f, 1.0f}, {14.0f, 180.0f, 3.0f} }};
// The glyph mirrored by its UVs, taken from the 64 px row, and where the
// original and the flipped copy go: whole pixels, so the check maps pixel to
// pixel. 'g' is neither left-right nor top-bottom symmetric.
constexpr char kMirrorGlyph = 'g';
constexpr float kMirrorX = 340.0f;
constexpr float kMirrorY = 150.0f;
constexpr float kMirrorGap = 16.0f;
// The empty-texel item: 0.4 of the atlas drawn over 32 px is 12.8 atlas
// pixels per screen pixel, a glyph under 3 px. FontAtlas packs shelves from
// the top left, and this string fills one, so the bottom right is untouched.
constexpr float kEmptyX = 500.0f;
constexpr float kEmptyY = 150.0f;
constexpr float kEmptySize = 32.0f;
constexpr float kEmptyUv0 = 0.55f;
constexpr float kEmptyUv1 = 0.95f;
// The fused copy of the scene sits this much lower: a whole number of
// pixels, so it samples the atlas exactly where the standalone copy does.
constexpr int kFusedOffset = 270;

// A pixel is an edge pixel when it is partly covered.
constexpr float kPartialLow  = 0.02f;
constexpr float kPartialHigh = 0.98f;
// The SDF ramp is not the reference's box filter, so a correct kernel lands
// near the reference rather than on it: 1.02, 1.07 and 1.09 times its edge
// pixels (14 px, 64 px, 14 px stretched) on an RTX 4090, 1.02, 1.11 and 1.09
// on llvmpipe, with a mean difference of 0.065 to 0.083 on them. The
// quarter-pixel ramp gives 0.31 and 0.21 times and 0.20; a ramp one atlas
// pixel wide 0.55 and 2.19 times and 0.14; one width for both axes 1.62 times
// on the stretched row.
constexpr float kMinEdgeRatio = 0.7f;
constexpr float kMaxEdgeRatio = 1.4f;
constexpr float kMaxEdgeError = 0.12f;

template<typename T>
void CreateItemBuffer(GraphicsBuffer<T, true>& buffer, std::uint32_t count) {
    buffer.Create(
        VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT,
        VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, count);
}

// Per-pixel coverage, 0 for the background and 1 for fully covered.
struct Coverage {
    int width = 0;
    int height = 0;
    std::vector<float> value;

    float At(int x, int y) const { return value[static_cast<std::size_t>(y) * width + x]; }
};

int Code(float coverage) {
    return static_cast<int>(std::lround(coverage * 255.0f));
}

// White text over an opaque black background on an 8-bit UNORM swapchain:
// the red channel is the coverage the kernels composited.
Coverage LoadFrameCoverage(const fs::path& path) {
    Coverage frame;
    int channels = 0;
    stbi_uc* pixels = stbi_load(path.string().c_str(), &frame.width, &frame.height, &channels, STBI_rgb_alpha);
    if (pixels == nullptr) return frame;
    frame.value.resize(static_cast<std::size_t>(frame.width) * frame.height);
    for (std::size_t i = 0; i < frame.value.size(); ++i) frame.value[i] = pixels[i * 4] / 255.0f;
    stbi_image_free(pixels);
    return frame;
}

// stb_truetype's area coverage of kText laid out the way UIRenderer::ShapeText
// lays it out, stretched as the test stretches the row: the same pen positions
// (the atlas's advances, scaled and summed in the same order), each glyph
// rasterised at its own sub-pixel offset and composited "over" as the kernels
// do.
Coverage ReferenceCoverage(Font& font, const FontAtlas& atlas, const Row& row) {
    Coverage ref{static_cast<int>(kWidth), static_cast<int>(kHeight),
                 std::vector<float>(static_cast<std::size_t>(kWidth) * kHeight, 0.0f)};
    const float scale      = stbtt_ScaleForPixelHeight(&font.font, row.size);
    const float atlasScale = row.size / FontAtlas::kBaseSize;
    float cursor = 0.0f;
    for (const char c : kText) {
        const int codepoint = static_cast<unsigned char>(c);
        const Glyph* glyph = atlas.Lookup(font, static_cast<std::uint32_t>(codepoint));
        if (glyph == nullptr) return ref;
        const float pen   = kTextX + cursor * row.stretch;
        const float shift = pen - std::floor(pen);
        const float scaleX = scale * row.stretch;
        int x0 = 0, y0 = 0, x1 = 0, y1 = 0;
        stbtt_GetCodepointBitmapBoxSubpixel(&font.font, codepoint, scaleX, scale, shift, 0.0f, &x0, &y0, &x1, &y1);
        const int w = x1 - x0;
        const int h = y1 - y0;
        if (w > 0 && h > 0) {
            std::vector<unsigned char> bitmap(static_cast<std::size_t>(w) * h);
            stbtt_MakeCodepointBitmapSubpixel(&font.font, bitmap.data(), w, h, w, scaleX, scale, shift, 0.0f, codepoint);
            const int left = static_cast<int>(std::floor(pen)) + x0;
            const int top  = static_cast<int>(row.baseline) + y0;
            for (int y = 0; y < h; ++y) {
                for (int x = 0; x < w; ++x) {
                    const int px = left + x;
                    const int py = top + y;
                    if (px < 0 || py < 0 || px >= ref.width || py >= ref.height) continue;
                    const float a = bitmap[static_cast<std::size_t>(y) * w + x] / 255.0f;
                    float& dst = ref.value[static_cast<std::size_t>(py) * ref.width + px];
                    dst = a + dst * (1.0f - a);
                }
            }
        }
        cursor += glyph->advance * atlasScale;
    }
    return ref;
}

// Half-open pixel box.
struct Box {
    int x0 = 0, y0 = 0, x1 = 0, y1 = 0;
};

// The reference's covered pixels, plus a margin for the kernels' ramp, which
// reaches half a pixel past the outline.
Box CoveredBox(const Coverage& ref) {
    Box box{ref.width, ref.height, 0, 0};
    for (int y = 0; y < ref.height; ++y) {
        for (int x = 0; x < ref.width; ++x) {
            if (ref.At(x, y) <= 0.0f) continue;
            box.x0 = std::min(box.x0, x);
            box.y0 = std::min(box.y0, y);
            box.x1 = std::max(box.x1, x + 1);
            box.y1 = std::max(box.y1, y + 1);
        }
    }
    constexpr int kMargin = 2;
    box.x0 = std::max(box.x0 - kMargin, 0);
    box.y0 = std::max(box.y0 - kMargin, 0);
    box.x1 = std::min(box.x1 + kMargin, ref.width);
    box.y1 = std::min(box.y1 + kMargin, ref.height);
    return box;
}

bool Partial(float coverage) {
    return coverage > kPartialLow && coverage < kPartialHigh;
}

struct EdgeStats {
    int frameEdges = 0;       // partly covered pixels in the frame
    int referenceEdges = 0;   // partly covered pixels in the reference
    float meanError = 0.0f;   // mean |frame - reference| over the reference's edge pixels
};

// Compares the frame, `dy` pixels lower, with the reference inside `box`.
EdgeStats CompareEdges(const Coverage& frame, int dy, const Coverage& ref, const Box& box) {
    EdgeStats stats;
    double error = 0.0;
    for (int y = box.y0; y < box.y1; ++y) {
        for (int x = box.x0; x < box.x1; ++x) {
            const float f = frame.At(x, y + dy);
            const float r = ref.At(x, y);
            if (Partial(f)) ++stats.frameEdges;
            if (Partial(r)) {
                ++stats.referenceEdges;
                error += std::abs(f - r);
            }
        }
    }
    if (stats.referenceEdges > 0) stats.meanError = static_cast<float>(error / stats.referenceEdges);
    return stats;
}

void CheckEdges(std::string_view what, const EdgeStats& stats) {
    const float ratio = stats.referenceEdges > 0
        ? static_cast<float>(stats.frameEdges) / static_cast<float>(stats.referenceEdges)
        : 0.0f;
    Check(ratio >= kMinEdgeRatio && ratio <= kMaxEdgeRatio,
          std::format("{}: {} antialiased edge pixels, stb_truetype's rasterizer {} ({:.2f} times, allowed {:.1f} to {:.1f})",
                      what, stats.frameEdges, stats.referenceEdges, ratio, kMinEdgeRatio, kMaxEdgeRatio));
    Check(stats.referenceEdges > 0 && stats.meanError <= kMaxEdgeError,
          std::format("{}: coverage on the reference's edge pixels is {:.3f} from it on average (allowed {:.2f})",
                      what, stats.meanError, kMaxEdgeError));
}

std::string RowName(const Row& row) {
    return row.stretch == 1.0f ? std::format("{} px", row.size)
                               : std::format("{} px stretched {}x", row.size, row.stretch);
}

int Run() {
    const fs::path dir = ExeDir();
    const fs::path outDir = fs::temp_directory_path() / "crafter-text-antialiasing";
    std::error_code ec;
    fs::remove_all(outDir, ec);
    fs::create_directories(outDir);
    const fs::path png = outDir / "frame.png";

    Device::Initialize();
    Window window(kWidth, kHeight, "TextAntialiasing test");
    // CaptureNextFrame writes 8-bit RGBA, swapping BGRA, and the coverage is
    // read from those bytes. A surface whose 8-bit formats can't be storage
    // images gets another format, which neither understands.
    if (window.colorFormat != VK_FORMAT_B8G8R8A8_UNORM && window.colorFormat != VK_FORMAT_R8G8B8A8_UNORM) {
        std::println("SKIP  the swapchain format {} isn't 8-bit RGBA or BGRA, which the captured frame needs",
                     string_VkFormat(window.colorFormat));
        return 77;
    }
    VkCommandBuffer init = window.StartInit();

    GraphicsDescriptorHeap heap;
    heap.Initialize(/*images*/ 8, /*buffers*/ 8, /*samplers*/ 4);
    window.descriptorHeap = &heap;

    Font font(dir / "font.ttf");
    FontAtlas atlas;
    atlas.Initialize(init);

    UIRenderer ui;
    ui.fontAtlas = &atlas;
    ui.Initialize(window, heap, init,
                  dir / "ui-quads.comp.spv", dir / "ui-circles.comp.spv",
                  dir / "ui-images.comp.spv", dir / "ui-text.comp.spv",
                  dir / "ui-fused.comp.spv");
    window.passes.push_back(&ui);

    GraphicsBuffer<QuadItem,  true> background;
    GraphicsBuffer<GlyphItem, true> standalone;
    GraphicsBuffer<GlyphItem, true> fused;
    CreateItemBuffer(background, 1);
    CreateItemBuffer(standalone, kGlyphCapacity);
    CreateItemBuffer(fused, kGlyphCapacity);
    BufferSlot backgroundSlot = ui.RegisterBuffer(background);
    BufferSlot standaloneSlot = ui.RegisterBuffer(standalone);
    BufferSlot fusedSlot      = ui.RegisterBuffer(fused);

    // Filled on the frame: the rows, then the mirrored pair and the
    // empty-texel item. The fused buffer holds the same items kFusedOffset
    // lower.
    std::uint32_t glyphCount = 0;
    std::array<std::uint32_t, kRows.size()> rowCounts{};
    GlyphItem original{};
    GlyphItem flipped{};
    EventListener<UIBuildArgs> buildSub(&ui.onBuild, [&](UIBuildArgs a) {
        VkCommandBuffer cmd = a.cmd;
        // The window never clears its image. Opaque black under white text
        // makes the red channel the coverage.
        background.value[0] = QuadItem{ 0, 0, 4096, 4096,  0.0f, 0.0f, 0.0f, 1.0f,  0, 0, 0, 0,  0, 0, 0, 0 };

        GlyphItem* items = standalone.value;
        glyphCount = 0;
        std::uint32_t mirrorIndex = kGlyphCapacity;
        for (std::size_t r = 0; r < kRows.size(); ++r) {
            const Row& row = kRows[r];
            const std::uint32_t first = glyphCount;
            rowCounts[r] = ui.ShapeText(font, row.size, kTextX, row.baseline, kText, {1.0f, 1.0f, 1.0f, 1.0f},
                                        items + first, kGlyphCapacity - first);
            glyphCount += rowCounts[r];
            if (row.stretch != 1.0f) {
                for (std::uint32_t i = first; i < glyphCount; ++i) {
                    items[i].x = kTextX + (items[i].x - kTextX) * row.stretch;
                    items[i].w *= row.stretch;
                }
            }
            // Every character of kText has ink, so item i is character i.
            if (row.size == 64.0f && row.stretch == 1.0f && rowCounts[r] == kText.size()) {
                mirrorIndex = first + static_cast<std::uint32_t>(kText.find(kMirrorGlyph));
            }
        }
        if (mirrorIndex < kGlyphCapacity && glyphCount + 3 <= kGlyphCapacity) {
            original = items[mirrorIndex];
            original.x = kMirrorX;
            original.y = kMirrorY;
            flipped = original;
            flipped.x = kMirrorX + original.w + kMirrorGap;
            std::swap(flipped.u0, flipped.u1);
            std::swap(flipped.v0, flipped.v1);
            items[glyphCount++] = original;
            items[glyphCount++] = flipped;

            GlyphItem empty{};
            empty.x = kEmptyX;
            empty.y = kEmptyY;
            empty.w = kEmptySize;
            empty.h = kEmptySize;
            empty.u0 = kEmptyUv0;
            empty.v0 = kEmptyUv0;
            empty.u1 = kEmptyUv1;
            empty.v1 = kEmptyUv1;
            empty.r = empty.g = empty.b = empty.a = 1.0f;
            items[glyphCount++] = empty;
        }
        for (std::uint32_t i = 0; i < glyphCount; ++i) {
            fused.value[i] = items[i];
            fused.value[i].y += kFusedOffset;
        }

        constexpr VkAccessFlags        read  = VK_ACCESS_SHADER_READ_BIT;
        constexpr VkPipelineStageFlags stage = VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT;
        background.FlushDevice(cmd, read, stage);
        standalone.FlushDevice(cmd, read, stage);
        fused.FlushDevice(cmd, read, stage);

        ui.DispatchQuads(cmd, backgroundSlot, 1);
        ui.DispatchText(cmd, standaloneSlot, glyphCount);
        ui.DispatchFused(cmd, {}, {}, {}, {fusedSlot, glyphCount});
    });

    // The first frame is enough: DispatchText and DispatchFused upload the
    // glyphs ShapeText rasterised into the atlas before they draw, in the
    // same command buffer, and the capture waits for the frame.
    window.FinishInit();
    window.CaptureNextFrame(png);
    window.Render();
    // The buffers, atlas and heap go out of scope while the frame may still
    // be presenting; settle the queue first.
    Device::CheckVkResult(vkQueueWaitIdle(Device::queue));

    if (std::getenv("CRAFTER_GRAPHICS_VALIDATION") != nullptr) {
        Check(Device::validationErrorCount == 0,
              std::format("no Vulkan validation errors while rendering and capturing ({} seen)", Device::validationErrorCount));
    }
    const bool shaped = std::ranges::all_of(rowCounts, [](std::uint32_t n) { return n == kText.size(); });
    Check(shaped && glyphCount == kText.size() * kRows.size() + 3,
          std::format("every glyph was shaped ({} items for {} rows of {} glyphs and 3 more)",
                      glyphCount, kRows.size(), kText.size()));

    const Coverage frame = LoadFrameCoverage(png);
    Check(!frame.value.empty(), "the frame was captured");
    if (frame.value.empty() || !shaped) return EXIT_FAILURE;

    // Both copies of the scene, one above the other, inside the frame.
    const float sceneRight  = std::max(kEmptyX + kEmptySize, flipped.x + flipped.w);
    const float sceneBottom = std::max(kEmptyY + kEmptySize, flipped.y + flipped.h);
    const bool fits = sceneRight <= static_cast<float>(kWidth) && sceneBottom <= static_cast<float>(kFusedOffset)
                   && frame.width >= static_cast<int>(kWidth) && frame.height >= 2 * kFusedOffset;
    Check(fits, std::format("the {}x{} frame holds both copies of the scene, {} by {} px each",
                            frame.width, frame.height, sceneRight, sceneBottom));
    if (!fits) return EXIT_FAILURE;

    for (const Row& row : kRows) {
        const Coverage ref = ReferenceCoverage(font, atlas, row);
        const Box box = CoveredBox(ref);
        CheckEdges(std::format("DispatchText, {}", RowName(row)), CompareEdges(frame, 0, ref, box));
        CheckEdges(std::format("DispatchFused, {}", RowName(row)), CompareEdges(frame, kFusedOffset, ref, box));
    }

    // The flipped copy against the original read backwards on both axes.
    // One code value of slack: the two sample the same atlas points through
    // different float arithmetic.
    {
        const int w  = static_cast<int>(original.w);
        const int h  = static_cast<int>(original.h);
        const int ox = static_cast<int>(original.x);
        const int fx = static_cast<int>(flipped.x);
        const int y0 = static_cast<int>(original.y);
        int maxDiff = 0;
        float peak = 0.0f;
        for (int y = 0; y < h; ++y) {
            for (int x = 0; x < w; ++x) {
                const float a = frame.At(ox + x, y0 + y);
                const float b = frame.At(fx + (w - 1 - x), y0 + (h - 1 - y));
                maxDiff = std::max(maxDiff, std::abs(Code(a) - Code(b)));
                peak = std::max(peak, a);
            }
        }
        Check(peak >= 0.99f && maxDiff <= 1,
              std::format("a '{}' turned by swapping its UVs draws its mirror image (max difference {}, peak coverage {:.2f})",
                          kMirrorGlyph, maxDiff, peak));
    }

    // Texels holding 0 draw nothing, however far they are shrunk.
    {
        int maxCode = 0;
        for (int y = static_cast<int>(kEmptyY); y < static_cast<int>(kEmptyY + kEmptySize); ++y) {
            for (int x = static_cast<int>(kEmptyX); x < static_cast<int>(kEmptyX + kEmptySize); ++x) {
                maxCode = std::max(maxCode, Code(frame.At(x, y)));
            }
        }
        Check(maxCode == 0, std::format("empty atlas texels shrunk to {:.1f} atlas pixels per screen pixel draw nothing (max {})",
                                        (kEmptyUv1 - kEmptyUv0) * FontAtlas::kAtlasSize / kEmptySize, maxCode));
    }

    // The whole fused copy against the whole standalone one. One code value
    // of slack: the two kernels share the per-pixel math, but the driver
    // compiles them separately and may contract or reorder it differently on
    // an antialiased edge.
    {
        int maxDiff = 0;
        for (int y = 0; y < kFusedOffset; ++y) {
            for (int x = 0; x < static_cast<int>(kWidth); ++x) {
                maxDiff = std::max(maxDiff, std::abs(Code(frame.At(x, y)) - Code(frame.At(x, y + kFusedOffset))));
            }
        }
        Check(maxDiff <= 1, std::format("DispatchFused draws what DispatchText draws (max difference {})", maxDiff));
    }

    if (failures != 0) std::println("      frame kept for inspection: {}", png.string());
    return failures == 0 ? EXIT_SUCCESS : EXIT_FAILURE;
}

} // namespace

int main() {
    const int result = Run();
    if (result != 77) std::println("{}", result == EXIT_SUCCESS ? "all checks passed" : "check(s) failed");
    return result;
}