//! Demonstrates percentage-closer soft shadows (PCSS).

use std::f32::consts::PI;

#[cfg(feature = "free_camera")]
use bevy::camera_controller::free_camera::{FreeCamera, FreeCameraPlugin};
use bevy::{
    anti_alias::taa::TemporalAntiAliasing,
    camera::{
        primitives::{CubemapFrusta, Frustum},
        visibility::{CubemapVisibleEntities, VisibleMeshEntities},
    },
    core_pipeline::prepass::{DepthPrepass, MotionVectorPrepass},
    feathers::{theme::UiTheme, FeathersPlugins},
    light::{ShadowFilteringMethod, Skybox},
    math::vec3,
    prelude::*,
    render::camera::TemporalJitter,
    ui_widgets::{radio_self_update, ValueChange},
};

#[path = "../helpers/radio.rs"]
mod radio;

#[path = "../helpers/theme.rs"]
mod theme;

/// The size of the light, which affects the size of the penumbras.
const LIGHT_RADIUS: f32 = 10.0;

/// The intensity of the point and spot lights.
const POINT_LIGHT_INTENSITY: f32 = 1_000_000_000.0;

/// The range in meters of the point and spot lights.
const POINT_LIGHT_RANGE: f32 = 110.0;

/// The depth bias for directional and spot lights. This value is set higher
/// than the default to avoid shadow acne.
const DIRECTIONAL_SHADOW_DEPTH_BIAS: f32 = 0.20;

/// The depth bias for point lights. This value is set higher than the default to
/// avoid shadow acne.
///
/// Unfortunately, there is a bit of Peter Panning with this value, because of
/// the distance and angle of the light. This can't be helped in this scene
/// without increasing the shadow map size beyond reasonable limits.
const POINT_SHADOW_DEPTH_BIAS: f32 = 0.35;

/// The near Z value for the shadow map, in meters. This is set higher than the
/// default in order to achieve greater resolution in the shadow map for point
/// and spot lights.
const SHADOW_MAP_NEAR_Z: f32 = 50.0;

/// The current application settings (light type, shadow filter, and the status
/// of PCSS).
#[derive(Default, Resource)]
struct AppStatus {
    /// The type of light presently in the scene: either directional or point.
    light_type: LightType,
    /// The type of shadow filter: Gaussian or temporal.
    shadow_filter: ShadowFilter,
    /// Whether soft shadows are enabled.
    soft_shadows: SoftShadows,
}

/// The type of light presently in the scene: directional, point, or spot.
#[derive(Clone, Copy, Default, PartialEq)]
enum LightType {
    /// A directional light, with a cascaded shadow map.
    #[default]
    Directional,
    /// A point light, with a cube shadow map.
    Point,
    /// A spot light, with a cube shadow map.
    Spot,
}

/// The type of shadow filter.
///
/// Generally, `Gaussian` is preferred when temporal antialiasing isn't in use,
/// while `Temporal` is preferred when TAA is in use. In this example, this
/// setting also turns TAA on and off.
#[derive(Clone, Copy, Default, PartialEq)]
enum ShadowFilter {
    /// The non-temporal Gaussian filter (Castano '13 for directional lights, an
    /// analogous alternative for point and spot lights).
    #[default]
    NonTemporal,
    /// The temporal Gaussian filter (Jimenez '14 for directional lights, an
    /// analogous alternative for point and spot lights).
    Temporal,
}

/// Whether soft shadows is enabled. Defaults to `true`.
#[derive(Clone, Copy, Deref, PartialEq)]
struct SoftShadows(bool);

impl Default for SoftShadows {
    fn default() -> Self {
        Self(true)
    }
}

/// The example application entry point.
fn main() {
    #[cfg(not(feature = "free_camera"))]
    println!("Enable feature free_camera to add a free camera to this example");

    App::new()
        .init_resource::<AppStatus>()
        .insert_resource(UiTheme(theme::basic_example_theme(Color::BLACK)))
        .add_plugins((
            DefaultPlugins.set(WindowPlugin {
                primary_window: Some(Window {
                    title: "Bevy Percentage Closer Soft Shadows Example".into(),
                    ..default()
                }),
                ..default()
            }),
            #[cfg(feature = "free_camera")]
            FreeCameraPlugin,
            FeathersPlugins,
        ))
        .add_systems(Startup, setup)
        .add_observer(handle_light_type_change)
        .add_observer(handle_shadow_filter_change)
        .add_observer(handle_pcss_toggle)
        .add_observer(radio_self_update)
        .run();
}

/// Creates all the objects in the scene.
fn setup(mut commands: Commands, asset_server: Res<AssetServer>, app_status: Res<AppStatus>) {
    spawn_camera(&mut commands, &asset_server);
    spawn_light(&mut commands, &app_status);
    spawn_gltf_scene(&mut commands, &asset_server);
    spawn_buttons(&mut commands);
}

/// Spawns the camera, with the initial shadow filtering method.
fn spawn_camera(commands: &mut Commands, asset_server: &AssetServer) {
    commands
        .spawn((
            Camera3d::default(),
            Transform::from_xyz(-12.912 * 0.7, 4.466 * 0.7, -10.624 * 0.7).with_rotation(
                Quat::from_euler(EulerRot::YXZ, -134.76 / 180.0 * PI, -0.175, 0.0),
            ),
            #[cfg(feature = "free_camera")]
            FreeCamera::default(),
        ))
        .insert(ShadowFilteringMethod::Gaussian)
        // `TemporalJitter` is needed for TAA. Note that it does nothing without
        // `TemporalAntiAliasSettings`.
        .insert(TemporalJitter::default())
        // We want MSAA off for TAA to work properly.
        .insert(Msaa::Off)
        // The depth prepass is needed for TAA.
        .insert(DepthPrepass)
        // The motion vector prepass is needed for TAA.
        .insert(MotionVectorPrepass)
        // Add a nice skybox.
        .insert(Skybox {
            image: Some(asset_server.load("environment_maps/sky_skybox.ktx2")),
            brightness: 500.0,
            rotation: Quat::IDENTITY,
        });
}

/// Spawns the initial light.
fn spawn_light(commands: &mut Commands, app_status: &AppStatus) {
    // Because this light can become a directional light, point light, or spot
    // light depending on the settings, we add the union of the components
    // necessary for this light to behave as all three of those.
    commands
        .spawn((
            create_directional_light(app_status),
            Transform::from_rotation(Quat::from_array([
                0.6539259,
                -0.34646285,
                0.36505926,
                -0.5648683,
            ]))
            .with_translation(vec3(57.693, 34.334, -6.422)),
        ))
        // These two are needed for point lights.
        .insert(CubemapVisibleEntities::default())
        .insert(CubemapFrusta::default())
        // These two are needed for spot lights.
        .insert(VisibleMeshEntities::default())
        .insert(Frustum::default());
}

/// Loads and spawns the glTF palm tree scene.
fn spawn_gltf_scene(commands: &mut Commands, asset_server: &AssetServer) {
    commands.spawn(WorldAssetRoot(
        asset_server.load("models/PalmTree/PalmTree.gltf#Scene0"),
    ));
}

/// Spawns all the buttons at the bottom of the screen.
fn spawn_buttons(commands: &mut Commands) {
    commands.spawn_scene(bsn! {
        @radio::main_ui_node_scene()
        Children [
            @radio::feathers_option_buttons(
                "Light Type",
                &[
                    (LightType::Directional, "Directional"),
                    (LightType::Point, "Point"),
                    (LightType::Spot, "Spot"),
                ],
                0,
            )
            --
            @radio::feathers_option_buttons(
                "Shadow Filter",
                &[
                    (
                        ShadowFilter::NonTemporal,
                        "Non-Temporal",
                    ),
                    (ShadowFilter::Temporal, "Temporal"),
                ],
                0,
            )
            --
            @radio::feathers_option_buttons(
                "Soft Shadows",
                &[
                    (SoftShadows(true), "On"),
                    (SoftShadows(false), "Off"),
                ],
                0,
            )
        ]
    });
}

/// Handles requests from the user to change the type of light.
fn handle_light_type_change(
    event: On<ValueChange<Entity>>,
    new_value_query: Query<&radio::RadioButtonOptionValue<LightType>>,
    mut commands: Commands,
    mut lights: Query<Entity, Or<(With<DirectionalLight>, With<PointLight>, With<SpotLight>)>>,
    mut app_status: ResMut<AppStatus>,
) {
    let Ok(radio::RadioButtonOptionValue(light_type)) = new_value_query.get(event.value) else {
        return;
    };
    app_status.light_type = *light_type;

    for light in lights.iter_mut() {
        let mut light_commands = commands.entity(light);
        light_commands
            .remove::<DirectionalLight>()
            .remove::<PointLight>()
            .remove::<SpotLight>();
        match light_type {
            LightType::Point => {
                light_commands.insert(create_point_light(&app_status));
            }
            LightType::Spot => {
                light_commands.insert(create_spot_light(&app_status));
            }
            LightType::Directional => {
                light_commands.insert(create_directional_light(&app_status));
            }
        }
    }
}

/// Handles requests from the user to change the shadow filter method.
///
/// This system is also responsible for enabling and disabling TAA as
/// appropriate.
fn handle_shadow_filter_change(
    event: On<ValueChange<Entity>>,
    new_value_query: Query<&radio::RadioButtonOptionValue<ShadowFilter>>,
    mut commands: Commands,
    mut cameras: Query<(Entity, &mut ShadowFilteringMethod)>,
    mut app_status: ResMut<AppStatus>,
) {
    let Ok(radio::RadioButtonOptionValue(shadow_filter)) = new_value_query.get(event.value) else {
        return;
    };

    app_status.shadow_filter = *shadow_filter;

    for (camera, mut shadow_filtering_method) in cameras.iter_mut() {
        match shadow_filter {
            ShadowFilter::NonTemporal => {
                *shadow_filtering_method = ShadowFilteringMethod::Gaussian;
                commands.entity(camera).remove::<TemporalAntiAliasing>();
            }
            ShadowFilter::Temporal => {
                *shadow_filtering_method = ShadowFilteringMethod::Temporal;
                commands
                    .entity(camera)
                    .insert(TemporalAntiAliasing::default());
            }
        }
    }
}

/// Handles requests from the user to toggle soft shadows on and off.
fn handle_pcss_toggle(
    event: On<ValueChange<Entity>>,
    new_value_query: Query<&radio::RadioButtonOptionValue<SoftShadows>>,
    mut lights: Query<AnyOf<(&mut DirectionalLight, &mut PointLight, &mut SpotLight)>>,
    mut app_status: ResMut<AppStatus>,
) {
    let Ok(radio::RadioButtonOptionValue(soft_shadows)) = new_value_query.get(event.value) else {
        return;
    };
    app_status.soft_shadows = *soft_shadows;

    // Recreating the lights is the simplest way to toggle soft shadows.
    for (directional_light, point_light, spot_light) in lights.iter_mut() {
        if let Some(mut directional_light) = directional_light {
            *directional_light = create_directional_light(&app_status);
        }
        if let Some(mut point_light) = point_light {
            *point_light = create_point_light(&app_status);
        }
        if let Some(mut spot_light) = spot_light {
            *spot_light = create_spot_light(&app_status);
        }
    }
}

/// Creates the [`DirectionalLight`] component with the appropriate settings.
fn create_directional_light(app_status: &AppStatus) -> DirectionalLight {
    DirectionalLight {
        shadow_maps_enabled: true,
        soft_shadow_size: if *app_status.soft_shadows {
            Some(LIGHT_RADIUS)
        } else {
            None
        },
        shadow_depth_bias: DIRECTIONAL_SHADOW_DEPTH_BIAS,
        ..default()
    }
}

/// Creates the [`PointLight`] component with the appropriate settings.
fn create_point_light(app_status: &AppStatus) -> PointLight {
    PointLight {
        intensity: POINT_LIGHT_INTENSITY,
        range: POINT_LIGHT_RANGE,
        shadow_maps_enabled: true,
        radius: LIGHT_RADIUS,
        soft_shadows_enabled: *app_status.soft_shadows,
        shadow_depth_bias: POINT_SHADOW_DEPTH_BIAS,
        shadow_map_near_z: SHADOW_MAP_NEAR_Z,
        ..default()
    }
}

/// Creates the [`SpotLight`] component with the appropriate settings.
fn create_spot_light(app_status: &AppStatus) -> SpotLight {
    SpotLight {
        intensity: POINT_LIGHT_INTENSITY,
        range: POINT_LIGHT_RANGE,
        radius: LIGHT_RADIUS,
        shadow_maps_enabled: true,
        soft_shadows_enabled: *app_status.soft_shadows,
        shadow_depth_bias: DIRECTIONAL_SHADOW_DEPTH_BIAS,
        shadow_map_near_z: SHADOW_MAP_NEAR_Z,
        ..default()
    }
}