const PI: f32 = 3.14159265358979323846;

fn srgb_to_linear(srgb: vec4<f32>) -> vec4<f32> {
    return vec4<f32>(pow(srgb.xyz, vec3<f32>(2.2)), srgb.w);
}

fn srgb_to_linear3(srgb: vec3<f32>) -> vec3<f32> {
    return pow(srgb, vec3<f32>(2.2));
}

fn DistributionGGX(N: vec3<f32>, H: vec3<f32>, roughness: f32) -> f32 {
    let a = roughness * roughness;
    let a2 = a * a;
    let NdotH = max(dot(N, H), 0.0);
    let NdotH2 = NdotH * NdotH;
    let nom = a2;
    let denom = (NdotH2 * (a2 - 1.0) + 1.0);
    return nom / (PI * denom * denom);
}

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

fn GeometrySmith(N: vec3<f32>, V: vec3<f32>, L: vec3<f32>, roughness: f32) -> f32 {
    let NdotV = max(dot(N, V), 0.0);
    let NdotL = max(dot(N, L), 0.0);
    let ggx2 = GeometrySchlickGGX(NdotV, roughness);
    let ggx1 = GeometrySchlickGGX(NdotL, roughness);
    return ggx1 * ggx2;
}

fn fresnelSchlick(cosTheta: f32, F0: vec3<f32>) -> vec3<f32> {
    return F0 + (1.0 - F0) * pow(max(1.0 - cosTheta, 0.0), 5.0);
}

fn fresnelSchlickRoughness(cosTheta: f32, F0: vec3<f32>, roughness: f32) -> vec3<f32> {
    return F0 + (max(vec3<f32>(1.0 - roughness), F0) - F0) * pow(max(1.0 - cosTheta, 0.0), 5.0);
}

fn getRangeAttenuation(range: f32, distance: f32) -> f32 {
    if range <= 0.0 {
        return 1.0;
    }
    let clamped_distance = max(distance, 0.01);
    return max(min(1.0 - pow(distance / range, 4.0), 1.0), 0.0) / (clamped_distance * clamped_distance);
}

fn getSpotAttenuation(pointToLight: vec3<f32>, spotDirection: vec3<f32>, outerConeCos: f32, innerConeCos: f32) -> f32 {
    let actualCos = dot(normalize(spotDirection), normalize(-pointToLight));
    if actualCos > outerConeCos {
        if actualCos < innerConeCos {
            return smoothstep(outerConeCos, innerConeCos, actualCos);
        }
        return 1.0;
    }
    return 0.0;
}

fn getVolumeTransmissionRay(normal: vec3<f32>, view: vec3<f32>, thickness: f32, ior: f32, model_scale: vec3<f32>) -> vec3<f32> {
    let refraction_vector = refract(-view, normal, 1.0 / ior);
    return normalize(refraction_vector) * thickness * model_scale;
}

fn applyVolumeAttenuation(radiance: vec3<f32>, transmission_distance: f32, attenuation_color: vec3<f32>, attenuation_distance: f32) -> vec3<f32> {
    if attenuation_distance <= 0.0 {
        return radiance;
    }
    let attenuation_coefficient = -log(max(attenuation_color, vec3<f32>(0.0001))) / attenuation_distance;
    let transmittance = exp(-attenuation_coefficient * transmission_distance);
    return transmittance * radiance;
}

fn applyIorToRoughness(roughness: f32, ior: f32) -> f32 {
    return roughness * clamp(ior * 2.0 - 2.0, 0.0, 1.0);
}

fn getTransmissionSample(reflection: vec3<f32>, roughness: f32, ior: f32) -> vec3<f32> {
    let transmission_roughness = applyIorToRoughness(roughness, ior);
    let MAX_REFLECTION_LOD = 4.0;
    return textureSampleLevel(prefiltered_env, prefiltered_sampler, reflection, transmission_roughness * MAX_REFLECTION_LOD).rgb;
}

fn getIBLVolumeRefraction(
    normal: vec3<f32>,
    view: vec3<f32>,
    roughness: f32,
    base_color: vec3<f32>,
    F0: vec3<f32>,
    ior: f32,
    thickness: f32,
    attenuation_color: vec3<f32>,
    attenuation_distance: f32,
    model_scale: vec3<f32>
) -> vec3<f32> {
    let transmission_ray = getVolumeTransmissionRay(normal, view, thickness, ior, model_scale);
    let refracted_ray_exit = -normalize(transmission_ray);
    let transmitted_light = getTransmissionSample(refracted_ray_exit, roughness, ior);
    let attenuated_color = applyVolumeAttenuation(transmitted_light, length(transmission_ray), attenuation_color, attenuation_distance);
    let NdotV = clamp(dot(normal, view), 0.001, 1.0);
    let brdf = textureSampleLevel(brdf_lut, brdf_lut_sampler, vec2<f32>(NdotV, roughness), 0.0).rg;
    let specular_color = F0 * brdf.x + brdf.y;
    return (1.0 - specular_color) * attenuated_color * base_color;
}

struct VertexInput {
    @location(0) position: vec3<f32>,
    @location(1) normal: vec3<f32>,
    @location(2) tex_coords: vec2<f32>,
    @location(3) tangent: vec4<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>,
    @location(5) world_tangent: 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,
    global_unlit: f32,
    shadow_normal_bias: 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,
    has_normal_texture: u32,
    has_metallic_roughness_texture: u32,
    has_occlusion_texture: u32,
    normal_scale: f32,
    occlusion_strength: f32,
    roughness: f32,
    metallic: f32,
    unlit: u32,
    normal_map_flags: u32,
    uv_scale: vec2<f32>,
    transmission_factor: f32,
    has_transmission_texture: u32,
    thickness: f32,
    has_thickness_texture: u32,
    _align_attenuation: vec2<u32>,
    attenuation_color: vec3<f32>,
    attenuation_distance: f32,
    ior: f32,
    specular_factor: f32,
    _align_specular: vec2<u32>,
    specular_color_factor: vec3<f32>,
    has_specular_texture: u32,
    has_specular_color_texture: u32,
    emissive_strength: f32,
    _end_padding: vec2<u32>,
};

const NORMAL_MAP_FLIP_Y: u32 = 1u;
const NORMAL_MAP_TWO_COMPONENT: u32 = 2u;

struct ObjectData {
    transform_index: u32,
    mesh_id: u32,
    material_id: u32,
    batch_id: u32,
    morph_weights: array<f32, 8>,
    morph_target_count: u32,
    morph_displacement_offset: u32,
    mesh_vertex_offset: u32,
    mesh_vertex_count: u32,
    entity_id: u32,
    _entity_padding: array<u32, 3>,
};

struct MorphDisplacement {
    position: vec3<f32>,
    _pad0: f32,
    normal: vec3<f32>,
    _pad1: f32,
    tangent: vec3<f32>,
    _pad2: f32,
};

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(1) @binding(6)
var<storage, read> morph_displacements: array<MorphDisplacement>;

@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(2) @binding(4)
var normal_texture: texture_2d<f32>;

@group(2) @binding(5)
var normal_sampler: sampler;

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

@group(2) @binding(7)
var metallic_roughness_sampler: sampler;

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

@group(2) @binding(9)
var occlusion_sampler: sampler;

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

@group(2) @binding(11)
var transmission_sampler: sampler;

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

@group(2) @binding(13)
var thickness_sampler: sampler;

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

@group(2) @binding(15)
var specular_sampler: sampler;

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

@group(2) @binding(17)
var specular_color_sampler: sampler;

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

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

@group(3) @binding(2)
var spotlight_shadow_atlas: texture_depth_2d;

@group(3) @binding(3)
var spotlight_shadow_sampler: sampler;

@group(3) @binding(4)
var<storage, read> spotlight_shadows: array<vec4<f32>>;

@group(3) @binding(5)
var brdf_lut: texture_2d<f32>;

@group(3) @binding(6)
var brdf_lut_sampler: sampler;

@group(3) @binding(7)
var irradiance_map: texture_cube<f32>;

@group(3) @binding(8)
var irradiance_sampler: sampler;

@group(3) @binding(9)
var prefiltered_env: texture_cube<f32>;

@group(3) @binding(10)
var prefiltered_sampler: sampler;

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

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

    var position = in.position;
    var normal = in.normal;
    var tangent = in.tangent.xyz;

    if object.morph_target_count > 0u {
        let local_vertex_index = vertex_index - object.mesh_vertex_offset;

        for (var target_index = 0u; target_index < object.morph_target_count; target_index = target_index + 1u) {
            let weight = object.morph_weights[target_index];
            if abs(weight) > 0.0001 {
                let displacement_index = object.morph_displacement_offset
                    + target_index * object.mesh_vertex_count
                    + local_vertex_index;
                let displacement = morph_displacements[displacement_index];
                position = position + displacement.position * weight;
                normal = normal + displacement.normal * weight;
                tangent = tangent + displacement.tangent * weight;
            }
        }
        normal = normalize(normal);
        tangent = normalize(tangent);
    }

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

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

    return out;
}

fn getLightIntensity(light: Light, point_to_light: vec3<f32>, world_pos: vec3<f32>) -> vec3<f32> {
    var range_attenuation = 1.0;
    var spot_attenuation = 1.0;

    if light.light_type != LIGHT_TYPE_DIRECTIONAL {
        range_attenuation = getRangeAttenuation(light.range, length(point_to_light));
    }
    if light.light_type == LIGHT_TYPE_SPOT {
        spot_attenuation = getSpotAttenuation(point_to_light, light.direction.xyz, light.outer_cone, light.inner_cone);
    }

    return range_attenuation * spot_attenuation * light.color.rgb;
}

fn get_normal(
    world_normal: vec3<f32>,
    world_pos: vec3<f32>,
    tex_coord: vec2<f32>,
    normal_sample: vec3<f32>,
    has_normal_texture: bool,
    normal_scale: f32
) -> vec3<f32> {
    let tangent_normal = normal_sample * 2.0 - 1.0;

    let q1 = dpdx(world_pos);
    let q2 = dpdy(world_pos);
    let st1 = dpdx(tex_coord);
    let st2 = dpdy(tex_coord);

    let N = normalize(world_normal);
    let T = normalize(q1 * st2.y - q2 * st1.y);
    let B = -normalize(cross(N, T));
    let TBN = mat3x3<f32>(T, B, N);

    let mapped_normal = normalize(TBN * tangent_normal) * vec3<f32>(vec2<f32>(normal_scale), 1.0);
    return select(N, mapped_normal, has_normal_texture);
}


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, @builtin(front_facing) is_front_face: bool) -> OITOutput {
    let material = materials[in.material_id];
    let scaled_tex_coords = in.tex_coords * material.uv_scale;

    let normal_tex_sample = textureSampleLevel(normal_texture, normal_sampler, scaled_tex_coords, 0.0).xyz;

    let normal = get_normal(
        in.normal,
        in.world_pos,
        scaled_tex_coords,
        normal_tex_sample,
        material.has_normal_texture != 0u,
        material.normal_scale
    );

    let base_tex_sample = textureSample(base_texture, base_sampler, scaled_tex_coords);
    let emissive_tex_sample = textureSample(emissive_texture, emissive_sampler, scaled_tex_coords);

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

    let albedo = base_color.rgb;

    let mr_sample = textureSampleLevel(metallic_roughness_texture, metallic_roughness_sampler, scaled_tex_coords, 0.0);
    var metallic = material.metallic;
    var roughness = material.roughness;
    if material.has_metallic_roughness_texture != 0u {
        roughness = roughness * mr_sample.g;
        metallic = metallic * mr_sample.b;
    }

    let occlusion_sample = textureSampleLevel(occlusion_texture, occlusion_sampler, scaled_tex_coords, 0.0).r;
    var occlusion = 1.0;
    if material.has_occlusion_texture != 0u {
        occlusion = occlusion_sample;
    }

    var transmission_factor = material.transmission_factor;
    if material.has_transmission_texture != 0u {
        let transmission_sample = textureSampleLevel(transmission_texture, transmission_sampler, scaled_tex_coords, 0.0).r;
        transmission_factor = transmission_factor * transmission_sample;
    }

    var thickness = material.thickness;
    if material.has_thickness_texture != 0u {
        let thickness_sample = textureSampleLevel(thickness_texture, thickness_sampler, scaled_tex_coords, 0.0).g;
        thickness = thickness * thickness_sample;
    }

    var specular_factor = material.specular_factor;
    if material.has_specular_texture != 0u {
        let specular_sample = textureSampleLevel(specular_texture, specular_sampler, scaled_tex_coords, 0.0).a;
        specular_factor = specular_factor * specular_sample;
    }

    var specular_color_factor = material.specular_color_factor;
    if material.has_specular_color_texture != 0u {
        let specular_color_sample = textureSampleLevel(specular_color_texture, specular_color_sampler, scaled_tex_coords, 0.0).rgb;
        specular_color_factor = specular_color_factor * srgb_to_linear3(specular_color_sample);
    }

    var emission = material.emissive_factor * material.emissive_strength;
    if material.has_emissive_texture != 0u {
        emission = emission * srgb_to_linear3(emissive_tex_sample.rgb);
    }

    var final_color: vec3<f32>;

    if material.unlit != 0u || uniforms.global_unlit > 0.5 {
        final_color = albedo + emission;
    } else {
        let N = normal;
        let V = normalize(uniforms.camera_position.xyz - in.world_pos);
        let R = reflect(-V, N);

        let ior = material.ior;
        let dielectric_f0 = pow((ior - 1.0) / (ior + 1.0), 2.0);
        let dielectric_specular_f0 = min(vec3<f32>(dielectric_f0) * specular_color_factor, vec3<f32>(1.0));
        var F0 = mix(dielectric_specular_f0, albedo, metallic);

        var Lo = vec3<f32>(0.0);

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

            var point_to_light: vec3<f32>;
            if light.light_type != LIGHT_TYPE_DIRECTIONAL {
                point_to_light = light.position.xyz - in.world_pos;
            } else {
                point_to_light = -light.direction.xyz;
            }

            let L = normalize(point_to_light);
            let H = normalize(V + L);

            let radiance = getLightIntensity(light, point_to_light, in.world_pos);

            let NDF = DistributionGGX(N, H, roughness);
            let G = GeometrySmith(N, V, L, roughness);
            let F = fresnelSchlick(max(dot(H, V), 0.0), F0);

            let numerator = NDF * G * F;
            let denominator = 4.0 * max(dot(N, V), 0.0) * max(dot(N, L), 0.0) + 0.001;
            let specular = numerator / denominator;

            let kS = F;
            var kD = vec3<f32>(1.0) - kS;
            kD = kD * (1.0 - metallic);

            let NdotL = max(dot(N, L), 0.0);

            let light_contribution = (kD * albedo / PI + specular) * radiance * NdotL;

            Lo = Lo + light_contribution;
        }

        let F = fresnelSchlickRoughness(max(dot(N, V), 0.0), F0, roughness);
        let kS = F;
        var kD = 1.0 - kS;
        kD = kD * (1.0 - metallic);

        let irradiance = textureSampleLevel(irradiance_map, irradiance_sampler, N, 0.0).rgb;
        var diffuse_color = irradiance * albedo;

        let MAX_REFLECTION_LOD = 4.0;
        let prefiltered_color = textureSampleLevel(prefiltered_env, prefiltered_sampler, R, roughness * MAX_REFLECTION_LOD).rgb;
        let brdf = textureSampleLevel(brdf_lut, brdf_lut_sampler, vec2<f32>(max(dot(N, V), 0.0), roughness), 0.0).rg;
        let specular_ibl = prefiltered_color * (F * brdf.x + brdf.y) * specular_factor;

        if transmission_factor > 0.0 {
            let model_scale = vec3<f32>(1.0);
            let transmission = getIBLVolumeRefraction(
                N, V, roughness, albedo, F0, ior,
                thickness, material.attenuation_color, material.attenuation_distance, model_scale
            );
            diffuse_color = mix(diffuse_color, transmission, transmission_factor);
        }

        var ambient = kD * diffuse_color + specular_ibl;

        ambient = mix(ambient, ambient * occlusion, material.occlusion_strength);

        let color = ambient + Lo;

        final_color = color + emission;
    }

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