//! Demonstrates how to trigger various rendering errors, and how bevy can recover from them.

use bevy::{
    input::keyboard::Key,
    prelude::*,
    render::{
        error_handler::{RenderErrorHandler, RenderErrorPolicy},
        extract_resource::{ExtractResource, ExtractResourcePlugin},
        render_resource::{
            BufferDescriptor, BufferUsages, CommandEncoderDescriptor, ComputePassDescriptor,
            Extent3d, PipelineLayoutDescriptor, PollType, RawComputePipelineDescriptor,
            ShaderModuleDescriptor, ShaderSource, TextureDescriptor, TextureDimension,
            TextureFormat, TextureUsages,
        },
        renderer::{RenderDevice, RenderQueue},
        Render, RenderApp,
    },
};

fn main() {
    let mut app = App::new();
    app.add_plugins((
        DefaultPlugins,
        ExtractResourcePlugin::<RenderError>::default(),
    ))
    .add_systems(Startup, setup)
    .add_systems(Update, (update_camera, input))
    .init_resource::<RenderError>()
    .sub_app_mut(RenderApp)
    .add_systems(Render, cause_error);
    app.run();
}

fn setup(
    mut commands: Commands,
    mut meshes: ResMut<Assets<Mesh>>,
    mut materials: ResMut<Assets<StandardMaterial>>,
) {
    // circular base
    commands.spawn((
        Mesh3d(meshes.add(Circle::new(4.0))),
        MeshMaterial3d(materials.add(Color::WHITE)),
        Transform::from_rotation(Quat::from_rotation_x(-std::f32::consts::FRAC_PI_2)),
    ));
    // cube
    commands.spawn((
        Mesh3d(meshes.add(Cuboid::new(2.0, 2.0, 2.0))),
        MeshMaterial3d(materials.add(Color::srgb_u8(124, 144, 255))),
        Transform::from_xyz(0.0, 1.0, 0.0),
    ));
    // light
    commands.spawn((
        PointLight {
            shadow_maps_enabled: true,
            ..default()
        },
        Transform::from_xyz(4.0, 8.0, 4.0),
    ));
    // camera
    commands.spawn((
        Camera3d::default(),
        Transform::from_xyz(-2.5, 4.5, 9.0).looking_at(Vec3::ZERO, Vec3::Y),
    ));
    // help text
    commands.spawn((
        Text::new(
            "Test at your own risk: you may need to restart your computer to fully recover\n\
            Press O to trigger an OutOfMemory error\n\
            Press V to trigger a Validation error\n\
            Press D to Destroy the render device (causes device lost error)\n\
            Press L to Loop infinitely in a compute shader (causes device lost error)\n\
            \n\
            Press 1 to ignore errors, pretending nothing happened and continue rendering.\n\
            Press 2 to panic on error.\n\
            Press 3 to signals app exit on error.\n\
            Press 4 to keeps the app alive, but stops rendering further on error.\n\
            Press 5 to attempt renderer recovery.\n\
            ",
        ),
        Node {
            position_type: PositionType::Absolute,
            top: px(12),
            left: px(12),
            ..default()
        },
    ));
}

fn update_camera(mut camera: Query<&mut Transform, With<Camera>>, time: Res<Time>) {
    for mut t in camera.iter_mut() {
        let (s, c) = ops::sin_cos(time.elapsed_secs() * 0.3);
        *t = Transform::from_xyz(s * 10.0, 4.5, c * 10.0).looking_at(Vec3::ZERO, Vec3::Y);
    }
}

#[derive(Resource, ExtractResource, Clone, Default)]
#[extract_app(RenderApp)]
enum RenderError {
    #[default]
    None,
    OutOfMemory,
    Validation,
    DeviceLost,
    Loop,
}

fn input(
    input: Res<ButtonInput<Key>>,
    mut error: ResMut<RenderError>,
    mut handler: ResMut<RenderErrorHandler>,
) {
    *error = RenderError::None;
    if input.just_pressed(Key::Character("o".into())) {
        *error = RenderError::OutOfMemory;
    }
    if input.just_pressed(Key::Character("v".into())) {
        *error = RenderError::Validation;
    }
    if input.just_pressed(Key::Character("d".into())) {
        *error = RenderError::DeviceLost;
    }
    if input.just_pressed(Key::Character("l".into())) {
        *error = RenderError::Loop;
    }

    if input.just_pressed(Key::Character("1".into())) {
        *handler = RenderErrorHandler(|_, _, _| RenderErrorPolicy::Ignore);
    }
    if input.just_pressed(Key::Character("2".into())) {
        *handler = RenderErrorHandler(|error, _, _| panic!("Rendering error {error:?}"));
    }
    if input.just_pressed(Key::Character("3".into())) {
        *handler = RenderErrorHandler(|_, main_world, _| {
            main_world.write_message(AppExit::error());
            RenderErrorPolicy::StopRendering
        });
    }
    if input.just_pressed(Key::Character("4".into())) {
        *handler = RenderErrorHandler(|_, _, _| RenderErrorPolicy::StopRendering);
    }
    if input.just_pressed(Key::Character("5".into())) {
        *handler = RenderErrorHandler(|_, _, _| RenderErrorPolicy::Recover(default()));
    }
}

fn cause_error(error: If<Res<RenderError>>, device: Res<RenderDevice>, queue: Res<RenderQueue>) {
    match **error {
        RenderError::None => {}
        RenderError::OutOfMemory => {
            let mut textures = Vec::new();
            for _ in 0..64 {
                textures.push(device.create_texture(&TextureDescriptor {
                    label: None,
                    size: Extent3d {
                        width: 8192,
                        height: 8192,
                        depth_or_array_layers: 1,
                    },
                    mip_level_count: 1,
                    sample_count: 1,
                    dimension: TextureDimension::D2,
                    format: TextureFormat::Rgba16Float,
                    usage: TextureUsages::RENDER_ATTACHMENT,
                    view_formats: &[],
                }));
            }
        }
        RenderError::Validation => {
            device.create_buffer(&BufferDescriptor {
                label: None,
                size: 1 << 63,
                usage: BufferUsages::COPY_SRC,
                mapped_at_creation: false,
            });
        }
        RenderError::DeviceLost => {
            device.wgpu_device().destroy();
            device.poll(PollType::wait_indefinitely()).unwrap();
        }
        RenderError::Loop => {
            let sm = device.create_and_validate_shader_module(ShaderModuleDescriptor {
                label: Some("shader"),
                source: ShaderSource::Wgsl(
                    "@compute @workgroup_size(1, 1, 1) fn main() { loop { workgroupBarrier(); } }"
                        .into(),
                ),
            });

            let pipeline_layout = device.create_pipeline_layout(&PipelineLayoutDescriptor {
                label: Some("pipeline_layout"),
                bind_group_layouts: &[],
                immediate_size: 0,
            });

            let pipeline = device.create_compute_pipeline(&RawComputePipelineDescriptor {
                label: Some("pipeline"),
                layout: Some(&pipeline_layout),
                module: &sm,
                entry_point: Some("main"),
                compilation_options: Default::default(),
                cache: None,
            });

            let mut encoder = device.create_command_encoder(&CommandEncoderDescriptor::default());
            {
                let mut cpass = encoder.begin_compute_pass(&ComputePassDescriptor::default());
                cpass.set_pipeline(&pipeline);
                cpass.dispatch_workgroups(1, 1, 1);
            }
            device.poll(PollType::wait_indefinitely()).unwrap();
            queue.submit([encoder.finish()]);
        }
    }
}