use bevy::{
    camera_controller::free_camera::{FreeCamera, FreeCameraPlugin},
    core_pipeline::{
        core_3d::CORE_3D_DEPTH_FORMAT,
        deferred::{
            copy_lighting_id::copy_deferred_lighting_id, node::late_deferred_prepass,
            DEFERRED_LIGHTING_PASS_ID_FORMAT, DEFERRED_PREPASS_FORMAT,
        },
        prepass::{DeferredPrepass, DepthPrepass, ViewPrepassTextures},
        Core3d, Core3dSystems, FullscreenShader,
    },
    pbr::{
        per_view_shadow_pass, shared_shadow_pass, DefaultOpaqueRendererMethod, ShadowView,
        ViewLightEntities, LATE_SHADOW_PASS,
    },
    prelude::*,
    render::{
        globals::{GlobalsBuffer, GlobalsUniform},
        render_resource::{binding_types::uniform_buffer, *},
        renderer::{RenderContext, ViewQuery},
        view::{ViewDepthStencilTexture, ViewUniform, ViewUniformOffset, ViewUniforms},
        RenderApp, RenderStartup,
    },
};

fn main() {
    App::new()
        // Render everything through the deferred pipeline
        .insert_resource(DefaultOpaqueRendererMethod::deferred())
        .add_plugins((DefaultPlugins, DeferredRaymarchPlugin, FreeCameraPlugin))
        .add_systems(Startup, setup)
        .run();
}

const SHADER_ASSET_PATH: &str = "shaders/deferred_raymarch.wesl";

fn setup(
    mut commands: Commands,
    mut meshes: ResMut<Assets<Mesh>>,
    mut materials: ResMut<Assets<StandardMaterial>>,
) {
    commands.spawn((
        Camera3d::default(),
        Transform::from_xyz(4.0, 3.0, 6.0).looking_at(Vec3::new(0.0, 0.2, 0.0), Vec3::Y),
        // Deferred rendering requires MSAA to be off.
        Msaa::Off,
        DepthPrepass,
        DeferredPrepass,
        AmbientLight {
            brightness: 200.0,
            ..default()
        },
        FreeCamera::default(),
    ));

    // A ground plane that catches the SDF's shadow
    commands.spawn((
        Mesh3d(meshes.add(Plane3d::default().mesh().size(20.0, 20.0))),
        MeshMaterial3d(materials.add(Color::srgb(0.3, 0.5, 0.3))),
        Transform::from_xyz(0.0, -1.5, 0.0),
    ));

    // A "regular" mesh cube
    commands.spawn((
        Mesh3d(meshes.add(Cuboid::from_length(1.2))),
        MeshMaterial3d(materials.add(StandardMaterial {
            base_color: Color::srgb(0.2, 0.4, 0.9),
            perceptual_roughness: 0.4,
            ..default()
        })),
        Transform::from_xyz(2.2, -0.9, 0.5),
    ));

    commands.spawn((
        DirectionalLight {
            illuminance: 8_000.0,
            shadow_maps_enabled: true,
            ..default()
        },
        Transform::from_xyz(4.0, 8.0, 4.0).looking_at(Vec3::ZERO, Vec3::Y),
    ));
}

struct DeferredRaymarchPlugin;

impl Plugin for DeferredRaymarchPlugin {
    fn build(&self, app: &mut App) {
        let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
            return;
        };

        render_app
            .add_systems(RenderStartup, init_raymarch_pipelines)
            .add_systems(
                Core3d,
                (
                    // The gbuffer write must run after the mesh deferred prepass has
                    // filled the gbuffer, but *before* `copy_deferred_lighting_id`
                    // bakes the lighting-pass ids into the depth routing texture the
                    // lighting pass reads
                    raymarch_gbuffer_pass
                        .in_set(Core3dSystems::Prepass)
                        .after(late_deferred_prepass)
                        .before(copy_deferred_lighting_id),
                    // Write the SDF into the shadow maps after the mesh shadow passes
                    // have drawn, so it casts shadows like any other caster
                    raymarch_directional_shadow_pass
                        .after(per_view_shadow_pass::<LATE_SHADOW_PASS>)
                        .before(Core3dSystems::MainPass),
                    raymarch_shared_shadow_pass
                        .after(shared_shadow_pass::<LATE_SHADOW_PASS>)
                        .before(Core3dSystems::MainPass),
                ),
            );
    }
}

#[derive(Resource)]
struct RaymarchGBufferPipeline {
    layout: BindGroupLayoutDescriptor,
    pipeline_id: CachedRenderPipelineId,
}

#[derive(Resource)]
struct RaymarchShadowPipeline {
    layout: BindGroupLayoutDescriptor,
    pipeline_id: CachedRenderPipelineId,
}

fn init_raymarch_pipelines(
    mut commands: Commands,
    asset_server: Res<AssetServer>,
    fullscreen_shader: Res<FullscreenShader>,
    pipeline_cache: Res<PipelineCache>,
) {
    let layout = BindGroupLayoutDescriptor::new(
        "raymarch_bind_group_layout",
        &BindGroupLayoutEntries::sequential(
            ShaderStages::FRAGMENT,
            (
                uniform_buffer::<ViewUniform>(true),
                uniform_buffer::<GlobalsUniform>(false),
            ),
        ),
    );

    let shader = asset_server.load::<Shader>(SHADER_ASSET_PATH);
    let vertex_state = fullscreen_shader.to_vertex_state();

    // Writing depth lets the SDF sort against meshes
    let depth_stencil = DepthStencilState {
        format: CORE_3D_DEPTH_FORMAT,
        depth_write_enabled: Some(true),
        depth_compare: Some(CompareFunction::GreaterEqual),
        stencil: StencilState::default(),
        bias: DepthBiasState::default(),
    };

    let gbuffer_pipeline_id = pipeline_cache.queue_render_pipeline(RenderPipelineDescriptor {
        label: Some("raymarch_gbuffer_pipeline".into()),
        layout: vec![layout.clone()],
        vertex: vertex_state.clone(),
        fragment: Some(FragmentState {
            shader: shader.clone(),
            entry_point: Some("fragment".into()),
            targets: vec![
                Some(ColorTargetState {
                    format: DEFERRED_PREPASS_FORMAT,
                    blend: None,
                    write_mask: ColorWrites::ALL,
                }),
                Some(ColorTargetState {
                    format: DEFERRED_LIGHTING_PASS_ID_FORMAT,
                    blend: None,
                    write_mask: ColorWrites::ALL,
                }),
            ],
            ..default()
        }),
        depth_stencil: Some(depth_stencil.clone()),
        ..default()
    });

    let shadow_pipeline_id = pipeline_cache.queue_render_pipeline(RenderPipelineDescriptor {
        label: Some("raymarch_shadow_pipeline".into()),
        layout: vec![layout.clone()],
        vertex: vertex_state,
        fragment: Some(FragmentState {
            shader,
            entry_point: Some("fragment_shadow".into()),
            targets: vec![],
            ..default()
        }),
        depth_stencil: Some(depth_stencil),
        ..default()
    });

    commands.insert_resource(RaymarchGBufferPipeline {
        layout: layout.clone(),
        pipeline_id: gbuffer_pipeline_id,
    });
    commands.insert_resource(RaymarchShadowPipeline {
        layout,
        pipeline_id: shadow_pipeline_id,
    });
}

fn raymarch_bind_group(
    ctx: &RenderContext,
    pipeline_cache: &PipelineCache,
    layout: &BindGroupLayoutDescriptor,
    view_uniforms: &ViewUniforms,
    globals: &GlobalsBuffer,
) -> Option<BindGroup> {
    let view_binding = view_uniforms.uniforms.binding()?;
    let globals_binding = globals.buffer.binding()?;
    Some(ctx.render_device().create_bind_group(
        "raymarch_bind_group",
        &pipeline_cache.get_bind_group_layout(layout),
        &BindGroupEntries::sequential((view_binding, globals_binding)),
    ))
}

fn raymarch_gbuffer_pass(
    view: ViewQuery<(
        &ViewUniformOffset,
        &ViewDepthStencilTexture,
        &ViewPrepassTextures,
    )>,
    pipeline: Option<Res<RaymarchGBufferPipeline>>,
    pipeline_cache: Res<PipelineCache>,
    view_uniforms: Res<ViewUniforms>,
    globals: Res<GlobalsBuffer>,
    mut ctx: RenderContext,
) {
    let Some(pipeline) = pipeline else {
        return;
    };
    let Some(render_pipeline) = pipeline_cache.get_render_pipeline(pipeline.pipeline_id) else {
        return;
    };

    let (view_uniform_offset, view_depth, view_prepass_textures) = view.into_inner();

    let (Some(deferred), Some(lighting_pass_id)) = (
        &view_prepass_textures.deferred,
        &view_prepass_textures.deferred_lighting_pass_id,
    ) else {
        return;
    };

    let Some(bind_group) = raymarch_bind_group(
        &ctx,
        &pipeline_cache,
        &pipeline.layout,
        &view_uniforms,
        &globals,
    ) else {
        return;
    };

    {
        // We load rather than clear because we only want to overwrite the pixels which the deferred mesh
        // prepass didn't write
        let mut pass = ctx.begin_tracked_render_pass(RenderPassDescriptor {
            label: Some("raymarch_gbuffer_pass"),
            color_attachments: &[
                Some(deferred.get_attachment()),
                Some(lighting_pass_id.get_attachment()),
            ],
            depth_stencil_attachment: Some(view_depth.get_attachment(StoreOp::Store)),
            timestamp_writes: None,
            occlusion_query_set: None,
            multiview_mask: None,
        });

        pass.set_render_pipeline(render_pipeline);
        pass.set_bind_group(0, &bind_group, &[view_uniform_offset.offset]);
        pass.draw(0..3, 0..1);
    }

    // The deferred lighting pass reconstructs world position from the prepass depth
    // texture, not the depth attachment we just wrote, so we have to do a copy.
    if let Some(prepass_depth) = &view_prepass_textures.depth {
        ctx.command_encoder().copy_texture_to_texture(
            view_depth.texture().as_image_copy(),
            prepass_depth.texture.texture.as_image_copy(),
            view_prepass_textures.size,
        );
    }
}

fn raymarch_directional_shadow_pass(
    view: ViewQuery<&ViewLightEntities>,
    shadow_views: Query<(&ShadowView, &ViewUniformOffset)>,
    pipeline: Option<Res<RaymarchShadowPipeline>>,
    pipeline_cache: Res<PipelineCache>,
    view_uniforms: Res<ViewUniforms>,
    globals: Res<GlobalsBuffer>,
    mut ctx: RenderContext,
) {
    let Some(pipeline) = pipeline else {
        return;
    };
    let Some(render_pipeline) = pipeline_cache.get_render_pipeline(pipeline.pipeline_id) else {
        return;
    };

    let view_lights = view.into_inner();
    for light_entity in view_lights.lights.iter().copied() {
        let Ok((shadow_view, view_uniform_offset)) = shadow_views.get(light_entity) else {
            continue;
        };
        draw_raymarch_shadow(
            &mut ctx,
            &pipeline_cache,
            &pipeline.layout,
            &view_uniforms,
            &globals,
            render_pipeline,
            shadow_view,
            view_uniform_offset,
        );
    }
}

fn raymarch_shared_shadow_pass(
    view: ViewQuery<(&ShadowView, &ViewUniformOffset)>,
    pipeline: Option<Res<RaymarchShadowPipeline>>,
    pipeline_cache: Res<PipelineCache>,
    view_uniforms: Res<ViewUniforms>,
    globals: Res<GlobalsBuffer>,
    mut ctx: RenderContext,
) {
    let Some(pipeline) = pipeline else {
        return;
    };
    let Some(render_pipeline) = pipeline_cache.get_render_pipeline(pipeline.pipeline_id) else {
        return;
    };

    let (shadow_view, view_uniform_offset) = view.into_inner();
    draw_raymarch_shadow(
        &mut ctx,
        &pipeline_cache,
        &pipeline.layout,
        &view_uniforms,
        &globals,
        render_pipeline,
        shadow_view,
        view_uniform_offset,
    );
}

fn draw_raymarch_shadow(
    ctx: &mut RenderContext,
    pipeline_cache: &PipelineCache,
    layout: &BindGroupLayoutDescriptor,
    view_uniforms: &ViewUniforms,
    globals: &GlobalsBuffer,
    render_pipeline: &RenderPipeline,
    shadow_view: &ShadowView,
    view_uniform_offset: &ViewUniformOffset,
) {
    let Some(bind_group) = raymarch_bind_group(ctx, pipeline_cache, layout, view_uniforms, globals)
    else {
        return;
    };

    let mut pass = ctx.begin_tracked_render_pass(RenderPassDescriptor {
        label: Some("raymarch_shadow_pass"),
        color_attachments: &[],
        depth_stencil_attachment: Some(shadow_view.depth_attachment.get_attachment(StoreOp::Store)),
        timestamp_writes: None,
        occlusion_query_set: None,
        multiview_mask: None,
    });

    pass.set_render_pipeline(render_pipeline);
    pass.set_bind_group(0, &bind_group, &[view_uniform_offset.offset]);
    pass.draw(0..3, 0..1);
}

//! Raymarches a signed distance field and packs it into the deferred gbuffer exactly as
//! a `StandardMaterial` would (via `deferred_gbuffer_from_pbr_input`), so Bevy's deferred
//! PBR lighting shades it like any other geometry. This example assumes familiarity with
//! raymarching.

import bevy_core_pipeline::fullscreen_vertex_shader::FullscreenVertexOutput;
import bevy_pbr::render::mesh_view_bindings::view;
import bevy_pbr::render::view_transformations::{
    uv_to_ndc,
    position_ndc_to_view,
    position_ndc_to_world,
    direction_view_to_world,
    position_world_to_view,
    view_z_to_depth_ndc,
};
import bevy_pbr::deferred::functions::deferred_gbuffer_from_pbr_input;
import bevy_pbr::render::pbr_types::{pbr_input_new, STANDARD_MATERIAL_FLAGS_ALPHA_MODE_OPAQUE};
import bevy_pbr::render::mesh_types::MESH_FLAGS_SHADOW_RECEIVER_BIT;
import bevy_render::globals::Globals;

@group(0) @binding(1) var<uniform> globals: Globals;

const FAR: f32 = 100.0;
const SURFACE_EPSILON: f32 = 0.001;
const MAX_STEPS: u32 = 160u;
const STEP_SCALE: f32 = 0.5;

// Swap this for any SDF you like, the gbuffer integration below is independent of the
// shape.
const GYROID_SCALE: f32 = 5.0;

fn gyroid_field(p: vec3<f32>) -> f32 {
    let q = p * GYROID_SCALE + globals.time * 0.6;
    return dot(sin(q), cos(q.yzx));
}

fn map(p: vec3<f32>) -> f32 {
    let shell = (abs(gyroid_field(p)) - 0.14) / GYROID_SCALE;
    let ball = length(p) - 1.3;
    return max(ball, shell);
}

fn sdf_normal(p: vec3<f32>) -> vec3<f32> {
    let e = vec2<f32>(1.0, -1.0) * 0.0005;
    return normalize(
        e.xyy * map(p + e.xyy) +
        e.yyx * map(p + e.yyx) +
        e.yxy * map(p + e.yxy) +
        e.xxx * map(p + e.xxx)
    );
}

fn raymarch(ray_origin: vec3<f32>, ray_dir: vec3<f32>) -> f32 {
    var t = 0.0;
    for (var i = 0u; i < MAX_STEPS; i++) {
        let d = map(ray_origin + ray_dir * t);
        if d < SURFACE_EPSILON {
            return t;
        }
        t += d * STEP_SCALE;
        if t > FAR {
            break;
        }
    }
    return FAR;
}

struct Ray {
    origin: vec3<f32>,
    dir: vec3<f32>,
}

fn ray_for_uv(uv: vec2<f32>) -> Ray {
    let ndc = uv_to_ndc(uv);
    var ray: Ray;
    // Check ortho projection
    if view.clip_from_view[3].w == 1.0 {
        ray.origin = position_ndc_to_world(vec3<f32>(ndc, 1.0));
        ray.dir = normalize(direction_view_to_world(vec3<f32>(0.0, 0.0, -1.0)));
    } else {
        ray.origin = view.world_position;
        ray.dir = normalize(direction_view_to_world(position_ndc_to_view(vec3<f32>(ndc, 1.0))));
    }
    return ray;
}

fn depth_for_world_pos(world_pos: vec3<f32>) -> f32 {
    return view_z_to_depth_ndc(position_world_to_view(world_pos).z);
}

/// Our render pass has two color targets, matching the deferred prepass:
///   location 0: the packed gbuffer (Rgba32Uint)
///   location 1: the deferred lighting pass id (R8Uint)
struct GBufferOutput {
    @location(0) deferred: vec4<u32>,
    @location(1) deferred_lighting_pass_id: u32,
    @builtin(frag_depth) depth: f32,
}

@fragment
fn fragment(in: FullscreenVertexOutput) -> GBufferOutput {
    let ray = ray_for_uv(in.uv);
    let t = raymarch(ray.origin, ray.dir);

    // discard so we leave the mesh gbuffer intact when we miss
    if t >= FAR {
        discard;
    }

    let world_pos = ray.origin + ray.dir * t;
    let normal = sdf_normal(world_pos);

    // Fill in a PbrInput as a StandardMaterial fragment would
    var pbr_input = pbr_input_new();
    pbr_input.frag_coord = vec4<f32>(in.position.xy, depth_for_world_pos(world_pos), 1.0);
    pbr_input.world_position = vec4<f32>(world_pos, 1.0);
    pbr_input.world_normal = normal;
    pbr_input.N = normal;
    pbr_input.V = normalize(view.world_position - world_pos);

    // Per-pixel base color
    let field = gyroid_field(world_pos);
    let color = 0.5 + 0.5 * cos(
        6.2831 * (field * 0.4 + globals.time * 0.05) + vec3<f32>(0.0, 0.8, 1.6)
    );
    pbr_input.material.base_color = vec4<f32>(color, 1.0);
    pbr_input.material.metallic = 0.7;
    pbr_input.material.perceptual_roughness = 0.25;
    pbr_input.material.flags = STANDARD_MATERIAL_FLAGS_ALPHA_MODE_OPAQUE;
    // Let the surface receive shadows cast by other geometry.
    pbr_input.flags = MESH_FLAGS_SHADOW_RECEIVER_BIT;

    var out: GBufferOutput;
    out.deferred = deferred_gbuffer_from_pbr_input(pbr_input);
    // The lighting pass reads this per-pixel id to choose which lighting shader runs.
    //  1 is Bevy's built-in PBR deferred shader.
    out.deferred_lighting_pass_id = 1u;
    out.depth = pbr_input.frag_coord.z;
    return out;
}

// When rendering into a light's shadow map we only need depth
@fragment
fn fragment_shadow(in: FullscreenVertexOutput) -> @builtin(frag_depth) f32 {
    let ray = ray_for_uv(in.uv);
    let t = raymarch(ray.origin, ray.dir);
    if t >= FAR {
        discard;
    }
    return depth_for_world_pos(ray.origin + ray.dir * t);
}