//! Chunked render and collider hydration for authored height-grid terrain.

use avian3d::prelude::ColliderConstructor;
use bevy::asset::RenderAssetUsages;
use bevy::light::{NotShadowCaster, NotShadowReceiver};
use bevy::mesh::{Indices, PrimitiveTopology};
use bevy::prelude::*;

use crate::{
    authoring_component_active, load_resolved_material_from_path, AuthoringComponentStates,
    HydratedTerrainMaterialBinding, InspectorOrder, LevelObject, RaytracingExcluded, TerrainDesc,
    COMPONENT_TERRAIN_DESC,
};

use super::materials::material_from_desc;

#[derive(Component, Reflect, Debug, Clone, Copy, PartialEq, Eq)]
#[reflect(Component, Default, Debug)]
pub struct HydratedTerrainChunk {
    pub owner: Entity,
    pub chunk_x: u32,
    pub chunk_z: u32,
}

impl Default for HydratedTerrainChunk {
    fn default() -> Self {
        Self {
            owner: Entity::PLACEHOLDER,
            chunk_x: 0,
            chunk_z: 0,
        }
    }
}

#[allow(clippy::type_complexity)]
pub fn hydrate_terrain(
    mut commands: Commands,
    asset_server: Res<AssetServer>,
    mut meshes: ResMut<Assets<Mesh>>,
    mut materials: ResMut<Assets<StandardMaterial>>,
    terrains: Query<
        (
            Entity,
            &TerrainDesc,
            Option<&AuthoringComponentStates>,
            Option<&InspectorOrder>,
        ),
        (
            With<LevelObject>,
            Or<(
                Added<TerrainDesc>,
                Changed<TerrainDesc>,
                Changed<AuthoringComponentStates>,
                Changed<InspectorOrder>,
            )>,
        ),
    >,
    children: Query<&Children>,
    chunks: Query<(), With<HydratedTerrainChunk>>,
) {
    for (entity, terrain, states, order) in &terrains {
        despawn_terrain_chunks(&mut commands, entity, &children, &chunks);
        if authoring_component_active(states, order, COMPONENT_TERRAIN_DESC) {
            spawn_terrain_chunks(
                &mut commands,
                &asset_server,
                &mut meshes,
                &mut materials,
                entity,
                terrain,
            );
        }
    }
}

pub fn cleanup_removed_terrain(
    mut commands: Commands,
    mut removed: RemovedComponents<TerrainDesc>,
    children: Query<&Children>,
    chunks: Query<(), With<HydratedTerrainChunk>>,
) {
    for entity in removed.read() {
        despawn_terrain_chunks(&mut commands, entity, &children, &chunks);
    }
}

pub fn spawn_terrain_chunks(
    commands: &mut Commands,
    asset_server: &AssetServer,
    meshes: &mut Assets<Mesh>,
    materials: &mut Assets<StandardMaterial>,
    owner: Entity,
    terrain: &TerrainDesc,
) {
    if let Err(error) = terrain.validate() {
        warn!("Terrain hydration skipped invalid actor {owner:?}: {error}");
        return;
    }

    let effective_layers = terrain.effective_material_layers();
    let material = effective_layers
        .first()
        .and_then(|layer| layer.material.as_ref())
        .and_then(|reference| reference.source_path.as_deref())
        .and_then(|path| match load_resolved_material_from_path(path) {
            Ok((desc, _)) => Some(material_from_desc(asset_server, &desc)),
            Err(error) => {
                warn!("Terrain {owner:?} material could not resolve: {error}");
                None
            }
        })
        .unwrap_or_else(|| StandardMaterial {
            base_color: Color::srgb(0.24, 0.29, 0.25),
            perceptual_roughness: 0.92,
            ..default()
        });
    let material = materials.add(material);
    let quad_count = terrain.resolution - 1;
    let chunks_per_axis = quad_count.div_ceil(terrain.chunk_quads);

    for chunk_z in 0..chunks_per_axis {
        for chunk_x in 0..chunks_per_axis {
            let start_x = chunk_x * terrain.chunk_quads;
            let start_z = chunk_z * terrain.chunk_quads;
            let end_x = (start_x + terrain.chunk_quads).min(quad_count);
            let end_z = (start_z + terrain.chunk_quads).min(quad_count);
            let Some(mesh) = terrain_chunk_mesh(terrain, start_x, start_z, end_x, end_z) else {
                continue;
            };
            let mesh = meshes.add(mesh);
            let mut chunk = commands.spawn((
                HydratedTerrainChunk {
                    owner,
                    chunk_x,
                    chunk_z,
                },
                Name::new(format!("Terrain Chunk {chunk_x},{chunk_z}")),
                Mesh3d(mesh),
                MeshMaterial3d(material.clone()),
                HydratedTerrainMaterialBinding {
                    owner,
                    layers: terrain.effective_material_layers(),
                },
                Transform::default(),
                Visibility::Visible,
                RaytracingExcluded,
                ChildOf(owner),
            ));
            if terrain.generate_colliders {
                chunk.insert(ColliderConstructor::TrimeshFromMesh);
            }
            if !terrain.cast_shadows {
                chunk.insert(NotShadowCaster);
            }
            if !terrain.receive_shadows {
                chunk.insert(NotShadowReceiver);
            }
        }
    }
}

fn despawn_terrain_chunks(
    commands: &mut Commands,
    entity: Entity,
    children: &Query<&Children>,
    chunks: &Query<(), With<HydratedTerrainChunk>>,
) {
    let Ok(children) = children.get(entity) else {
        return;
    };
    for child in children.iter() {
        if chunks.get(child).is_ok() {
            commands.entity(child).despawn();
        }
    }
}

pub fn terrain_chunk_mesh(
    terrain: &TerrainDesc,
    start_x: u32,
    start_z: u32,
    end_x: u32,
    end_z: u32,
) -> Option<Mesh> {
    if start_x >= end_x
        || start_z >= end_z
        || end_x >= terrain.resolution
        || end_z >= terrain.resolution
    {
        return None;
    }
    let width = end_x - start_x + 1;
    let depth = end_z - start_z + 1;
    let mut positions = Vec::with_capacity((width * depth) as usize);
    let mut normals = Vec::with_capacity((width * depth) as usize);
    let mut uvs = Vec::with_capacity((width * depth) as usize);
    let mut colors = Vec::with_capacity((width * depth) as usize);
    let quad_count = (terrain.resolution - 1) as f32;
    let half_extent = quad_count * terrain.sample_spacing * 0.5;

    let sample = |x: u32, z: u32| {
        terrain.heights[(z * terrain.resolution + x) as usize] * terrain.height_scale
    };
    for z in start_z..=end_z {
        for x in start_x..=end_x {
            let left = sample(x.saturating_sub(1), z);
            let right = sample((x + 1).min(terrain.resolution - 1), z);
            let down = sample(x, z.saturating_sub(1));
            let up = sample(x, (z + 1).min(terrain.resolution - 1));
            let normal = Vec3::new(left - right, 2.0 * terrain.sample_spacing, down - up)
                .normalize_or_zero();
            positions.push([
                x as f32 * terrain.sample_spacing - half_extent,
                sample(x, z),
                z as f32 * terrain.sample_spacing - half_extent,
            ]);
            normals.push(normal.to_array());
            uvs.push([x as f32 / quad_count, z as f32 / quad_count]);
            colors.push(terrain.normalized_material_weights((z * terrain.resolution + x) as usize));
        }
    }

    let mut indices = Vec::with_capacity(((width - 1) * (depth - 1) * 6) as usize);
    for z in 0..depth - 1 {
        for x in 0..width - 1 {
            let a = z * width + x;
            let b = a + 1;
            let c = a + width;
            let d = c + 1;
            indices.extend_from_slice(&[a, c, b, b, c, d]);
        }
    }

    let mut mesh = Mesh::new(
        PrimitiveTopology::TriangleList,
        RenderAssetUsages::default(),
    );
    mesh.insert_attribute(Mesh::ATTRIBUTE_POSITION, positions);
    mesh.insert_attribute(Mesh::ATTRIBUTE_NORMAL, normals);
    mesh.insert_attribute(Mesh::ATTRIBUTE_UV_0, uvs);
    mesh.insert_attribute(Mesh::ATTRIBUTE_COLOR, colors);
    mesh.insert_indices(Indices::U32(indices));
    if let Err(error) = mesh.generate_tangents() {
        warn!("Terrain chunk tangent generation failed: {error:?}");
    }
    Some(mesh)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ActorKind;

    #[test]
    fn chunk_mesh_has_expected_topology() {
        let terrain = TerrainDesc::flat(5);
        let mesh = terrain_chunk_mesh(&terrain, 0, 0, 4, 4).unwrap();
        assert_eq!(mesh.count_vertices(), 25);
        assert_eq!(mesh.indices().unwrap().len(), 96);
    }

    #[test]
    fn descriptor_rejects_mismatched_and_non_finite_heights() {
        let mut terrain = TerrainDesc::flat(3);
        terrain.heights.pop();
        assert!(terrain.validate().unwrap_err().contains("requires 9"));
        terrain.heights.push(f32::NAN);
        assert!(terrain.validate().unwrap_err().contains("finite"));
    }

    #[test]
    fn descriptor_round_trips_through_ron() {
        let mut terrain = TerrainDesc::flat(5);
        terrain.heights[12] = 0.75;
        terrain.sample_spacing = 2.0;
        terrain.material_layers = vec![
            crate::TerrainMaterialLayer {
                material: Some(crate::EditorAssetRef::new(
                    "ground",
                    "material:source",
                    "Ground",
                )),
                uv_scale: 6.0,
            },
            crate::TerrainMaterialLayer {
                material: Some(crate::EditorAssetRef::new(
                    "rock",
                    "material:source",
                    "Rock",
                )),
                uv_scale: 12.0,
            },
        ];
        terrain.ensure_material_weights();
        terrain.material_weights[12] = [64, 191, 0, 0];
        let text = ron::to_string(&terrain).unwrap();
        let restored: TerrainDesc = ron::from_str(&text).unwrap();
        assert_eq!(restored, terrain);
        assert!(restored.validate().is_ok());
    }

    #[test]
    fn chunk_mesh_transports_normalized_material_weights_as_vertex_colors() {
        let mut terrain = TerrainDesc::flat(3);
        terrain.ensure_material_weights();
        terrain.material_weights[4] = [64, 191, 0, 0];
        let mesh = terrain_chunk_mesh(&terrain, 0, 0, 2, 2).unwrap();
        let Some(bevy::mesh::VertexAttributeValues::Float32x4(colors)) =
            mesh.attribute(Mesh::ATTRIBUTE_COLOR)
        else {
            panic!("terrain chunks should carry RGBA layer weights");
        };
        assert_eq!(colors.len(), 9);
        assert!((colors[4][0] - 64.0 / 255.0).abs() < 0.0001);
        assert!((colors[4][1] - 191.0 / 255.0).abs() < 0.0001);
    }

    #[test]
    fn descriptor_rejects_non_normalized_material_weights() {
        let mut terrain = TerrainDesc::flat(3);
        terrain.material_weights = vec![[128, 128, 0, 0]; 9];
        assert!(terrain.validate().unwrap_err().contains("sum to 255"));

        terrain.material_weights = vec![[127, 0, 128, 0]; 9];
        assert!(terrain
            .validate()
            .unwrap_err()
            .contains("unassigned layer channel"));
    }

    #[test]
    fn hydration_generates_runtime_only_chunks_and_cleans_up() {
        let mut app = App::new();
        app.add_plugins((MinimalPlugins, AssetPlugin::default()))
            .init_asset::<Mesh>()
            .init_asset::<StandardMaterial>()
            .add_systems(Update, (hydrate_terrain, cleanup_removed_terrain).chain());
        let owner = app
            .world_mut()
            .spawn((
                LevelObject,
                ActorKind::Terrain,
                TerrainDesc {
                    resolution: 5,
                    heights: vec![0.0; 25],
                    chunk_quads: 2,
                    ..Default::default()
                },
                Transform::default(),
            ))
            .id();

        app.update();
        app.update();
        let children = app.world().get::<Children>(owner).unwrap();
        assert_eq!(children.len(), 4);
        for child in children.iter() {
            assert!(app.world().get::<HydratedTerrainChunk>(child).is_some());
            assert!(app.world().get::<LevelObject>(child).is_none());
            assert!(app.world().get::<Mesh3d>(child).is_some());
            assert!(app.world().get::<ColliderConstructor>(child).is_some());
        }

        app.world_mut().entity_mut(owner).remove::<TerrainDesc>();
        app.update();
        app.update();
        assert!(app
            .world()
            .get::<Children>(owner)
            .is_none_or(Children::is_empty));
    }
}
