use bevy::pbr::ExtendedMaterial;
use bevy::pbr::MaterialExtension;
use bevy::pbr::OpaqueRendererMethod;
use bevy::prelude::*;
use bevy::render::render_resource::*;

type MyMat = ExtendedMaterial<StandardMaterial, MatExt>;

fn main() {
    App::new()
        .add_plugins((
            DefaultPlugins.set(ImagePlugin::default_nearest()),
            MaterialPlugin::<MyMat>::default(),
        ))
        .add_systems(Startup, setup)
        .add_systems(Update, rotate)
        .add_systems(Update, toggle_material.run_if(
            |keys: Res<Input<KeyCode>>| -> bool {
                keys.just_pressed(KeyCode::Space)
            })
        )
        .add_systems(Update, init_materials.run_if(|events: Query<Entity, Added<Transform>>| -> bool {
            !events.is_empty()
        }))
        .run();
}

fn setup(mut commands: Commands, assets: Res<AssetServer>) {
    commands.spawn(SceneBundle {
        scene: assets.load("models/Martian Chess.glb#Scene0"),
        ..default()
    });

    commands.spawn(PointLightBundle {
        point_light: PointLight {
            intensity: 10.0,
            ..default()
        },
        transform: Transform::from_xyz(-1.0, 4.5, 4.0).looking_at(Vec3::ZERO, Vec3::Y),
        ..default()
    });

    commands.spawn(Camera3dBundle {
        transform: Transform::from_xyz(-1.0, 1.5, 4.0).looking_at(Vec3::ZERO, Vec3::Y),
        ..default()
    });
}

#[derive(Asset, AsBindGroup, Reflect, Debug, Clone)]
struct MatExt {
    #[uniform(100)]
    cutoff: f32,
}

impl MaterialExtension for MatExt {
    fn fragment_shader() -> ShaderRef {
        "examples/shaders/dissolve.wgsl".into()
    }
}


#[derive(Debug, Component)]
struct Backup<T: Component>(T);

fn init_materials(
    events: Query<(Entity, &Handle<StandardMaterial>), Added<Handle<StandardMaterial>>>,
    standard_materials: Res<Assets<StandardMaterial>>,
    mut materials: ResMut<Assets<MyMat>>,
    mut commands: Commands,
) {
    info!("initializing materials");

    events.iter().for_each(|(entity, std_handle)| {
        // Extension we will add to existing gltf-sourced materials
        let extension = MatExt { cutoff: 1.0 };
        // Base material we will extend for the duration of the dissolve effect
        let mut base = standard_materials
            .get(std_handle)
            .expect("Resolve material data")
            .clone();

        base.opaque_render_method = OpaqueRendererMethod::Auto;
        base.alpha_mode = AlphaMode::Mask(0.5);

        let ext_handle = materials.add(ExtendedMaterial { base, extension });

        commands
            .entity(entity)
            .insert(Backup(ext_handle.clone()));
    });
}

fn toggle_material(
    query: Query<Entity, Or<(With<Handle<StandardMaterial>>, With<Handle<MyMat>>)>>,
    std_mat: Query<(&Handle<StandardMaterial>, Option<&Backup<Handle<MyMat>>>)>,
    ext_mat: Query<(&Handle<MyMat>, Option<&Backup<Handle<StandardMaterial>>>)>,
    mut commands: Commands,
) {
    query.iter().for_each(|entity| {
        // Entity has standard material, and had extended material
        // Swap these materials
        if let Ok(
            (std_handle, Some(Backup(ext_handle)))
        ) = std_mat.get(entity) {
            info!("Swapping standard material for extended material");
            commands
                .entity(entity)
                .insert(ext_handle.clone())
                .insert(Backup(std_handle.clone()))
                .remove::<Backup<Handle<MyMat>>>()
                .remove::<Handle<StandardMaterial>>();

            return
        }

        // In this branch we have the extended material assigned to the object
        if let Ok(
            (ext_handle, Some(Backup(std_handle)))
        ) = ext_mat.get(entity) {
            // Entity has standard material, and had extended material // Swap these materials
            info!("Swapping extended material for standard material");

            commands
                .entity(entity)
                .insert(std_handle.clone())
                .insert(Backup(ext_handle.clone()))
                .remove::<Backup<Handle<StandardMaterial>>>()
                .remove::<Handle<MyMat>>();

            return
        }

        panic!("What is happening?")
    });
}

fn rotate(
    mut query: Query<&mut Transform, With<Handle<Mesh>>>,
    time: Res<Time>,
    mut materials: ResMut<Assets<MyMat>>,
) {
    query.iter_mut().for_each(|mut t| {
        t.rotate_local_y(time.delta_seconds() / 2.0);
    });
    materials.iter_mut().for_each(|(_id, m)| {
        m.extension.cutoff = time.elapsed_seconds().sin().abs();
    })
}