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) shadow_pos: vec4<f32>,
    @location(5) instance_tint: vec4<f32>,
};

struct Uniforms {
    view: mat4x4<f32>,
    projection: mat4x4<f32>,
    camera_position: vec4<f32>,
    num_lights: vec4<u32>,
    ambient_light: vec4<f32>,
    light_view_projection: mat4x4<f32>,
    shadow_bias: f32,
    shadows_enabled: f32,
    _padding: vec2<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>,
};

const ALPHA_MODE_OPAQUE: u32 = 0u;
const ALPHA_MODE_MASK: u32 = 1u;
const ALPHA_MODE_BLEND: u32 = 2u;

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;

@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;

@group(3) @binding(0)
var shadow_texture: texture_depth_2d;

@group(3) @binding(1)
var shadow_sampler: sampler_comparison;

@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];

    let light_space_pos = uniforms.light_view_projection * world_position;
    out.shadow_pos = light_space_pos;

    return out;
}

fn calculate_shadow_factor(shadow_pos: vec4<f32>) -> f32 {
    if uniforms.shadows_enabled < 0.5 {
        return 1.0;
    }

    let shadow_coords_ndc = shadow_pos.xyz / shadow_pos.w;
    let shadow_coords = vec3<f32>(
        shadow_coords_ndc.xy * vec2<f32>(0.5, -0.5) + vec2<f32>(0.5, 0.5),
        shadow_coords_ndc.z
    );

    let clamped_coords = clamp(shadow_coords, vec3<f32>(0.0), vec3<f32>(1.0));

    let size = f32(textureDimensions(shadow_texture).x);
    let texel_size = 1.0 / size;

    let poisson_disk = array<vec2<f32>, 16>(
        vec2<f32>(-0.94201624, -0.39906216),
        vec2<f32>(0.94558609, -0.76890725),
        vec2<f32>(-0.094184101, -0.92938870),
        vec2<f32>(0.34495938, 0.29387760),
        vec2<f32>(-0.91588581, 0.45771432),
        vec2<f32>(-0.81544232, -0.87912464),
        vec2<f32>(-0.38277543, 0.27676845),
        vec2<f32>(0.97484398, 0.75648379),
        vec2<f32>(0.44323325, -0.97511554),
        vec2<f32>(0.53742981, -0.47373420),
        vec2<f32>(-0.26496911, -0.41893023),
        vec2<f32>(0.79197514, 0.19090188),
        vec2<f32>(-0.24188840, 0.99706507),
        vec2<f32>(-0.81409955, 0.91437590),
        vec2<f32>(0.19984126, 0.78641367),
        vec2<f32>(0.14383161, -0.14100790)
    );

    var visibility = 0.0;
    let filter_radius = 0.5;

    for (var i = 0; i < 16; i = i + 1) {
        let offset = poisson_disk[i] * texel_size * filter_radius;
        visibility = visibility + textureSampleCompare(
            shadow_texture,
            shadow_sampler,
            clamped_coords.xy + offset,
            clamped_coords.z + uniforms.shadow_bias
        );
    }

    let in_bounds = f32(
        shadow_coords.x >= 0.0 && shadow_coords.x <= 1.0 &&
        shadow_coords.y >= 0.0 && shadow_coords.y <= 1.0 &&
        shadow_coords.z >= 0.0 && shadow_coords.z <= 1.0
    );

    return mix(1.0, visibility / 16.0, in_bounds);
}

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;
}

@fragment
fn fs_main(in: VertexOutput) -> @location(0) vec4<f32> {
    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 alpha = base_color.a;

    if material.alpha_mode == ALPHA_MODE_OPAQUE {
        alpha = 1.0;
    } else if material.alpha_mode == ALPHA_MODE_MASK {
        if alpha < material.alpha_cutoff {
            discard;
        }
        alpha = 1.0;
    }

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

    let shadow_factor = calculate_shadow_factor(in.shadow_pos);

    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 {
                let light_contribution = calculate_directional_light(
                    light, normal, view_dir, base_color.rgb, material.roughness, material.metallic
                );
                total_light = total_light + light_contribution * shadow_factor;
            } 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;
    }

    return vec4<f32>(final_color, alpha);
}
