#version 450
layout(location = 0) in vec3 in_position;
layout(location = 1) in vec3 in_normal;
layout(location = 2) in vec2 in_uv;
layout(location = 3) in vec4 in_color;
layout(location = 4) in uvec4 in_joints;
layout(location = 5) in vec4 in_weights;
layout(location = 6) in vec2 in_mask_uv;
layout(location = 7) in float in_rhw;

layout(location = 0) out vec4 out_color;
layout(location = 1) out vec2 out_uv;
layout(location = 2) out vec2 out_mask_uv;
layout(location = 3) out float out_fog;

layout(set = 1, binding = 0, std430) readonly buffer BonePalette {
    mat4 bones[];
} bone_palette;

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;
    mat4 world;
    uvec4 lighting_flags;
    vec4 material_ambient;
    vec4 material_emissive;
    vec4 global_ambient;
    vec4 light_position_type[2];
    vec4 light_direction_range[2];
    vec4 light_diffuse[2];
    vec4 light_ambient[2];
    vec4 light_attenuation[2];
    vec4 light_spot[2];
} fixed_state;

void main() {
    vec3 pos = in_position;
    vec3 nrm = in_normal;
    if (draw.light_diffuse.w > 0.5) {
        uint bone_base = uint(draw.light_diffuse.w - 0.5);
        vec3 skinned_pos = vec3(0.0);
        vec3 skinned_nrm = vec3(0.0);
        float total_w = 0.0;
        for (int k = 0; k < 4; ++k) {
            float w = in_weights[k];
            if (w > 0.0) {
                mat4 B = bone_palette.bones[bone_base + in_joints[k]];
                skinned_pos += w * (B * vec4(pos, 1.0)).xyz;
                skinned_nrm += w * (mat3(B) * nrm);
                total_w += w;
            }
        }
        if (total_w > 0.0) {
            pos = skinned_pos;
            nrm = skinned_nrm;
        }
    }
    // D3D row-vector matrices are uploaded row-major. GLSL reads the bytes as their transpose.
    gl_Position = draw.mvp * vec4(pos, 1.0);
    if (draw.light_diffuse.w < -0.5) {
        float clip_w = in_rhw > 0.0 ? 1.0 / in_rhw : 1.0;
        gl_Position.xyz *= clip_w;
        gl_Position.w = clip_w;
    }
    gl_Position.y = -gl_Position.y;
    vec3 n = length(nrm) > 1e-4 ? normalize(nrm) : vec3(0.0, 0.0, 1.0);
    out_color = in_color;
    if (fixed_state.lighting_flags.x != 0u) {
        vec3 world_pos = (fixed_state.world * vec4(pos, 1.0)).xyz;
        vec3 world_normal = normalize(mat3(fixed_state.world) * n);
        vec3 mat_diffuse = (fixed_state.lighting_flags.y != 0u && fixed_state.lighting_flags.w == 1u)
            ? in_color.rgb : draw.tint_color.rgb;
        vec3 mat_ambient = (fixed_state.lighting_flags.y != 0u && fixed_state.lighting_flags.z == 1u)
            ? in_color.rgb : fixed_state.material_ambient.rgb;
        vec3 ambient_sum = fixed_state.global_ambient.rgb;
        vec3 diffuse_sum = vec3(0.0);
        for (int i = 0; i < 2; ++i) {
            int type = int(fixed_state.light_position_type[i].w + 0.5);
            if (type == 0) continue;
            vec3 L;
            float strength = 1.0;
            if (type == 3) {
                L = normalize(-fixed_state.light_direction_range[i].xyz);
            } else {
                vec3 to_light = fixed_state.light_position_type[i].xyz - world_pos;
                float distance_to_light = length(to_light);
                if (distance_to_light > fixed_state.light_direction_range[i].w || distance_to_light < 0.0001)
                    continue;
                L = to_light / distance_to_light;
                vec4 attenuation = fixed_state.light_attenuation[i];
                float denominator = attenuation.x + attenuation.y * distance_to_light +
                                    attenuation.z * distance_to_light * distance_to_light;
                strength = denominator > 0.0001 ? min(1.0 / denominator, 1.0) : 1.0;
                if (type == 2) {
                    vec3 spot_dir = normalize(fixed_state.light_direction_range[i].xyz);
                    float cosine = dot(-L, spot_dir);
                    float inner = cos(fixed_state.light_spot[i].x * 0.5);
                    float outer = cos(fixed_state.light_spot[i].y * 0.5);
                    float cone = clamp((cosine - outer) / max(inner - outer, 0.0001), 0.0, 1.0);
                    strength *= pow(cone, max(attenuation.w, 0.0001));
                }
            }
            ambient_sum += fixed_state.light_ambient[i].rgb * strength;
            diffuse_sum += fixed_state.light_diffuse[i].rgb * max(dot(world_normal, L), 0.0) * strength;
        }
        vec3 lit = fixed_state.material_emissive.rgb + mat_ambient * ambient_sum + mat_diffuse * diffuse_sum;
        float alpha = (fixed_state.lighting_flags.y != 0u && fixed_state.lighting_flags.w == 1u)
            ? in_color.a : draw.tint_color.a;
        out_color = vec4(clamp(lit, 0.0, 1.0), alpha);
    }
    out_uv = in_uv;
    if (draw.light_dir.w > 1.001) {
        vec3 view_normal = normalize(mat3(fixed_state.world_view) * n);
        vec3 view_pos = (fixed_state.world_view * vec4(pos, 1.0)).xyz;
        vec3 view_dir = normalize(-view_pos);
        vec3 reflected = reflect(-view_dir, view_normal);
        out_mask_uv = reflected.xy * vec2(0.5, -0.5) + vec2(0.5);
    } else {
        out_mask_uv = in_mask_uv;
    }
    out_fog = 1.0;
    if (fixed_state.flags.x != 0u && fixed_state.flags.w != 0u && draw.light_diffuse.w >= -0.5) {
        vec3 eye = (fixed_state.world_view * vec4(pos, 1.0)).xyz;
        float distance_to_eye = fixed_state.fog_params.w > 0.5 ? length(eye) : abs(eye.z);
        if (fixed_state.flags.w == 3u) {
            float span = fixed_state.fog_params.y - fixed_state.fog_params.x;
            if (span > 0.0001)
                out_fog = clamp((fixed_state.fog_params.y - distance_to_eye) / span, 0.0, 1.0);
        } else if (fixed_state.flags.w == 1u) {
            out_fog = clamp(exp(-fixed_state.fog_params.z * distance_to_eye), 0.0, 1.0);
        } else if (fixed_state.flags.w == 2u) {
            float fog_distance = fixed_state.fog_params.z * distance_to_eye;
            out_fog = clamp(exp(-fog_distance * fog_distance), 0.0, 1.0);
        }
    }
}