use super::RaytracingMesh3d;
use bevy_asset::{AssetId, Assets};
use bevy_derive::Deref;
use bevy_ecs::{
    component::Component,
    resource::Resource,
    system::{Commands, Query},
};
use bevy_pbr::{MeshMaterial3d, PreviousGlobalTransform, StandardMaterial};
use bevy_platform::collections::HashMap;
use bevy_render::{Extract, extract_resource::ExtractResource, sync_world::RenderEntity};
use bevy_transform::components::GlobalTransform;
use blacksite_surface::SurfaceMaterial;

pub fn extract_raytracing_scene(
    instances: Extract<
        Query<(
            RenderEntity,
            &RaytracingMesh3d,
            Option<&MeshMaterial3d<StandardMaterial>>,
            Option<&MeshMaterial3d<SurfaceMaterial>>,
            &GlobalTransform,
            Option<&PreviousGlobalTransform>,
        )>,
    >,
    mut commands: Commands,
) {
    for (render_entity, mesh, standard, surface, transform, previous_frame_transform) in &instances
    {
        let material = if let Some(surface) = surface {
            RaytracingMaterial3d::Surface(surface.0.clone())
        } else if let Some(standard) = standard {
            RaytracingMaterial3d::Standard(standard.0.clone())
        } else {
            continue;
        };
        let mut commands = commands.entity(render_entity);

        match previous_frame_transform.cloned() {
            Some(previous_frame_transform) => commands.insert((
                mesh.clone(),
                material.clone(),
                *transform,
                previous_frame_transform,
            )),
            None => commands.insert((mesh.clone(), material.clone(), *transform)),
        };
    }
}

/// Render-world material identity selected during extraction.
#[derive(Component, Clone, Debug)]
pub enum RaytracingMaterial3d {
    Standard(bevy_asset::Handle<StandardMaterial>),
    Surface(bevy_asset::Handle<SurfaceMaterial>),
}

#[derive(Resource, Deref, Default)]
pub struct StandardMaterialAssets(HashMap<AssetId<StandardMaterial>, StandardMaterial>);

impl ExtractResource for StandardMaterialAssets {
    type Source = Assets<StandardMaterial>;

    fn extract_resource(source: &Self::Source) -> Self {
        Self(
            source
                .iter()
                .map(|(asset_id, material)| (asset_id, material.clone()))
                .collect(),
        )
    }
}

#[derive(Resource, Deref, Default)]
pub struct SurfaceMaterialAssets(HashMap<AssetId<SurfaceMaterial>, SurfaceMaterial>);

impl ExtractResource for SurfaceMaterialAssets {
    type Source = Assets<SurfaceMaterial>;

    fn extract_resource(source: &Self::Source) -> Self {
        Self(
            source
                .iter()
                .map(|(asset_id, material)| (asset_id, material.clone()))
                .collect(),
        )
    }
}
