#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>,
    model_inverse: mat4x4<f32>,
    power: f32,
    iterations: i32,
    max_distance: f32,
    max_steps: i32,
    ray_epsilon: f32,
    normal_epsilon: f32,
}

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

fn mandelbulb_sdf(p: vec3<f32>) -> f32 {
    var z = p;
    var dr = 1.0;
    var r = 0.0;
    let power = volu_material.power;
    let iterations = volu_material.iterations;

    for (var i = 0; i < iterations; i++) {
        r = length(z);
        if r > 2.0 {
            break;
        }

        let theta = acos(clamp(z.z / r, -1.0, 1.0));
        let phi = atan2(z.y, z.x);
        dr = pow(r, power - 1.0) * power * dr + 1.0;

        let zr = pow(r, power);
        let new_theta = theta * power;
        let new_phi = phi * power;

        z = zr * vec3<f32>(
            sin(new_theta) * cos(new_phi),
            sin(new_theta) * sin(new_phi),
            cos(new_theta)
        ) + p;
    }

    return 0.5 * log(max(r, 0.001)) * r / dr;
}

fn mandelbulb_normal(p: vec3<f32>) -> vec3<f32> {
    let e = volu_material.normal_epsilon;
    let dx = mandelbulb_sdf(p + vec3<f32>(e, 0.0, 0.0)) - mandelbulb_sdf(p - vec3<f32>(e, 0.0, 0.0));
    let dy = mandelbulb_sdf(p + vec3<f32>(0.0, e, 0.0)) - mandelbulb_sdf(p - vec3<f32>(0.0, e, 0.0));
    let dz = mandelbulb_sdf(p + vec3<f32>(0.0, 0.0, e)) - mandelbulb_sdf(p - vec3<f32>(0.0, 0.0, e));
    return normalize(vec3<f32>(dx, dy, dz));
}

fn raymarch(ray_origin: vec3<f32>, ray_direction: vec3<f32>) -> f32 {
    var t = 0.0;
    let max_distance = volu_material.max_distance;
    let max_steps = volu_material.max_steps;
    let epsilon = volu_material.ray_epsilon;

    for (var i = 0; i < max_steps; i++) {
        let current_pos = ray_origin + t * ray_direction;
        let distance = mandelbulb_sdf(current_pos);

        if distance < epsilon {
            return t;
        }

        t += distance;

        if t > max_distance {
            break;
        }
    }

    return -1.0;
}

@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 local_origin = (volu_material.model_inverse * vec4<f32>(ray_origin, 1.0)).xyz;
    let local_direction = normalize((volu_material.model_inverse * vec4<f32>(ray_direction, 0.0)).xyz);

    let t = raymarch(local_origin, local_direction);

    var out: FragmentOutput;

    if t > 0.0 {
        let local_hit_point = ray_origin + t * ray_direction;
        let local_normal = mandelbulb_normal(local_hit_point);

        let world_hit_point = (transpose(volu_material.model_inverse) * vec4<f32>(local_hit_point, 1.0)).xyz;
        let world_normal = normalize((transpose(volu_material.model_inverse) * vec4<f32>(local_normal, 0.0)).xyz);

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

        var pbr_input = pbr_input_from_standard_material(in, is_front);

        let n = world_normal * 0.5 + 0.5;
        pbr_input.material.base_color = alpha_discard(pbr_input.material, vec4<f32>(n, 1.0));

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