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

// Crafter.Graphics UI shader contract — shared by every standard UI compute
// shader and intended to be #included by user-authored shaders that want to
// dispatch alongside them. Layouts here are FROZEN: only additive changes
// (using the reserved `flags` bits or `_pad`).

#extension GL_EXT_shader_image_load_formatted : enable
#extension GL_EXT_shader_explicit_arithmetic_types_int16 : enable

// ─── binding model ──────────────────────────────────────────────────────
// Every shader that includes this file is compiled twice from one source.
// The default build indexes VK_EXT_descriptor_heap directly. The second
// build defines CRAFTER_GRAPHICS_DESCRIPTOR_SETS and reads the same slots
// from one bindless descriptor set, for drivers without the heap extension
// (output: <name>.sets.spv, picked by ComputeShader::Load at runtime).
//
// Declare bindless resources with these macros instead of writing
// `descriptor_heap` or a set/binding pair, and both builds stay in sync.
// Slot numbers mean the same thing in both: UI_HEAP_IMAGES and
// UI_HEAP_TEXTURES share one image slot space, and any number of buffer
// blocks may alias UI_HEAP_BUFFERS.
#ifdef CRAFTER_GRAPHICS_DESCRIPTOR_SETS
#define UI_HEAP_IMAGES   set = 0, binding = 0
#define UI_HEAP_TEXTURES set = 0, binding = 1
#define UI_HEAP_SAMPLERS set = 0, binding = 2
#define UI_HEAP_BUFFERS  set = 0, binding = 3
#else
#extension GL_EXT_descriptor_heap : enable
#define UI_HEAP_IMAGES   descriptor_heap
#define UI_HEAP_TEXTURES descriptor_heap
#define UI_HEAP_SAMPLERS descriptor_heap
#define UI_HEAP_BUFFERS  descriptor_heap
#endif
#extension GL_EXT_nonuniform_qualifier : enable
#extension GL_EXT_buffer_reference : enable

// ─── bindless heap declarations ─────────────────────────────────────────
// The same heap slot can be read as either uiImages[] (storage image) or
// uiTextures[] (sampled image) depending on which descriptor was written
// at that slot. Samplers live in a separate sampler heap.
layout(UI_HEAP_IMAGES)   uniform image2D    uiImages[];
layout(UI_HEAP_TEXTURES) uniform texture2D  uiTextures[];
layout(UI_HEAP_SAMPLERS) uniform sampler    uiSamplers[];

// ─── push-constant header ───────────────────────────────────────────────
// Every UI dispatch's push-constant struct begins with this. User shaders
// MUST embed it as the first member so UIRenderer::FillHeader works.
struct UIDispatchHeader {
    uint  outImage;       // heap slot of the swapchain image (this frame)
    uint  itemBuffer;     // heap slot of the item SSBO
    uvec2 surfaceSize;    // window pixel size
    vec4  clipRectPx;     // (xy, wh) — every standard shader honors this
    uint  itemCount;
    uint  frameIdx;
    uint  flags;          // feature bits — high bit reserved by the renderer
    uint  _pad;           // reserved — keep zeroed
};

// ─── header flag bits ───────────────────────────────────────────────────
// `flags` is otherwise free for user-defined feature bits, but the renderer
// reserves the high bit: UI_FLAG_CLIP is set by UIRenderer::FillHeader when
// the clip rect does not already cover the whole surface. uiResolveScreenPixel
// gates its four clipRectPx compares on this bit so the overwhelmingly common
// full-surface case (clipRect == {0,0,1e9,1e9}) skips them entirely. The
// branch is push-constant-uniform across the dispatch, so it never diverges.
const uint UI_FLAG_CLIP = 0x80000000u;

// ─── standard item structs ──────────────────────────────────────────────
// These match the C++ Crafter::QuadItem / CircleItem / ImageItem / GlyphItem
// byte-for-byte under std430.
struct QuadItem    { vec4 rect; vec4 color; vec4 corners; vec4 outline; };
// rect    = (x, y, w, h) in pixels
// color   = filled body RGBA (premultiplied alpha not assumed)
// corners = per-corner radius in px (TL, TR, BR, BL); 0 = sharp
// outline = (thickness, R, G, B); thickness > 0 paints an outline of given color

struct CircleItem  { vec4 centerRadius; vec4 color; vec4 outline; };
// centerRadius = (cx, cy, radius, _)
// outline.x = thickness (0 = filled), .yzw = outline RGB

struct ImageItem   { vec4 rect; vec4 uv; vec4 tint; uvec4 slots; };
// rect  = (x, y, w, h)
// uv    = (u0, v0, u1, v1) into the source texture
// tint  = multiplied with the sampled color
// slots = (textureHeapSlot, samplerHeapSlot, _, _)

struct GlyphItem   { vec4 rect; vec4 uv; vec4 color; };
// rect  = (x, y, w, h) on screen
// uv    = (u0, v0, u1, v1) into the SDF font atlas
// color = glyph color (alpha modulated by SDF)

// ─── SSBO heap views ────────────────────────────────────────────────────
// One declaration per item type; each shader uses the one matching its
// dispatch. Indexed by hdr.itemBuffer.
layout(UI_HEAP_BUFFERS, std430) readonly buffer UIQuadBuf   { QuadItem   items[]; }   uiQuadHeap[];
layout(UI_HEAP_BUFFERS, std430) readonly buffer UICircleBuf { CircleItem items[]; }   uiCircleHeap[];
layout(UI_HEAP_BUFFERS, std430) readonly buffer UIImageBuf  { ImageItem  items[]; }   uiImageHeap[];
layout(UI_HEAP_BUFFERS, std430) readonly buffer UIGlyphBuf  { GlyphItem  items[]; }   uiGlyphHeap[];

// ─── pixel-tile dispatch model ─────────────────────────────────────────
// Standard shaders dispatch one workgroup per 8×8 screen tile. Each thread
// owns ONE pixel and iterates ALL items in order, accumulating the result
// in a local register, then stores once at the end. This guarantees correct
// z-order within a single dispatch (no inter-workgroup race on imageLoad/
// imageStore) and gives the user simple semantics: "items render in array
// order, later items overdraw earlier ones".
//
// Caller dispatches `(ceil(W/8), ceil(H/8), 1)` — no need to know the max
// item size.

// Returns the screen pixel and validates against the surface and clip rect.
bool uiResolveScreenPixel(UIDispatchHeader hdr, out ivec2 screenPx) {
    uvec2 px = gl_GlobalInvocationID.xy;
    if (px.x >= hdr.surfaceSize.x || px.y >= hdr.surfaceSize.y) return false;
    if ((hdr.flags & UI_FLAG_CLIP) != 0u) {
        if (float(px.x) < hdr.clipRectPx.x || float(px.y) < hdr.clipRectPx.y) return false;
        if (float(px.x) >= hdr.clipRectPx.x + hdr.clipRectPx.z) return false;
        if (float(px.y) >= hdr.clipRectPx.y + hdr.clipRectPx.w) return false;
    }
    screenPx = ivec2(px);
    return true;
}

// Per-category clip test for the fused kernel (issue #47). The fused shader
// loads/stores the destination image once for the whole surface but each
// category carries its own clip rect, so the clip can only gate compositing —
// not the shared load/store. This mirrors the integer-pixel comparison
// uiResolveScreenPixel does (compare the pixel's top-left, not its center), so
// a fused category is pixel-identical to its standalone Dispatch* pass.
bool uiPixelInClipRect(uvec2 px, vec4 clip) {
    if (float(px.x) < clip.x || float(px.y) < clip.y) return false;
    if (float(px.x) >= clip.x + clip.z) return false;
    if (float(px.y) >= clip.y + clip.w) return false;
    return true;
}

// Non-premultiplied "src over dst" blend. Both operands and result are
// straight-alpha vec4. Use this when iterating items in a loop with a local
// accumulator.
vec4 uiBlendOver(vec4 dst, vec4 src) {
    float a = clamp(src.a, 0.0, 1.0);
    vec3 outRGB = mix(dst.rgb, src.rgb, a);
    float outA  = a + dst.a * (1.0 - a);
    return vec4(outRGB, outA);
}

// SDF for a rounded rect with per-corner radius. p is the point relative to
// the rect's center; halfSize is the rect half-extents; r is per-corner
// (TL, TR, BR, BL). Returns signed distance (negative inside).
float uiSdRoundRect(vec2 p, vec2 halfSize, vec4 r) {
    // Pick the radius for the quadrant p is in.
    r.xy = (p.x > 0.0) ? r.zy : r.wx;   // pick TR/BR vs TL/BL
    r.x  = (p.y > 0.0) ? r.x  : r.y;
    vec2 q = abs(p) - halfSize + r.x;
    return min(max(q.x, q.y), 0.0) + length(max(q, 0.0)) - r.x;
}

// ─── SDF text (FontAtlas) ───────────────────────────────────────────────
// FontAtlas keeps each glyph as a signed distance field: stb_truetype writes
// kOnEdgeValue + kPixelDistScale * d into each texel, d being the distance
// from the outline in atlas pixels, positive inside, truncated to a byte and
// clamped to 0..255. These are FontAtlas's two constants in the encoder's own
// units; the FontAtlasDecode test holds them, and the WGSL copy in
// dom-webgpu.js, to FontAtlas's. ui-text and the fused kernel's text phase
// decode through the helpers below, and so should a copy of either.
const float UI_SDF_ON_EDGE          = 128.0;  // FontAtlas::kOnEdgeValue
const float UI_SDF_PIXEL_DIST_SCALE = 32.0;   // FontAtlas::kPixelDistScale

// The widest edge ramp, in atlas pixels, the stored distances can carry. A
// byte holds from (UI_SDF_ON_EDGE - 0.5) / UI_SDF_PIXEL_DIST_SCALE outside
// the outline to (255.5 - UI_SDF_ON_EDGE) / UI_SDF_PIXEL_DIST_SCALE inside,
// about 4 atlas pixels each way, and a ramp wider than twice the nearer end
// never reaches 0 there: below about 4 px every glyph quad's padding, all 0,
// would show as a translucent box. For text that small a ramp narrower than
// one screen pixel is the lesser evil.
const float UI_SDF_MAX_BAND = 2.0 * min(UI_SDF_ON_EDGE - 0.5, 255.5 - UI_SDF_ON_EDGE)
                            / UI_SDF_PIXEL_DIST_SCALE;

// Distance from the outline in atlas pixels, negative inside, for an atlas
// sample (0..1). stb_truetype truncates, so byte b stands for [b, b + 1) of
// the encoded value and b + 0.5 is its unbiased middle: decoding b itself
// pulls every outline 1/64 atlas pixel inwards.
float uiSdfDistance(float sdf) {
    return (UI_SDF_ON_EDGE - (sdf * 255.0 + 0.5)) / UI_SDF_PIXEL_DIST_SCALE;
}

// What a glyph's per-pixel coverage needs, worked out once per glyph:
//   atlasPerScreen  the atlas pixels one screen pixel spans along x and y;
//   coverage        (k, b) of coverage = clamp(k * sdf + b, 0, 1) for a glyph
//                   scaled alike along x and y (within 1%), as ShapeText's
//                   always are: a ramp one screen pixel wide folded into one
//                   multiply-add. (0, 0) for one that isn't, whose ramp is
//                   then worked out per pixel.
struct UIGlyphRamp {
    vec2 atlasPerScreen;
    vec2 coverage;
};

// The atlas's size in texels: queried rather than taken from
// FontAtlas::kAtlasSize, so the two can't drift apart.
vec2 uiGlyphAtlasSize(uint texSlot, uint samplerSlot) {
    return vec2(textureSize(sampler2D(uiTextures[texSlot], uiSamplers[samplerSlot]), 0));
}

// `uv` is the GlyphItem's (u0, v0, u1, v1), `invRectSize` 1 / its (w, h).
UIGlyphRamp uiGlyphRamp(vec4 uv, vec2 invRectSize, vec2 atlasSize) {
    UIGlyphRamp ramp;
    // abs: a glyph turned by swapping its UVs keeps its scale.
    ramp.atlasPerScreen = abs(uv.zw - uv.xy) * atlasSize * invRectSize;
    ramp.coverage = vec2(0.0);
    float lo = min(ramp.atlasPerScreen.x, ramp.atlasPerScreen.y);
    float hi = max(ramp.atlasPerScreen.x, ramp.atlasPerScreen.y);
    if (hi <= lo * 1.01) {
        // 0.5 - uiSdfDistance(sdf) / band, the band one screen pixel.
        float inv = 1.0 / (UI_SDF_PIXEL_DIST_SCALE * clamp(hi, 0.0001, UI_SDF_MAX_BAND));
        ramp.coverage = vec2(255.0 * inv, 0.5 - (UI_SDF_ON_EDGE - 0.5) * inv);
    }
    return ramp;
}

// The glyph's coverage at `uv` in the atlas. The slots must be dynamically
// uniform, as push constants are: there is no nonuniformEXT here, which lets
// the compiler hoist the descriptor load out of the per-pixel glyph loop.
float uiGlyphCoverage(uint texSlot, uint samplerSlot, vec2 uv, UIGlyphRamp ramp, vec2 atlasSize) {
    float sdf = texture(sampler2D(uiTextures[texSlot], uiSamplers[samplerSlot]), uv).r;
    if (ramp.coverage.x != 0.0) return clamp(ramp.coverage.x * sdf + ramp.coverage.y, 0.0, 1.0);

    // Scaled differently along x and y: one screen pixel spans
    // atlasPerScreen.x atlas pixels across a vertical edge and
    // atlasPerScreen.y across a horizontal one, so the ramp's width follows
    // the direction the distance changes in, from one-texel differences. One
    // width for both would be several pixels wide across one of them.
    vec2 texel = 1.0 / atlasSize;
    vec2 gradient = vec2(
        texture(sampler2D(uiTextures[texSlot], uiSamplers[samplerSlot]), uv + vec2(texel.x, 0.0)).r,
        texture(sampler2D(uiTextures[texSlot], uiSamplers[samplerSlot]), uv + vec2(0.0, texel.y)).r) - sdf;
    float len = length(gradient);
    // No gradient where the field is clamped, 4 atlas pixels from any
    // outline, where the coverage is 0 or 1 at any width up to the cap.
    float band = len > 0.0 ? length(gradient / len * ramp.atlasPerScreen)
                           : max(ramp.atlasPerScreen.x, ramp.atlasPerScreen.y);
    return clamp(0.5 - uiSdfDistance(sdf) / clamp(band, 0.0001, UI_SDF_MAX_BAND), 0.0, 1.0);
}

// ─── shared-memory cooperative tile culling ─────────────────────────────
// Each shader still dispatches one 8×8 workgroup per screen tile, but rather
// than have all 64 threads walk the FULL item list — re-reading every item
// from the SSBO 64× and running the per-pixel reject for items that never
// touch the tile — the workgroup streams the list in chunks of 64. Each
// thread loads ONE item, tests its AABB against the whole tile, and the
// survivors are compacted (in original array order) into a shared list that
// every thread then iterates for its own pixel.
//
// This preserves draw order exactly: chunks run in array order and the
// in-chunk compaction is a stable in-order scan, so the per-pixel inner loop
// still sees items in buffer order (later overdraws earlier). It also drops
// SSBO traffic ~64× (each item read once per workgroup, not once per pixel)
// and shrinks the inner loop to only items that actually overlap the tile.
// The inner loop body is left byte-for-byte identical to the unculled path,
// so output is pixel-identical — the cull only decides which items reach it.

const uint UI_TILE_SIZE = 8u;   // matches local_size_x / local_size_y
const uint UI_CHUNK      = 64u; // = workgroup size (8×8); one item per thread

// This workgroup's tile bounds in pixels.
void uiTileBounds(out vec2 tileMin, out vec2 tileMax) {
    tileMin = vec2(gl_WorkGroupID.xy * UI_TILE_SIZE);
    tileMax = tileMin + vec2(float(UI_TILE_SIZE));
}

// Conservative AABB-vs-tile overlap. MUST be a superset of "some pixel in the
// tile passes the in-loop per-pixel reject" so an item a pixel needs is never
// dropped. Pixel centers live in (tileMin, tileMax), so testing the full tile
// box only ever keeps more, never fewer.
bool uiAabbOverlapsTile(vec2 lo, vec2 hi, vec2 tileMin, vec2 tileMax) {
    return lo.x < tileMax.x && hi.x > tileMin.x
        && lo.y < tileMax.y && hi.y > tileMin.y;
}