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) shadow_pos: vec4<f32>,
    @location(5) instance_tint: vec4<f32>,
    @location(6) @interpolate(linear) tex_coords_affine: vec2<f32>,
    @location(7) clip_w: f32,
    @location(8) view_depth: f32,
    @location(9) world_tangent: vec4<f32>,
    @location(10) @interpolate(flat) entity_id: u32,
};

struct FragmentOutput {
    @location(0) color: vec4<f32>,
    @location(1) entity_id: 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,
    snap_resolution: vec2<f32>,
    snap_enabled: u32,
    affine_enabled: u32,
    fog_color: vec3<f32>,
    fog_enabled: u32,
    fog_start: f32,
    fog_end: f32,
    cascade_count: u32,
    _padding2: f32,
    cascade_view_projections: array<mat4x4<f32>, 4>,
    cascade_split_distances: vec4<f32>,
    cascade_atlas_offsets: array<vec4<f32>, 4>,
    cascade_atlas_scale: vec4<f32>,
    time: f32,
    directional_light_direction: 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,
    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,
    _padding: vec2<u32>,
};

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

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,
    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,
    shadow_index: i32,
    _padding0: f32,
    _padding1: f32,
    _padding2: f32,
};

struct SpotlightShadowData {
    view_projection: mat4x4<f32>,
    atlas_offset: vec2<f32>,
    atlas_scale: vec2<f32>,
    bias: f32,
    _padding0: f32,
    _padding1: f32,
    _padding2: f32,
};

struct PointLightShadowData {
    position: vec3<f32>,
    range: f32,
    bias: f32,
    shadow_index: i32,
    _padding0: f32,
    _padding1: 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(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<SpotlightShadowData>;

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

@group(3) @binding(11)
var point_shadow_cubemap: texture_cube_array<f32>;

@group(3) @binding(12)
var point_shadow_sampler: sampler;

@group(3) @binding(13)
var<storage, read> point_light_shadows: array<PointLightShadowData>;

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

    let mat = materials[object.material_id];

    let view_position = uniforms.view * world_position;
    var clip_position = uniforms.projection * view_position;

    if uniforms.snap_enabled == 1u {
        let resolution = uniforms.snap_resolution;
        let snapped_x = round(clip_position.x * resolution.x / clip_position.w) * clip_position.w / resolution.x;
        let snapped_y = round(clip_position.y * resolution.y / clip_position.w) * clip_position.w / resolution.y;
        clip_position.x = snapped_x;
        clip_position.y = snapped_y;
    }

    out.position = clip_position;
    let world_normal = normalize(normal_matrix * normal);
    out.normal = world_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.tex_coords_affine = in.tex_coords * mat.uv_scale;
    out.clip_w = clip_position.w;
    out.view_depth = -view_position.z;
    out.instance_tint = instance_custom_data[object.transform_index];

    let shadow_offset_position = world_position.xyz + world_normal * uniforms.shadow_normal_bias;
    let light_space_pos = uniforms.light_view_projection * vec4<f32>(shadow_offset_position, 1.0);
    out.shadow_pos = light_space_pos;

    out.entity_id = object.entity_id;

    return out;
}

fn select_cascade(view_depth: f32) -> i32 {
    for (var cascade_index = 0; cascade_index < 4; cascade_index = cascade_index + 1) {
        if view_depth < uniforms.cascade_split_distances[cascade_index] {
            return cascade_index;
        }
    }
    return 3;
}

fn sample_shadow_cascade(world_pos: vec3<f32>, world_normal: vec3<f32>, cascade_index: i32) -> f32 {
    let direction_to_light = -normalize(uniforms.directional_light_direction.xyz);
    let texel_size_world = uniforms.cascade_atlas_offsets[cascade_index].z;
    let normal_offset = uniforms.shadow_normal_bias * texel_size_world * world_normal;
    let depth_offset = uniforms.shadow_bias * direction_to_light;
    let shadow_offset_position = world_pos + normal_offset + depth_offset;
    let shadow_pos = uniforms.cascade_view_projections[cascade_index] * vec4<f32>(shadow_offset_position, 1.0);

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

    let atlas_offset = uniforms.cascade_atlas_offsets[cascade_index].xy;
    let atlas_scale = uniforms.cascade_atlas_scale.xy;
    shadow_uv = shadow_uv * atlas_scale + atlas_offset;

    let shadow_depth = shadow_coords_ndc.z;

    let slot_min = atlas_offset;
    let slot_max = atlas_offset + atlas_scale;
    let in_bounds = f32(
        shadow_uv.x >= slot_min.x && shadow_uv.x <= slot_max.x &&
        shadow_uv.y >= slot_min.y && shadow_uv.y <= slot_max.y &&
        shadow_depth >= 0.0 && shadow_depth <= 1.0
    );

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

    let slot_margin = atlas_scale * 0.01;
    let safe_min = slot_min + slot_margin;
    let safe_max = slot_max - slot_margin;

    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 = 1.5;

    for (var sample_index = 0; sample_index < 16; sample_index = sample_index + 1) {
        let offset = poisson_disk[sample_index] * texel_size * filter_radius;
        let sample_uv = clamp(shadow_uv + offset, safe_min, safe_max);
        let sampled_depth = textureSampleLevel(shadow_texture, shadow_sampler, sample_uv, 0);
        visibility = visibility + select(0.0, 1.0, shadow_depth >= sampled_depth);
    }

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

fn calculate_shadow_factor(world_pos: vec3<f32>, world_normal: vec3<f32>, view_depth: f32) -> f32 {
    if uniforms.shadows_enabled < 0.5 {
        return 1.0;
    }

    let cascade_index = select_cascade(view_depth);
    let shadow_factor = sample_shadow_cascade(world_pos, world_normal, cascade_index);

    if cascade_index < 3 {
        let cascade_end = uniforms.cascade_split_distances[cascade_index];
        let blend_range = cascade_end * 0.25;
        let blend_start = cascade_end - blend_range;

        if view_depth > blend_start {
            let blend_factor = (view_depth - blend_start) / blend_range;
            let next_shadow_factor = sample_shadow_cascade(world_pos, world_normal, cascade_index + 1);
            return mix(shadow_factor, next_shadow_factor, blend_factor);
        }
    }

    return shadow_factor;
}

fn calculate_spotlight_shadow_factor(world_pos: vec3<f32>, normal: vec3<f32>, light_pos: vec3<f32>, shadow_index: i32) -> f32 {
    if shadow_index < 0 {
        return 1.0;
    }

    let shadow_data = spotlight_shadows[shadow_index];
    let light_space_pos = shadow_data.view_projection * vec4<f32>(world_pos, 1.0);

    if light_space_pos.w <= 0.0 {
        return 1.0;
    }

    let shadow_coords_ndc = light_space_pos.xyz / light_space_pos.w;

    var shadow_uv = shadow_coords_ndc.xy * vec2<f32>(0.5, -0.5) + vec2<f32>(0.5, 0.5);
    shadow_uv = shadow_uv * shadow_data.atlas_scale + shadow_data.atlas_offset;

    let shadow_depth = shadow_coords_ndc.z;

    let slot_min = shadow_data.atlas_offset;
    let slot_max = shadow_data.atlas_offset + shadow_data.atlas_scale;
    let in_bounds = f32(
        shadow_uv.x >= slot_min.x && shadow_uv.x <= slot_max.x &&
        shadow_uv.y >= slot_min.y && shadow_uv.y <= slot_max.y &&
        shadow_depth >= 0.0 && shadow_depth <= 1.0
    );

    let atlas_size = f32(textureDimensions(spotlight_shadow_atlas).x);
    let texel_size = 1.0 / atlas_size;
    let filter_radius = 0.5;

    let slot_margin = shadow_data.atlas_scale * 0.01;
    let safe_min = slot_min + slot_margin;
    let safe_max = slot_max - slot_margin;

    let light_dir = normalize(light_pos - world_pos);
    let slope_factor = 1.0 - abs(dot(normal, light_dir));
    let adjusted_bias = shadow_data.bias * (1.0 + slope_factor * 2.0);

    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 reference_depth = shadow_depth + adjusted_bias;
    for (var i = 0; i < 16; i = i + 1) {
        let offset = poisson_disk[i] * texel_size * filter_radius;
        let sample_uv = clamp(shadow_uv + offset, safe_min, safe_max);
        let sampled_depth = textureSampleLevel(spotlight_shadow_atlas, spotlight_shadow_sampler, sample_uv, 0);
        visibility = visibility + select(0.0, 1.0, reference_depth >= sampled_depth);
    }

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

fn calculate_point_light_shadow_factor(world_pos: vec3<f32>, shadow_index: i32) -> f32 {
    if shadow_index < 0 {
        return 1.0;
    }

    let shadow_data = point_light_shadows[shadow_index];
    let light_pos = shadow_data.position;
    let light_range = shadow_data.range;
    let bias = shadow_data.bias;

    let light_to_frag = world_pos - light_pos;
    let distance = length(light_to_frag);

    if distance > light_range {
        return 1.0;
    }

    let direction = normalize(light_to_frag);
    let normalized_distance = distance / light_range;

    let sampled_depth = textureSampleLevel(point_shadow_cubemap, point_shadow_sampler, direction, shadow_index, 0.0).r;

    let distance_with_bias = normalized_distance - bias * 0.1;

    if distance_with_bias <= sampled_depth {
        return 1.0;
    }
    return 0.0;
}

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 safe_normalize(v: vec3<f32>) -> vec3<f32> {
    let len_sq = dot(v, v);
    if len_sq < 0.0001 {
        return vec3<f32>(0.0, 1.0, 0.0);
    }
    return v / sqrt(len_sq);
}

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 = safe_normalize(world_normal);
    let T = safe_normalize(q1 * st2.y - q2 * st1.y);
    let B = -safe_normalize(cross(N, T));
    let TBN = mat3x3<f32>(T, B, N);

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


@fragment
fn fs_main(in: VertexOutput, @builtin(front_facing) is_front_face: bool) -> FragmentOutput {
    var output: FragmentOutput;
    let material = materials[in.material_id];

    var base_tex_coords: vec2<f32>;
    if uniforms.affine_enabled == 1u {
        base_tex_coords = in.tex_coords_affine;
    } else {
        base_tex_coords = in.tex_coords * material.uv_scale;
    }
    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 base_tex_sample = textureSample(base_texture, base_sampler, base_tex_coords);
    let emissive_tex_sample = textureSample(emissive_texture, emissive_sampler, base_tex_coords);

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

    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 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 emission = material.emissive_factor * material.emissive_strength;
    if material.has_emissive_texture != 0u {
        emission = emission * srgb_to_linear3(emissive_tex_sample.rgb);
    }

    let shadow_factor = calculate_shadow_factor(in.world_pos, in.normal, in.view_depth);

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

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

            if light.light_type == LIGHT_TYPE_DIRECTIONAL {
                light_contribution = light_contribution * shadow_factor;
            } else if light.light_type == LIGHT_TYPE_SPOT {
                let spotlight_shadow = calculate_spotlight_shadow_factor(in.world_pos, N, light.position.xyz, light.shadow_index);
                light_contribution = light_contribution * spotlight_shadow;
            } else if light.light_type == LIGHT_TYPE_POINT {
                let point_shadow = calculate_point_light_shadow_factor(in.world_pos, light.shadow_index);
                light_contribution = light_contribution * point_shadow;
            }

            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 = clamp(textureSampleLevel(irradiance_map, irradiance_sampler, N, 0.0).rgb, vec3<f32>(0.0), vec3<f32>(65000.0));
        var diffuse_color = irradiance * albedo;

        let MAX_REFLECTION_LOD = 4.0;
        let prefiltered_color = clamp(textureSampleLevel(prefiltered_env, prefiltered_sampler, R, roughness * MAX_REFLECTION_LOD).rgb, vec3<f32>(0.0), vec3<f32>(65000.0));
        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;
    }

    if uniforms.fog_enabled == 1u {
        let fog_factor = clamp((in.view_depth - uniforms.fog_start) / (uniforms.fog_end - uniforms.fog_start), 0.0, 1.0);
        final_color = mix(final_color, uniforms.fog_color, fog_factor);
    }

    final_color = clamp(final_color, vec3<f32>(0.0), vec3<f32>(65000.0));

    output.color = vec4<f32>(final_color, alpha);
    output.entity_id = bitcast<f32>(in.entity_id);
    return output;
}
