const PI: f32 = 3.14159265358979323846;

fn ggxDistribution(dot_nh: f32, roughness: f32) -> f32 {
    let a = roughness * roughness;
    let a2 = a * a;
    let dot_nh2 = dot_nh * dot_nh;
    let nom = a2;
    let denom = (dot_nh2 * (a2 - 1.0) + 1.0);
    return nom / (PI * denom * denom);
}

fn geomSmith(dot_prod: f32, roughness: f32) -> f32 {
    let r = (roughness + 1.0);
    let k = (r * r) / 8.0;
    let nom = dot_prod;
    let denom = dot_prod * (1.0 - k) + k;
    return nom / denom;
}

fn schlickFresnel(cos_theta: f32, f0: vec3<f32>) -> vec3<f32> {
    return f0 + (1.0 - f0) * pow(max(1.0 - cos_theta, 0.0), 5.0);
}

fn cookTorranceBRDF(
    light_dir: vec3<f32>,
    view_dir: vec3<f32>,
    normal: vec3<f32>,
    albedo: vec3<f32>,
    roughness: f32,
    metallic: f32
) -> vec3<f32> {
    let h = normalize(view_dir + light_dir);

    let dot_nv = max(dot(normal, view_dir), 0.0);
    let dot_nl = max(dot(normal, light_dir), 0.0);
    let dot_nh = max(dot(normal, h), 0.0);
    let dot_hv = max(dot(h, view_dir), 0.0);

    var f0 = vec3(0.04);
    f0 = mix(f0, albedo, metallic);

    let ndf = ggxDistribution(dot_nh, roughness);
    let g = geomSmith(dot_nv, roughness) * geomSmith(dot_nl, roughness);
    let f = schlickFresnel(dot_hv, f0);

    let numerator = ndf * g * f;
    let denominator = 4.0 * dot_nv * dot_nl + 0.001;
    let specular = numerator / denominator;

    let ks = f;
    var kd = vec3(1.0) - ks;
    kd *= 1.0 - metallic;

    let diffuse = kd * albedo / PI;

    return diffuse + specular;
}

struct VertexInput {
    @location(0) position: vec3<f32>,
    @location(1) normal: vec3<f32>,
    @location(2) tex_coords: vec2<f32>,
};

struct VertexOutput {
    @builtin(position) position: vec4<f32>,
    @location(0) normal: vec3<f32>,
    @location(1) world_pos: vec3<f32>,
    @location(2) @interpolate(flat) material_id: u32,
    @location(3) tex_coords: vec2<f32>,
    @location(4) instance_tint: vec4<f32>,
};

struct Uniforms {
    view: mat4x4<f32>,
    projection: mat4x4<f32>,
    camera_position: vec4<f32>,
    num_lights: vec4<u32>,
    ambient_light: vec4<f32>,
};

struct ModelMatrix {
    model: mat4x4<f32>,
    normal_matrix: mat3x3<f32>,
};

struct Material {
    base_color: vec4<f32>,
    emissive_factor: vec3<f32>,
    alpha_mode: u32,
    alpha_cutoff: f32,
    has_base_texture: u32,
    has_emissive_texture: u32,
    roughness: f32,
    metallic: f32,
    unlit: u32,
    _padding: vec2<f32>,
};

struct ObjectData {
    transform_index: u32,
    mesh_id: u32,
    material_id: u32,
    batch_id: u32,
};

struct Light {
    position: vec4<f32>,
    direction: vec4<f32>,
    color: vec4<f32>,
    light_type: u32,
    range: f32,
    inner_cone: f32,
    outer_cone: f32,
};

const LIGHT_TYPE_DIRECTIONAL: u32 = 0u;
const LIGHT_TYPE_POINT: u32 = 1u;
const LIGHT_TYPE_SPOT: u32 = 2u;

struct OITOutput {
    @location(0) accum: vec4<f32>,
    @location(1) reveal: f32,
};

@group(0) @binding(0)
var<uniform> uniforms: Uniforms;

@group(1) @binding(0)
var<storage, read> transforms: array<ModelMatrix>;

@group(1) @binding(1)
var<storage, read> materials: array<Material>;

@group(1) @binding(2)
var<storage, read> objects: array<ObjectData>;

@group(1) @binding(3)
var<storage, read> lights: array<Light>;

@group(1) @binding(4)
var<storage, read> visible_indices: array<u32>;

@group(1) @binding(5)
var<storage, read> instance_custom_data: array<vec4<f32>>;

@group(2) @binding(0)
var base_texture: texture_2d<f32>;

@group(2) @binding(1)
var base_sampler: sampler;

@group(2) @binding(2)
var emissive_texture: texture_2d<f32>;

@group(2) @binding(3)
var emissive_sampler: sampler;

@vertex
fn vs_main(
    in: VertexInput,
    @builtin(instance_index) instance_id: u32
) -> VertexOutput {
    var out: VertexOutput;

    let object_index = visible_indices[instance_id];
    let object = objects[object_index];

    let transform = transforms[object.transform_index];
    let model_matrix = transform.model;
    let normal_matrix = transform.normal_matrix;

    let world_position = model_matrix * vec4<f32>(in.position, 1.0);
    out.world_pos = world_position.xyz;
    out.position = uniforms.projection * uniforms.view * world_position;
    out.normal = normalize(normal_matrix * in.normal);
    out.material_id = object.material_id;
    out.tex_coords = in.tex_coords;
    out.instance_tint = instance_custom_data[object.transform_index];

    return out;
}

fn calculate_directional_light(
    light: Light,
    normal: vec3<f32>,
    view_dir: vec3<f32>,
    albedo: vec3<f32>,
    roughness: f32,
    metallic: f32
) -> vec3<f32> {
    let light_dir = normalize(-light.direction.xyz);
    let brdf = cookTorranceBRDF(light_dir, view_dir, normal, albedo, roughness, metallic);
    let ndotl = max(dot(normal, light_dir), 0.0);
    return light.color.rgb * brdf * ndotl;
}

fn calculate_point_light(
    light: Light,
    world_pos: vec3<f32>,
    normal: vec3<f32>,
    view_dir: vec3<f32>,
    albedo: vec3<f32>,
    roughness: f32,
    metallic: f32
) -> vec3<f32> {
    let light_vec = light.position.xyz - world_pos;
    let distance = length(light_vec);

    if distance > light.range {
        return vec3<f32>(0.0, 0.0, 0.0);
    }

    let light_dir = normalize(light_vec);
    let brdf = cookTorranceBRDF(light_dir, view_dir, normal, albedo, roughness, metallic);
    let ndotl = max(dot(normal, light_dir), 0.0);

    let attenuation = 1.0 / (1.0 + 0.09 * distance + 0.032 * distance * distance);
    let fade = 1.0 - smoothstep(light.range * 0.75, light.range, distance);

    return light.color.rgb * brdf * ndotl * attenuation * fade;
}

fn calculate_spot_light(
    light: Light,
    world_pos: vec3<f32>,
    normal: vec3<f32>,
    view_dir: vec3<f32>,
    albedo: vec3<f32>,
    roughness: f32,
    metallic: f32
) -> vec3<f32> {
    let light_vec = light.position.xyz - world_pos;
    let distance = length(light_vec);

    if distance > light.range {
        return vec3<f32>(0.0, 0.0, 0.0);
    }

    let light_dir = normalize(light_vec);
    let spot_dir = normalize(light.direction.xyz);
    let spot_cos = dot(-light_dir, spot_dir);

    if spot_cos < light.outer_cone {
        return vec3<f32>(0.0, 0.0, 0.0);
    }

    let brdf = cookTorranceBRDF(light_dir, view_dir, normal, albedo, roughness, metallic);
    let ndotl = max(dot(normal, light_dir), 0.0);

    let spot_intensity = smoothstep(light.outer_cone, light.inner_cone, spot_cos);
    let attenuation = 1.0 / (1.0 + 0.09 * distance + 0.032 * distance * distance);
    let fade = 1.0 - smoothstep(light.range * 0.75, light.range, distance);

    return light.color.rgb * brdf * ndotl * spot_intensity * attenuation * fade;
}

fn weight(z: f32, a: f32) -> f32 {
    let z_scale = 5.0;
    let z_power = 200.0;
    let max_weight = 3000.0;
    return a * max(0.01, min(max_weight, 10.0 / (1e-5 + pow(abs(z) / z_scale, z_power))));
}

@fragment
fn fs_main(in: VertexOutput) -> OITOutput {
    let normal = normalize(in.normal);
    let material = materials[in.material_id];

    let base_tex_sample = textureSample(base_texture, base_sampler, in.tex_coords);
    let emissive_tex_sample = textureSample(emissive_texture, emissive_sampler, in.tex_coords);

    var base_color = material.base_color * in.instance_tint;
    if material.has_base_texture != 0u {
        base_color = base_color * base_tex_sample;
    }

    var emissive = material.emissive_factor;
    if material.has_emissive_texture != 0u {
        emissive = emissive * emissive_tex_sample.rgb;
    }

    var final_color: vec3<f32>;

    if material.unlit != 0u {
        final_color = base_color.rgb + emissive;
    } else {
        let view_dir = normalize(uniforms.camera_position.xyz - in.world_pos);
        var total_light = vec3<f32>(0.0, 0.0, 0.0);

        for (var i = 0u; i < uniforms.num_lights.x; i = i + 1u) {
            let light = lights[i];

            if light.light_type == LIGHT_TYPE_DIRECTIONAL {
                total_light = total_light + calculate_directional_light(
                    light, normal, view_dir, base_color.rgb, material.roughness, material.metallic
                );
            } else if light.light_type == LIGHT_TYPE_POINT {
                total_light = total_light + calculate_point_light(
                    light, in.world_pos, normal, view_dir, base_color.rgb, material.roughness, material.metallic
                );
            } else if light.light_type == LIGHT_TYPE_SPOT {
                total_light = total_light + calculate_spot_light(
                    light, in.world_pos, normal, view_dir, base_color.rgb, material.roughness, material.metallic
                );
            }
        }

        total_light = total_light + uniforms.ambient_light.xyz * base_color.rgb;
        final_color = total_light + emissive;
    }

    let alpha = base_color.a;
    let w = weight(in.position.z, alpha);

    var out: OITOutput;
    out.accum = vec4<f32>(final_color, 1.0) * w;
    out.reveal = 1.0 - alpha;

    return out;
}
