#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 = 8) in float in_vertex_fog;

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 params; // bone base + 1 (0 = none), pretransformed, UI mask, unused
} draw;

// D3D8 fixed-function vertex state; see FixedFunctionState in main.cpp. Lights are in camera space.
layout(set = 1, binding = 1, std430) readonly buffer FixedFunctionState {
    uvec4 stage_color[2]; // op, arg1, arg2, D3DTSS_TEXCOORDINDEX
    uvec4 stage_alpha[2]; // op, arg1, arg2, D3DTSS_TEXTURETRANSFORMFLAGS
    vec4 fog_color;
    vec4 fog_params;
    vec4 texture_factor;
    mat4 world_view;
    uvec4 flags;
    mat4 normal_matrix;
    uvec4 lighting_flags;   // LIGHTING, COLORVERTEX, vertex has diffuse, NORMALIZENORMALS
    uvec4 material_sources; // DIFFUSE, AMBIENT, EMISSIVE material source, LOCALVIEWER
    uvec4 alpha_test;
    vec4 material_diffuse;
    vec4 material_ambient;
    vec4 material_emissive;
    vec4 global_ambient;
    mat4 texture_matrix[2];
    vec4 light_position_type[8];
    vec4 light_direction_range[8];
    vec4 light_diffuse[8];
    vec4 light_ambient[8];
    vec4 light_attenuation[8];
    vec4 light_spot[8];
} fixed_state;

// D3DMCS_COLOR1 picks the vertex diffuse color when COLORVERTEX is on and the vertex has one.
vec4 material_color(uint source, vec4 material) {
    if (fixed_state.lighting_flags.y != 0u && fixed_state.lighting_flags.z != 0u && source == 1u)
        return in_color;
    return material;
}

// D3D8 texture coordinate processing: D3DTSS_TEXCOORDINDEX generation, then D3DTS_TEXTUREn.
vec2 stage_texcoord(int stage, bool pretransformed, vec3 eye, vec3 n) {
    uint tci = fixed_state.stage_color[stage].w;
    uint gen = tci & 0xFFFF0000u;
    uint set_index = tci & 0xFFFFu;
    vec4 coord = vec4(0.0, 0.0, 1.0, 0.0);
    if (gen == 0u || pretransformed) {
        if (set_index == 0u) coord.xy = in_uv;
        else if (set_index == 1u) coord.xy = in_mask_uv;
    } else if (gen == 0x10000u) {          // CAMERASPACENORMAL
        coord = vec4(n, 1.0);
    } else if (gen == 0x20000u) {          // CAMERASPACEPOSITION
        coord = vec4(eye, 1.0);
    } else if (gen == 0x30000u) {          // CAMERASPACEREFLECTIONVECTOR
        vec3 e = fixed_state.material_sources.w != 0u
            ? (length(eye) > 0.0 ? -normalize(eye) : vec3(0.0)) : vec3(0.0, 0.0, -1.0);
        coord = vec4(2.0 * dot(e, n) * n - e, 1.0);
    }
    // Pretransformed (XYZRHW) vertices bypass D3D8 T&L, texture transforms included
    // (RecordingDevice's CPU texcoord path does the same).
    uint transform_flags = pretransformed ? 0u : fixed_state.stage_alpha[stage].w;
    uint elements = transform_flags & 0xFFu;
    vec4 r = coord;
    if (elements != 0u)
        r = fixed_state.texture_matrix[stage] * coord;
    if ((transform_flags & 0x100u) != 0u && elements >= 2u && elements <= 4u && r[elements - 1u] != 0.0)
        r.xy /= r[elements - 1u];
    return r.xy;
}

void main() {
    vec3 pos = in_position;
    vec3 nrm = in_normal;
    if (draw.params.x > 0.5) {
        uint bone_base = uint(draw.params.x - 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;
        }
    }
    bool pretransformed = draw.params.y > 0.5;
    // D3D row-vector matrices are uploaded row-major. GLSL reads the bytes as their transpose.
    gl_Position = draw.mvp * vec4(pos, 1.0);
    if (pretransformed) {
        // Vertices behind the eye carry a negative rhw (STP terrain): rebuilding clip space
        // with w = 1/rhw < 0 lets the clipper drop them instead of drawing w = 1 spikes.
        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;

    // A vertex without D3DFVF_DIFFUSE carries opaque white (upload_geometry).
    out_color = in_color;
    if (fixed_state.flags.x == 0u) {
        // 2D UI batch.
        out_uv = in_uv;
        out_mask_uv = in_mask_uv;
        out_fog = 1.0;
        return;
    }

    vec3 eye = (fixed_state.world_view * vec4(pos, 1.0)).xyz;
    vec3 n = mat3(fixed_state.normal_matrix) * nrm;
    if (fixed_state.lighting_flags.w != 0u && length(n) > 0.0)
        n = normalize(n);

    if (fixed_state.lighting_flags.x != 0u && !pretransformed) {
        vec4 mat_diffuse = material_color(fixed_state.material_sources.x, fixed_state.material_diffuse);
        vec4 mat_ambient = material_color(fixed_state.material_sources.y, fixed_state.material_ambient);
        vec4 mat_emissive = material_color(fixed_state.material_sources.z, fixed_state.material_emissive);
        vec3 ambient_sum = fixed_state.global_ambient.rgb;
        vec3 diffuse_sum = vec3(0.0);
        for (int i = 0; i < 8; ++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) { // D3DLIGHT_DIRECTIONAL
                L = -fixed_state.light_direction_range[i].xyz;
            } else {         // D3DLIGHT_POINT / D3DLIGHT_SPOT
                vec3 to_light = fixed_state.light_position_type[i].xyz - eye;
                float d = length(to_light);
                if (d > fixed_state.light_direction_range[i].w)
                    continue;
                L = d > 0.0 ? to_light / d : vec3(0.0);
                vec4 attenuation = fixed_state.light_attenuation[i];
                float denominator = attenuation.x + attenuation.y * d + attenuation.z * d * d;
                strength = denominator != 0.0 ? 1.0 / denominator : 1.0;
                if (type == 2) {
                    float rho = dot(-L, fixed_state.light_direction_range[i].xyz);
                    float cos_theta = fixed_state.light_spot[i].x;
                    float cos_phi = fixed_state.light_spot[i].y;
                    float spot = rho > cos_theta ? 1.0 : (rho <= cos_phi ? 0.0 :
                        pow((rho - cos_phi) / (cos_theta - cos_phi), attenuation.w));
                    strength *= spot;
                }
            }
            ambient_sum += fixed_state.light_ambient[i].rgb * strength;
            diffuse_sum += fixed_state.light_diffuse[i].rgb * max(dot(n, L), 0.0) * strength;
        }
        vec3 lit = mat_emissive.rgb + mat_ambient.rgb * ambient_sum + mat_diffuse.rgb * diffuse_sum;
        out_color = vec4(clamp(lit, 0.0, 1.0), clamp(mat_diffuse.a, 0.0, 1.0));
    }

    out_uv = stage_texcoord(0, pretransformed, eye, n);
    out_mask_uv = stage_texcoord(1, pretransformed, eye, n);

    out_fog = 1.0;
    if (fixed_state.flags.w == 4u) {
        // XYZRHW with FOGTABLEMODE NONE: D3D8 takes the vertex fog from the specular alpha.
        out_fog = in_vertex_fog;
    } else if (fixed_state.flags.w != 0u) {
        // Table fog on XYZRHW vertices runs on eye depth w = 1/rhw.
        float distance_to_eye = pretransformed ? gl_Position.w
            : (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);
        }
    }
}