#version 450
layout(location = 0) in vec4 in_color;
layout(location = 1) in vec2 in_uv;
layout(location = 2) in vec2 in_mask_uv;
layout(location = 3) in float in_fog;
layout(location = 0) out vec4 out_color;

layout(set = 0, binding = 0) uniform sampler2D tex_sampler;
layout(set = 0, binding = 1) uniform sampler2D mask_sampler;

layout(push_constant) uniform DrawConstants {
    mat4 mvp;
    vec4 params; // bone base + 1, pretransformed, UI mask, unused
} draw;

layout(set = 1, binding = 1, std430) readonly buffer FixedFunctionState {
    uvec4 stage_color[2];
    uvec4 stage_alpha[2];
    vec4 fog_color;
    vec4 fog_params;
    vec4 texture_factor;
    mat4 world_view;
    uvec4 flags;
    mat4 normal_matrix;
    uvec4 lighting_flags;
    uvec4 material_sources;
    uvec4 alpha_test; // ALPHATESTENABLE, ALPHAFUNC, ALPHAREF, unused
} fixed_state;

vec4 stage_arg(uint selector, vec4 diffuse, vec4 current, vec4 texel) {
    uint source = selector & 15u;
    vec4 value = source == 0u ? diffuse :
                 source == 1u ? current :
                 source == 2u ? texel :
                 source == 3u ? fixed_state.texture_factor : vec4(1.0);
    if ((selector & 16u) != 0u) value = vec4(1.0) - value;
    if ((selector & 32u) != 0u) value.rgb = vec3(value.a);
    return value;
}

vec4 apply_op(uint op, vec4 a, vec4 b, vec4 current, vec4 diffuse, vec4 texel) {
    if (op == 2u) return a;                          // SELECTARG1
    if (op == 3u) return b;                          // SELECTARG2
    if (op == 4u) return a * b;                      // MODULATE
    if (op == 5u) return 2.0 * a * b;                // MODULATE2X
    if (op == 6u) return 4.0 * a * b;                // MODULATE4X
    if (op == 7u) return a + b;                      // ADD
    if (op == 8u) return a + b - vec4(0.5);         // ADDSIGNED
    if (op == 9u) return 2.0 * (a + b - vec4(0.5)); // ADDSIGNED2X
    if (op == 10u) return a - b;                     // SUBTRACT
    if (op == 11u) return a + b - a * b;             // ADDSMOOTH
    if (op == 12u) return mix(b, a, diffuse.a);
    if (op == 13u) return mix(b, a, texel.a);
    if (op == 14u) return mix(b, a, fixed_state.texture_factor.a);
    if (op == 15u) return a + b * (1.0 - texel.a);
    if (op == 16u) return mix(b, a, current.a);
    if (op == 18u) return vec4(a.rgb + a.a * b.rgb, a.a);
    if (op == 19u) return vec4(a.rgb * b.rgb + vec3(a.a), a.a);
    if (op == 20u) return vec4((1.0 - a.a) * b.rgb + a.rgb, a.a);
    if (op == 21u) return vec4((vec3(1.0) - a.rgb) * b.rgb + vec3(a.a), a.a);
    return current;
}

vec4 apply_stage(uvec4 color_state, uvec4 alpha_state,
                 vec4 diffuse, vec4 current, vec4 texel) {
    if (color_state.x <= 1u) return current;
    vec4 color_arg1 = stage_arg(color_state.y, diffuse, current, texel);
    vec4 color_arg2 = stage_arg(color_state.z, diffuse, current, texel);
    vec4 result = apply_op(color_state.x, color_arg1, color_arg2, current, diffuse, texel);
    if (alpha_state.x > 1u) {
        vec4 alpha_arg1 = stage_arg(alpha_state.y, diffuse, current, texel);
        vec4 alpha_arg2 = stage_arg(alpha_state.z, diffuse, current, texel);
        result.a = apply_op(alpha_state.x, alpha_arg1, alpha_arg2, current, diffuse, texel).a;
    } else if (alpha_state.x == 1u) {
        result.a = current.a;
    }
    return clamp(result, 0.0, 1.0);
}

void main() {
    if (fixed_state.flags.x != 0u) {
        // D3D8 texture stage cascade: stage 1 runs unless its COLOROP is D3DTOP_DISABLE.
        vec4 diffuse = in_color;
        vec4 color = apply_stage(fixed_state.stage_color[0], fixed_state.stage_alpha[0],
                                 diffuse, diffuse, texture(tex_sampler, in_uv));
        if (fixed_state.stage_color[0].x > 1u && fixed_state.stage_color[1].x > 1u)
            color = apply_stage(fixed_state.stage_color[1], fixed_state.stage_alpha[1],
                                diffuse, color, texture(mask_sampler, in_mask_uv));
        if (fixed_state.alpha_test.x != 0u) {
            // D3DRS_ALPHAREF is an 8-bit reference compared against the 8-bit pixel alpha.
            uint a = uint(clamp(color.a, 0.0, 1.0) * 255.0 + 0.5);
            uint ref = fixed_state.alpha_test.z & 255u;
            uint func = fixed_state.alpha_test.y;
            bool passes = func == 1u ? false :
                          func == 2u ? a < ref :
                          func == 3u ? a == ref :
                          func == 4u ? a <= ref :
                          func == 5u ? a > ref :
                          func == 6u ? a != ref :
                          func == 7u ? a >= ref : true;
            if (!passes) discard;
        }
        color.rgb = mix(fixed_state.fog_color.rgb, color.rgb, in_fog);
        out_color = color;
        return;
    }
    // 2D UI batch (UIRenderCommand): vertex color times texture, optionally clipped by a mask.
    // The UI is laid out in logical pixels and magnified onto a HiDPI swapchain. Plain bilinear
    // smears each texel over the scale factor; "sharp bilinear" keeps texels square and only blends
    // across a one-pixel seam. Minified or 1:1 texels keep the original bilinear lookup.
    vec2 uv = in_uv;
    vec2 tex_size = vec2(textureSize(tex_sampler, 0));
    vec2 texel = uv * tex_size;
    vec2 texels_per_pixel = fwidth(texel);
    if (texels_per_pixel.x > 0.0 && texels_per_pixel.y > 0.0 &&
        texels_per_pixel.x < 0.99 && texels_per_pixel.y < 0.99) {
        vec2 seam = floor(texel + 0.5);
        texel = seam + clamp((texel - seam) / texels_per_pixel, -0.5, 0.5);
        uv = texel / tex_size;
    }
    vec4 color = in_color * texture(tex_sampler, uv);
    if (draw.params.z > 0.5) {
        if (in_mask_uv.x < 0.0 || in_mask_uv.x > 1.0 || in_mask_uv.y < 0.0 || in_mask_uv.y > 1.0)
            discard;
        vec4 mask_val = texture(mask_sampler, in_mask_uv);
        color.rgb *= mask_val.rgb;
        color.a = mask_val.a * in_color.a;
    }
    out_color = color;
}