#language 2026

#include "shared.slang"
#include "wgsl_workaround.slang"

struct VertexShaderInput {
    uint32_t vertex_index   : SV_VulkanVertexID;
    uint32_t instance_index : SV_VulkanInstanceID;

    // Per Instance
    [vk::location(0)] uint32_t2 pos;
    [vk::location(1)] uint32_t  kind;
};

struct VertexShaderOutput {
    float4 pos : SV_Position;

    float2   uv;
    uint32_t kind;
};

struct FragmentShaderOutput {
    float4 color : SV_Target;
};

static const float2 offset = float2(0, 0.25); 

[shader("vertex")]
VertexShaderOutput main_vertex(VertexShaderInput input) {
    var output: VertexShaderOutput;

    var vertex_pos = float2(0, 0);
    var vertex_uv  = float2(0, 0);

    switch (input.vertex_index) {
        case 0: { vertex_pos = float2(-1.0,  1.0); vertex_uv = float2(0, 0); } break;
        case 1: { vertex_pos = float2(-1.0, -1.0); vertex_uv = float2(0, 1); } break;
        case 2: { vertex_pos = float2( 1.0, -1.0); vertex_uv = float2(1, 1); } break;
        case 3: { vertex_pos = float2(-1.0,  1.0); vertex_uv = float2(0, 0); } break;
        case 4: { vertex_pos = float2( 1.0, -1.0); vertex_uv = float2(1, 1); } break;
        case 5: { vertex_pos = float2( 1.0,  1.0); vertex_uv = float2(1, 0); } break;
        default: {}
    }

    vertex_pos += offset;

    output.pos  = mul(float4(input.pos + vertex_pos, 0, 1), view_projection_matrix);
    output.uv   = vertex_uv;
    output.kind = input.kind;

    return output;
}

struct Block2 {
    CombinedTextureSampler   tree_atlas;
    StructuredBuffer<float4> tree_uvs;
};

struct Block3 {
    ConstantBuffer<float3> tint;
};

[vk::binding(0, 2)] ParameterBlock<Block2> block2;
[vk::binding(0, 3)]ParameterBlock<Block3> block3;

[shader("pixel")]
FragmentShaderOutput main_fragment(VertexShaderOutput input) {
    var output: FragmentShaderOutput;

    output.color = pixel_art_sample(block2.tree_atlas, input.uv, block2.tree_uvs[input.kind]);
    output.color = float4(output.color.rgb * block3.tint.rgb, output.color.a);

    return output;
}

float4 pixel_art_sample(CombinedTextureSampler input_texture, float2 input_uv, float4 tree_uv) {
    var dimensions : float2;
    input_texture.getTexture().GetDimensions(dimensions.x, dimensions.y);

    let texture_uv = lerp(tree_uv.xy, tree_uv.zw, input_uv);
    let sample_uv = (floor(texture_uv) + saturate(fract(texture_uv) / fwidth(texture_uv)) - 0.5) / dimensions;

    let uv = clamp(sample_uv, (tree_uv.xy + 0.5) / dimensions, (tree_uv.zw - 0.5) / dimensions);

    return input_texture.Sample(uv);
}