//! Immediate-mode filled debug geometry.
//!
//! Unlike Bevy's line-only `Gizmos`, this collects triangles into one dynamic
//! mesh. It is intended for small, transient debug overlays rather than game
//! geometry.

use std::f32::consts::{FRAC_PI_2, PI, TAU};

use bevy::{
    asset::{load_internal_asset, uuid_handle, RenderAssetUsages},
    camera::visibility::NoFrustumCulling,
    ecs::system::SystemParam,
    mesh::{Indices, MeshVertexBufferLayoutRef, PrimitiveTopology},
    pbr::{MaterialPipeline, MaterialPipelineKey},
    prelude::*,
    render::render_resource::{
        AsBindGroup, RenderPipelineDescriptor, SpecializedMeshPipelineError,
    },
    shader::ShaderRef,
};

const SPHERE_SECTORS: usize = 16;
const SPHERE_STACKS: usize = 8;
const CAPSULE_SECTORS: usize = 16;
const CAPSULE_HEMISPHERE_STEPS: usize = 4;

#[derive(Default)]
struct FilledGizmoGeometry {
    positions: Vec<[f32; 3]>,
    normals: Vec<[f32; 3]>,
    colors: Vec<[f32; 4]>,
    indices: Vec<u32>,
}

#[derive(Resource, Default)]
struct FilledGizmoBuffer {
    /// Geometry requested by the current schedule invocation.
    pending: FilledGizmoGeometry,
    /// The last completed fixed-timestep geometry. This is retained between
    /// render frames, just like Bevy's fixed-timestep gizmos.
    retained: FilledGizmoGeometry,
    mesh: Option<Handle<Mesh>>,
    fixed_commit: bool,
}

/// Immediate-mode API for drawing filled debug shapes.
#[derive(SystemParam)]
pub struct FilledGizmos<'w> {
    buffer: ResMut<'w, FilledGizmoBuffer>,
}

impl FilledGizmos<'_> {
    /// Draw a filled planar polygon. Vertices must be ordered around the
    /// perimeter and the polygon must be convex.
    pub fn polygon(&mut self, vertices: &[Vec3], color: impl Into<Color> + Copy) {
        if vertices.len() < 3 {
            return;
        }
        let color = color.into().to_linear().to_f32_array();
        let normal = (vertices[1] - vertices[0])
            .cross(vertices[2] - vertices[0])
            .normalize_or_zero();
        let base = self.buffer.pending.positions.len() as u32;
        for &vertex in vertices {
            self.push_vertex(vertex, normal, color);
        }
        for index in 2..vertices.len() as u32 {
            self.buffer
                .pending
                .indices
                .extend([base, base + index - 1, base + index]);
        }
    }

    pub fn sphere(&mut self, isometry: Isometry3d, radius: f32, color: impl Into<Color> + Copy) {
        if radius <= 0.0 {
            return;
        }
        let color = color.into().to_linear().to_f32_array();
        let base = self.buffer.pending.positions.len() as u32;
        for stack in 0..=SPHERE_STACKS {
            let latitude = PI * 0.5 - stack as f32 * PI / SPHERE_STACKS as f32;
            let (sin_lat, cos_lat) = latitude.sin_cos();
            for sector in 0..=SPHERE_SECTORS {
                let longitude = sector as f32 * TAU / SPHERE_SECTORS as f32;
                let (sin_lon, cos_lon) = longitude.sin_cos();
                let normal = Vec3::new(cos_lat * cos_lon, sin_lat, cos_lat * sin_lon);
                self.push_vertex(
                    isometry.transform_point(normal * radius).into(),
                    isometry.rotation * normal,
                    color,
                );
            }
        }
        let stride = SPHERE_SECTORS + 1;
        for stack in 0..SPHERE_STACKS {
            for sector in 0..SPHERE_SECTORS {
                let a = base + (stack * stride + sector) as u32;
                let b = a + stride as u32;
                self.buffer
                    .pending
                    .indices
                    .extend([a, b, a + 1, a + 1, b, b + 1]);
            }
        }
    }

    pub fn capsule(
        &mut self,
        isometry: Isometry3d,
        radius: f32,
        half_length: f32,
        color: impl Into<Color> + Copy,
    ) {
        let transform = Mat4::from_scale_rotation_translation(
            Vec3::ONE,
            isometry.rotation,
            isometry.translation.into(),
        );
        self.affine_capsule(transform, radius, half_length, color);
    }

    /// Draw a capsule transformed by an arbitrary affine transform.
    pub fn affine_capsule(
        &mut self,
        transform: Mat4,
        radius: f32,
        half_length: f32,
        color: impl Into<Color> + Copy,
    ) {
        if radius <= 0.0 {
            return;
        }
        let color = color.into().to_linear().to_f32_array();
        let base = self.buffer.pending.positions.len() as u32;
        let mut rings = Vec::new();
        for step in 1..=CAPSULE_HEMISPHERE_STEPS {
            let angle = -FRAC_PI_2 + step as f32 * FRAC_PI_2 / CAPSULE_HEMISPHERE_STEPS as f32;
            rings.push((
                radius * angle.cos(),
                -half_length + radius * angle.sin(),
                angle.sin(),
            ));
        }
        rings.push((radius, half_length, 0.0));
        for step in 1..CAPSULE_HEMISPHERE_STEPS {
            let angle = step as f32 * FRAC_PI_2 / CAPSULE_HEMISPHERE_STEPS as f32;
            rings.push((
                radius * angle.cos(),
                half_length + radius * angle.sin(),
                angle.sin(),
            ));
        }

        let mut ring_indices = Vec::new();
        let bottom_pole = self.push_capsule_vertex(
            &mut ring_indices,
            transform,
            Vec3::NEG_Y,
            -half_length,
            radius,
            color,
        );
        for &(ring_radius, y, normal_y) in &rings {
            let mut ring = Vec::with_capacity(CAPSULE_SECTORS);
            for sector in 0..CAPSULE_SECTORS {
                let angle = sector as f32 * TAU / CAPSULE_SECTORS as f32;
                let (sin, cos) = angle.sin_cos();
                let normal = Vec3::new(
                    cos * (1.0 - normal_y * normal_y).sqrt(),
                    normal_y,
                    sin * (1.0 - normal_y * normal_y).sqrt(),
                );
                ring.push(self.push_transformed_vertex(
                    transform,
                    Vec3::new(cos * ring_radius, y, sin * ring_radius),
                    normal,
                    color,
                ));
            }
            ring_indices.push(ring);
        }
        let top_pole = self.push_capsule_vertex(
            &mut ring_indices,
            transform,
            Vec3::Y,
            half_length,
            radius,
            color,
        );
        for sector in 0..CAPSULE_SECTORS {
            let next = (sector + 1) % CAPSULE_SECTORS;
            self.buffer.pending.indices.extend([
                bottom_pole,
                ring_indices[0][next],
                ring_indices[0][sector],
            ]);
            for ring in ring_indices.windows(2) {
                self.buffer.pending.indices.extend([
                    ring[0][sector],
                    ring[0][next],
                    ring[1][sector],
                    ring[1][sector],
                    ring[0][next],
                    ring[1][next],
                ]);
            }
            let last = ring_indices.len() - 1;
            self.buffer.pending.indices.extend([
                ring_indices[last][sector],
                ring_indices[last][next],
                top_pole,
            ]);
        }
        // Keep the base used by this shape observable to avoid accidental
        // future changes making the generated indices refer to another shape.
        debug_assert!(base <= bottom_pole);
    }

    fn push_capsule_vertex(
        &mut self,
        _rings: &mut Vec<Vec<u32>>,
        transform: Mat4,
        normal: Vec3,
        y: f32,
        radius: f32,
        color: [f32; 4],
    ) -> u32 {
        self.push_transformed_vertex(
            transform,
            Vec3::new(0.0, y + normal.y * radius, 0.0),
            normal,
            color,
        )
    }

    fn push_transformed_vertex(
        &mut self,
        transform: Mat4,
        position: Vec3,
        normal: Vec3,
        color: [f32; 4],
    ) -> u32 {
        let transformed_normal = transform
            .inverse()
            .transpose()
            .transform_vector3(normal)
            .normalize_or_zero();
        let index = self.buffer.pending.positions.len() as u32;
        self.push_vertex(
            transform.transform_point3(position),
            transformed_normal,
            color,
        );
        index
    }

    fn push_vertex(&mut self, position: Vec3, normal: Vec3, color: [f32; 4]) {
        self.buffer.pending.positions.push(position.to_array());
        self.buffer.pending.normals.push(normal.to_array());
        self.buffer.pending.colors.push(color);
    }
}

pub struct FilledGizmosPlugin;

const FILLED_GIZMO_SHADER_HANDLE: Handle<Shader> =
    uuid_handle!("0e4a979f-ea11-4e5a-91cd-19c89d4e9b53");

/// A deliberately minimal material: unlike `StandardMaterial`, its fragment
/// shader writes vertex colors directly and therefore bypasses PBR lighting
/// and camera tonemapping.
#[derive(Asset, AsBindGroup, Reflect, Clone, Debug)]
struct FilledGizmoMaterial {
    #[uniform(0)]
    color: LinearRgba,
}

impl Material for FilledGizmoMaterial {
    fn vertex_shader() -> ShaderRef {
        FILLED_GIZMO_SHADER_HANDLE.clone().into()
    }

    fn fragment_shader() -> ShaderRef {
        FILLED_GIZMO_SHADER_HANDLE.clone().into()
    }

    fn alpha_mode(&self) -> AlphaMode {
        AlphaMode::Blend
    }

    fn enable_prepass() -> bool {
        false
    }

    fn enable_shadows() -> bool {
        false
    }

    fn specialize(
        _pipeline: &MaterialPipeline,
        descriptor: &mut RenderPipelineDescriptor,
        layout: &MeshVertexBufferLayoutRef,
        _key: MaterialPipelineKey<Self>,
    ) -> Result<(), SpecializedMeshPipelineError> {
        descriptor.vertex.buffers = vec![layout.0.get_layout(&[
            Mesh::ATTRIBUTE_POSITION.at_shader_location(0),
            Mesh::ATTRIBUTE_COLOR.at_shader_location(1),
        ])?];
        descriptor.primitive.cull_mode = None;
        Ok(())
    }
}

impl Plugin for FilledGizmosPlugin {
    fn build(&self, app: &mut App) {
        load_internal_asset!(
            app,
            FILLED_GIZMO_SHADER_HANDLE,
            "filled_gizmo.wgsl",
            Shader::from_wgsl
        );
        app.add_plugins(MaterialPlugin::<FilledGizmoMaterial>::default())
            .init_resource::<FilledGizmoBuffer>()
            .add_systems(Startup, setup_filled_gizmo_mesh)
            .add_systems(FixedFirst, begin_fixed_filled_gizmos)
            .add_systems(FixedLast, commit_fixed_filled_gizmos)
            .add_systems(Last, upload_filled_gizmo_mesh);
    }
}

fn setup_filled_gizmo_mesh(
    mut commands: Commands,
    mut meshes: ResMut<Assets<Mesh>>,
    mut materials: ResMut<Assets<FilledGizmoMaterial>>,
    mut buffer: ResMut<FilledGizmoBuffer>,
) {
    let mesh = meshes.add(
        Mesh::new(
            PrimitiveTopology::TriangleList,
            RenderAssetUsages::default(),
        )
        .with_inserted_attribute(Mesh::ATTRIBUTE_POSITION, vec![[0.0; 3]; 3])
        .with_inserted_attribute(Mesh::ATTRIBUTE_NORMAL, vec![[0.0, 0.0, 1.0]; 3])
        .with_inserted_attribute(Mesh::ATTRIBUTE_COLOR, vec![[0.0, 0.0, 0.0, 0.0]; 3])
        .with_inserted_indices(Indices::U32(vec![0, 1, 2])),
    );
    let material = materials.add(FilledGizmoMaterial {
        color: LinearRgba::WHITE,
    });
    commands.spawn((
        Mesh3d(mesh.clone()),
        MeshMaterial3d(material),
        NoFrustumCulling,
    ));
    buffer.mesh = Some(mesh);
}

fn begin_fixed_filled_gizmos(mut buffer: ResMut<FilledGizmoBuffer>) {
    buffer.pending = FilledGizmoGeometry::default();
    buffer.fixed_commit = false;
}

fn commit_fixed_filled_gizmos(mut buffer: ResMut<FilledGizmoBuffer>) {
    let pending = std::mem::take(&mut buffer.pending);
    buffer.pending = std::mem::replace(&mut buffer.retained, pending);
    buffer.fixed_commit = true;
}

fn upload_filled_gizmo_mesh(
    mut meshes: ResMut<Assets<Mesh>>,
    mut buffer: ResMut<FilledGizmoBuffer>,
) {
    // No fixed tick and no render-frame requests means the existing mesh is
    // already the correct image. In particular, do not replace it with the
    // empty placeholder between fixed updates.
    if !buffer.fixed_commit && buffer.pending.positions.is_empty() {
        return;
    }
    let Some(handle) = buffer.mesh.clone() else {
        return;
    };

    let (positions, normals, colors, indices) = if buffer.fixed_commit {
        // Retained geometry must remain in the CPU-side buffer after upload;
        // otherwise the next render frame would have nothing to draw.
        (
            buffer.retained.positions.clone(),
            buffer.retained.normals.clone(),
            buffer.retained.colors.clone(),
            buffer.retained.indices.clone(),
        )
    } else {
        (
            std::mem::take(&mut buffer.pending.positions),
            std::mem::take(&mut buffer.pending.normals),
            std::mem::take(&mut buffer.pending.colors),
            std::mem::take(&mut buffer.pending.indices),
        )
    };

    if let Some(mut mesh) = meshes.get_mut(&handle) {
        if positions.is_empty() {
            mesh.insert_attribute(Mesh::ATTRIBUTE_POSITION, vec![[0.0; 3]; 3]);
            mesh.insert_attribute(Mesh::ATTRIBUTE_NORMAL, vec![[0.0, 0.0, 1.0]; 3]);
            mesh.insert_attribute(Mesh::ATTRIBUTE_COLOR, vec![[0.0, 0.0, 0.0, 0.0]; 3]);
            mesh.insert_indices(Indices::U32(vec![0, 1, 2]));
        } else {
            mesh.insert_attribute(Mesh::ATTRIBUTE_POSITION, positions);
            mesh.insert_attribute(Mesh::ATTRIBUTE_NORMAL, normals);
            mesh.insert_attribute(Mesh::ATTRIBUTE_COLOR, colors);
            mesh.insert_indices(Indices::U32(indices));
        }
    }
    buffer.pending = FilledGizmoGeometry::default();
    buffer.fixed_commit = false;
}