use bevy::{
    color::palettes::css::RED,
    pbr::{ExtendedMaterial, MaterialExtension},
    prelude::*,
    render::{render_resource::*, storage::ShaderStorageBuffer},
};

const SHADER_ASSET_PATH: &str = "shaders/material.wgsl";

type VoluMaterial = ExtendedMaterial<StandardMaterial, VoluExtension>;

fn main() {
    App::new()
        .add_plugins(DefaultPlugins)
        .add_plugins(MaterialPlugin::<VoluMaterial>::default())
        .add_systems(Startup, setup)
        .add_systems(Update, (rotate_things, update_volu_material))
        .run();
}

fn setup(
    mut commands: Commands,
    mut meshes: ResMut<Assets<Mesh>>,
    mut materials: ResMut<Assets<VoluMaterial>>,
    mut buffers: ResMut<Assets<ShaderStorageBuffer>>,
) {
    let transform = Transform::from_xyz(0.0, 0.0, 0.0);

    let storage = buffers.add(ShaderStorageBuffer::from(VoluStorage {
        mesh_translation: transform.translation,
        sphere_radius: 1.0,
        color: Vec4::default(),
    }));

    commands.spawn((
        Mesh3d(meshes.add(Cuboid::new(2.0, 2.0, 2.0))),
        MeshMaterial3d(materials.add(ExtendedMaterial {
            base: StandardMaterial {
                base_color: RED.into(),
                alpha_mode: AlphaMode::Blend,
                ..default()
            },
            extension: VoluExtension { storage },
        })),
        transform,
        Rotate,
    ));

    commands.spawn((
        DirectionalLight::default(),
        Transform::from_xyz(1.0, 1.0, 1.0).looking_at(Vec3::ZERO, Vec3::Y),
    ));

    commands.spawn((
        Camera3d::default(),
        Transform::from_xyz(-2.0, 2.5, 5.0).looking_at(Vec3::ZERO, Vec3::Y),
    ));
}

#[derive(Component)]
struct Rotate;

fn rotate_things(mut q: Query<&mut Transform, With<Rotate>>, time: Res<Time>) {
    for mut t in &mut q {
        t.rotate_y(time.delta_secs());
    }
}

fn update_volu_material(
    mut q: Query<(&mut Transform, &MeshMaterial3d<VoluMaterial>), Changed<Transform>>,
    mut volu_materials: ResMut<Assets<VoluMaterial>>,
    mut buffers: ResMut<Assets<ShaderStorageBuffer>>,
    time: Res<Time>,
) {
    for (mut transform, material) in q.iter_mut() {
        transform.translation.y = time.elapsed_secs().sin();

        let volu_material = volu_materials.get_mut(material.0.id()).unwrap();

        let buffer = buffers
            .get_mut(volu_material.extension.storage.id())
            .unwrap();

        let elapsed = time.elapsed_secs_wrapped();

        let volu_storage = VoluStorage {
            mesh_translation: transform.translation,
            sphere_radius: (elapsed.sin() + 3.) / 4.,
            color: vec4((elapsed.sin() + 1.) / 2., 0.0, 0.0, 1.0),
        };

        buffer.set_data(volu_storage);
    }
}

#[derive(Asset, AsBindGroup, Reflect, Debug, Clone)]
struct VoluExtension {
    #[storage(100, read_only)]
    storage: Handle<ShaderStorageBuffer>,
}

impl MaterialExtension for VoluExtension {
    fn fragment_shader() -> ShaderRef {
        SHADER_ASSET_PATH.into()
    }
}

#[derive(ShaderType, Clone, Default, Debug)]
pub struct VoluStorage {
    mesh_translation: Vec3,
    sphere_radius: f32,
    color: Vec4,
}