#import bevy_pbr::{
    forward_io::{VertexOutput, FragmentOutput},
    mesh_view_bindings::view,
    pbr_functions::{alpha_discard, apply_pbr_lighting, main_pass_post_lighting_processing},
    pbr_fragment::pbr_input_from_standard_material,
}

struct VoluMaterial {
    mesh_translation: vec3<f32>,
    sphere_radius: f32,
    color: vec4<f32>,
}

@group(2) @binding(100)
var<storage, read> volu_material: VoluMaterial;

fn sphere_sdf(p: vec3<f32>, center: vec3<f32>, radius: f32) -> f32 {
    return length(p - center) - radius;
}

fn raymarch(ray_origin: vec3<f32>, ray_direction: vec3<f32>) -> f32 {
    var t = 0.0;
    let max_distance = 100.0;
    let max_steps = 128;
    let epsilon = 0.001;

    for (var i = 0; i < max_steps; i++) {
        let current_pos = ray_origin + t * ray_direction;
        let distance = sphere_sdf(current_pos, volu_material.mesh_translation, volu_material.sphere_radius);

        if distance < epsilon {
            return t;
        }

        t += distance;

        if t > max_distance {
            break;
        }
    }

    return -1.0;
}

fn sphere_normal(p: vec3<f32>) -> vec3<f32> {
    let epsilon = 0.001;
    let center = volu_material.mesh_translation;
    let radius = volu_material.sphere_radius;

    let gradient = vec3<f32>(
        sphere_sdf(p + vec3<f32>(epsilon, 0.0, 0.0), center, radius) - sphere_sdf(p - vec3<f32>(epsilon, 0.0, 0.0), center, radius),
        sphere_sdf(p + vec3<f32>(0.0, epsilon, 0.0), center, radius) - sphere_sdf(p - vec3<f32>(0.0, epsilon, 0.0), center, radius),
        sphere_sdf(p + vec3<f32>(0.0, 0.0, epsilon), center, radius) - sphere_sdf(p - vec3<f32>(0.0, 0.0, epsilon), center, radius)
    );

    return normalize(gradient);
}

@fragment
fn fragment(
    vertex_output: VertexOutput,
    @builtin(front_facing) is_front: bool,
) -> FragmentOutput {
    var in = vertex_output;

    let ray_origin = view.world_position;
    let ray_direction = normalize(in.world_position.xyz - view.world_position);

    let t = raymarch(ray_origin, ray_direction);

    var out: FragmentOutput;

    if t > 0.0 {
        let hit_point = ray_origin + t * ray_direction;
        let normal = sphere_normal(hit_point);

        in.world_position = vec4<f32>(hit_point, 1.0);
        in.world_normal = normal;

        var pbr_input = pbr_input_from_standard_material(in, is_front);

        pbr_input.material.base_color = alpha_discard(pbr_input.material, pbr_input.material.base_color);

        out.color = apply_pbr_lighting(pbr_input);

        out.color = main_pass_post_lighting_processing(pbr_input, out.color);
    } else {
        out.color = vec4(1.0, 1.0, 1.0, 0.1);
    }

    return out;
}