#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 tint_color;
    vec4 ambient_emissive;
    vec4 light_dir;
    vec4 light_diffuse;
} draw;

layout(set = 1, binding = 1, std430) readonly buffer FixedFunctionState {
    uvec4 stage0_color;
    uvec4 stage0_alpha;
    uvec4 stage1_color;
    uvec4 stage1_alpha;
    vec4 fog_color;
    vec4 fog_params;
    vec4 texture_factor;
    mat4 world_view;
    uvec4 flags;
} 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() {
    vec4 tex0 = texture(tex_sampler, in_uv);
    if (fixed_state.flags.x == 2u) {
        vec3 shadow = tex0.rgb;
        if (fixed_state.flags.z != 0u && fixed_state.stage1_color.x == 4u)
            shadow *= texture(mask_sampler, in_mask_uv).rgb;
        shadow = mix(fixed_state.fog_color.rgb, shadow, in_fog);
        if (all(greaterThanEqual(shadow, vec3(0.997)))) discard;
        out_color = vec4(shadow, 1.0);
        return;
    }
    if (fixed_state.flags.x != 0u) {
        vec4 diffuse = in_color;
        vec4 color = apply_stage(fixed_state.stage0_color, fixed_state.stage0_alpha,
                                 diffuse, diffuse, tex0);
        if (fixed_state.flags.z != 0u && fixed_state.stage1_color.x > 1u) {
            vec4 tex1 = texture(mask_sampler, in_mask_uv);
            color = apply_stage(fixed_state.stage1_color, fixed_state.stage1_alpha,
                                diffuse, color, tex1);
        }
        if (fixed_state.stage1_alpha.w != 0u) {
            float ref = draw.ambient_emissive.a;
            uint func = fixed_state.stage0_alpha.w;
            bool passes = func == 1u ? false :
                          func == 2u ? color.a < ref :
                          func == 3u ? abs(color.a - ref) < (0.5 / 255.0) :
                          func == 4u ? color.a <= ref :
                          func == 5u ? color.a > ref :
                          func == 6u ? abs(color.a - ref) >= (0.5 / 255.0) :
                          func == 7u ? color.a >= ref : true;
            if (!passes) discard;
        }
        color.rgb = mix(fixed_state.fog_color.rgb, color.rgb, in_fog);
        out_color = color;
        return;
    }
    vec4 color = in_color * tex0;
    if (draw.light_dir.w < -0.4) {
        vec4 mask_val = texture(mask_sampler, in_mask_uv);
        if (draw.light_dir.w < -0.8) {
            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;
            }
            color.rgb *= mask_val.rgb;
            color.a = mask_val.a * in_color.a;
        } else {
            color.a = mask_val.a * in_color.a;
            if (color.a <= 0.003) discard;
        }
    } else if (draw.light_dir.w > 1.001) {
        // Specular sphere-map stage 1 (D3DTOP_MODULATEALPHA_ADDCOLOR)
        float spec_power = draw.light_dir.w - 1.0;
        vec4 spec_map = texture(mask_sampler, in_mask_uv);
        color.rgb = min(color.rgb + (tex0.a * spec_power) * spec_map.rgb, vec3(1.5));
        color.a = 1.0;
    }
    if (color.a <= draw.ambient_emissive.a) discard;
    out_color = color;
}