// 3D Primitive geometry catalog — meshlet-compatible vertex data for
// basic shapes used in 3D scene editing (cube, sphere, cylinder, cone,
// plane, torus, capsule). Each primitive generates indexed triangle mesh
// data compatible with both standard vertex pipelines and meshlet-based
// mesh shader pipelines.
//
// CODESPEC: Pre-allocated, no per-frame allocation. Geometry generated
// once at startup and uploaded to GPU buffers.

use crate::mesh_shader::{build_meshlets, Meshlet};

/// Vertex for 3D primitives: position + normal + color.
#[repr(C)]
#[derive(Copy, Clone, Debug, bytemuck::Pod, bytemuck::Zeroable)]
pub struct PrimitiveVertex {
    pub position: [f32; 3],
    pub normal: [f32; 3],
    pub color: [f32; 4],
}

/// PBR vertex for textured mesh rendering with tangent space.
/// 64 bytes stride. Used with pbr.wgsl.
#[repr(C)]
#[derive(Copy, Clone, Debug, bytemuck::Pod, bytemuck::Zeroable)]
pub struct PbrVertex {
    pub position: [f32; 3],
    pub normal: [f32; 3],
    pub tangent: [f32; 4], // .w = bitangent sign
    pub tex_coord: [f32; 2],
    pub color: [f32; 4],
}

/// Vertex buffer layout for PbrVertex, matching pbr.wgsl attributes.
pub const PBR_VERTEX_LAYOUT: wgpu::VertexBufferLayout<'static> = wgpu::VertexBufferLayout {
    array_stride: std::mem::size_of::<PbrVertex>() as u64,
    step_mode: wgpu::VertexStepMode::Vertex,
    attributes: &wgpu::vertex_attr_array![
        0 => Float32x3, // position
        1 => Float32x3, // normal
        2 => Float32x4, // tangent
        3 => Float32x2, // tex_coord
        4 => Float32x4, // color
    ],
};

/// A generated mesh with vertices, indices, and pre-built meshlets.
#[derive(Debug, Clone)]
pub struct PrimitiveMesh {
    pub name: &'static str,
    pub vertices: Vec<PrimitiveVertex>,
    pub indices: Vec<u32>,
    pub meshlets: Vec<Meshlet>,
}

/// Primitive shape type.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum PrimitiveShape {
    Cube,
    Sphere,
    Cylinder,
    Cone,
    Plane,
    Torus,
    Capsule,
    Pyramid,
    Wedge,
}

impl PrimitiveShape {
    pub const ALL: &'static [PrimitiveShape] = &[
        Self::Cube,
        Self::Sphere,
        Self::Cylinder,
        Self::Cone,
        Self::Plane,
        Self::Torus,
        Self::Capsule,
        Self::Pyramid,
        Self::Wedge,
    ];

    pub fn name(self) -> &'static str {
        match self {
            Self::Cube => "Cube",
            Self::Sphere => "Sphere",
            Self::Cylinder => "Cylinder",
            Self::Cone => "Cone",
            Self::Plane => "Plane",
            Self::Torus => "Torus",
            Self::Capsule => "Capsule",
            Self::Pyramid => "Pyramid",
            Self::Wedge => "Wedge",
        }
    }

    pub fn icon(self) -> &'static str {
        match self {
            Self::Cube => "[#]",
            Self::Sphere => "(O)",
            Self::Cylinder => "[|]",
            Self::Cone => "[^]",
            Self::Plane => "[_]",
            Self::Torus => "(@)",
            Self::Capsule => "[D]",
            Self::Pyramid => "[A]",
            Self::Wedge => "[/]",
        }
    }

    /// Generate the mesh for this primitive with a given color.
    pub fn generate(self, color: [f32; 4]) -> PrimitiveMesh {
        match self {
            Self::Cube => generate_cube(color),
            Self::Sphere => generate_sphere(16, 12, color),
            Self::Cylinder => generate_cylinder(16, color),
            Self::Cone => generate_cone(16, color),
            Self::Plane => generate_plane(color),
            Self::Torus => generate_torus(16, 8, color),
            Self::Capsule => generate_capsule(12, 8, color),
            Self::Pyramid => generate_pyramid(color),
            Self::Wedge => generate_wedge(color),
        }
    }

    /// Generate a mesh with custom tessellation parameters.
    /// `slices` controls circumferential segments, `stacks` controls latitudinal segments.
    /// For shapes that don't parameterize tessellation, this falls back to `generate()`.
    pub fn generate_with_tessellation(self, color: [f32; 4], slices: u32, stacks: u32) -> PrimitiveMesh {
        match self {
            Self::Sphere => generate_sphere(slices, stacks, color),
            Self::Cylinder => generate_cylinder(slices, color),
            Self::Cone => generate_cone(slices, color),
            Self::Torus => generate_torus(slices, stacks, color),
            Self::Capsule => generate_capsule(slices, stacks, color),
            _ => self.generate(color),
        }
    }
}

/// Default color for editor primitives (neutral gray).
pub const DEFAULT_PRIMITIVE_COLOR: [f32; 4] = [0.7, 0.7, 0.7, 1.0];

// Max meshlet limits per NVIDIA Turing mesh shader recommendation.
// Each meshlet represents a variable number of vertices and primitives. There are no
// restrictions regarding the connectivity of these primitives. However, they must stay
// below a maximum amount, specified within the shader code.
//
// We recommend using up to 64 vertices and 126 primitives. The '6' in 126 is not a
// typo. First-generation hardware allocates primitive indices in 128-byte granularity
// and needs to reserve 4 bytes for the primitive count. Therefore 3 × 126 + 4
// maximizes the fit into a 3 × 128 = 384 bytes block. Going beyond 126 triangles
// would allocate the next 128 bytes. 84 and 40 are other maxima that work well.
//
// Source: https://developer.nvidia.com/blog/introduction-turing-mesh-shaders/
const MESHLET_MAX_VERTS: u32 = 64;
const MESHLET_MAX_TRIS: u32 = 126;

fn finalize_mesh(name: &'static str, vertices: Vec<PrimitiveVertex>, indices: Vec<u32>) -> PrimitiveMesh {
    let meshlets = build_meshlets(&indices, MESHLET_MAX_VERTS, MESHLET_MAX_TRIS);
    PrimitiveMesh {
        name,
        vertices,
        indices,
        meshlets,
    }
}

fn push_vert(verts: &mut Vec<PrimitiveVertex>, pos: [f32; 3], normal: [f32; 3], color: [f32; 4]) -> u32 {
    let idx = verts.len() as u32;
    verts.push(PrimitiveVertex {
        position: pos,
        normal,
        color,
    });
    idx
}

// ── Cube ──────────────────────────────────────────────────────────────

fn generate_cube(color: [f32; 4]) -> PrimitiveMesh {
    let mut verts = Vec::with_capacity(24);
    let mut idxs = Vec::with_capacity(36);
    let s = 0.5_f32;

    // 6 faces × 4 verts each (unique normals per face)
    let faces: &[([f32; 3], [[f32; 3]; 4])] = &[
        ([0.0, 0.0, 1.0], [[-s, -s, s], [s, -s, s], [s, s, s], [-s, s, s]]), // front
        ([0.0, 0.0, -1.0], [[s, -s, -s], [-s, -s, -s], [-s, s, -s], [s, s, -s]]), // back
        ([0.0, 1.0, 0.0], [[-s, s, s], [s, s, s], [s, s, -s], [-s, s, -s]]), // top
        ([0.0, -1.0, 0.0], [[-s, -s, -s], [s, -s, -s], [s, -s, s], [-s, -s, s]]), // bottom
        ([1.0, 0.0, 0.0], [[s, -s, s], [s, -s, -s], [s, s, -s], [s, s, s]]), // right
        ([-1.0, 0.0, 0.0], [[-s, -s, -s], [-s, -s, s], [-s, s, s], [-s, s, -s]]), // left
    ];

    for &(normal, corners) in faces {
        let base = verts.len() as u32;
        for pos in corners {
            verts.push(PrimitiveVertex {
                position: pos,
                normal,
                color,
            });
        }
        idxs.extend_from_slice(&[base, base + 1, base + 2, base, base + 2, base + 3]);
    }

    finalize_mesh("Cube", verts, idxs)
}

// ── Sphere ────────────────────────────────────────────────────────────

fn generate_sphere(slices: u32, stacks: u32, color: [f32; 4]) -> PrimitiveMesh {
    let mut verts = Vec::new();
    let mut idxs = Vec::new();
    let r = 0.5_f32;

    for j in 0..=stacks {
        let phi = std::f32::consts::PI * j as f32 / stacks as f32;
        for i in 0..=slices {
            let theta = 2.0 * std::f32::consts::PI * i as f32 / slices as f32;
            let x = phi.sin() * theta.cos();
            let y = phi.cos();
            let z = phi.sin() * theta.sin();
            push_vert(&mut verts, [x * r, y * r, z * r], [x, y, z], color);
        }
    }

    for j in 0..stacks {
        for i in 0..slices {
            let a = j * (slices + 1) + i;
            let b = a + slices + 1;
            idxs.extend_from_slice(&[a, b, a + 1, a + 1, b, b + 1]);
        }
    }

    finalize_mesh("Sphere", verts, idxs)
}

// ── Cylinder ──────────────────────────────────────────────────────────

fn generate_cylinder(segments: u32, color: [f32; 4]) -> PrimitiveMesh {
    let mut verts = Vec::new();
    let mut idxs = Vec::new();
    let r = 0.5_f32;
    let h = 0.5_f32;

    // Side
    for i in 0..=segments {
        let theta = 2.0 * std::f32::consts::PI * i as f32 / segments as f32;
        let nx = theta.cos();
        let nz = theta.sin();
        push_vert(&mut verts, [nx * r, h, nz * r], [nx, 0.0, nz], color);
        push_vert(&mut verts, [nx * r, -h, nz * r], [nx, 0.0, nz], color);
    }
    for i in 0..segments {
        let a = i * 2;
        let b = a + 1;
        let c = a + 2;
        let d = a + 3;
        idxs.extend_from_slice(&[a, b, c, c, b, d]);
    }

    // Top cap
    let top_center = push_vert(&mut verts, [0.0, h, 0.0], [0.0, 1.0, 0.0], color);
    for i in 0..segments {
        let theta = 2.0 * std::f32::consts::PI * i as f32 / segments as f32;
        push_vert(
            &mut verts,
            [theta.cos() * r, h, theta.sin() * r],
            [0.0, 1.0, 0.0],
            color,
        );
    }
    for i in 0..segments {
        let a = top_center + 1 + i;
        let b = top_center + 1 + (i + 1) % segments;
        idxs.extend_from_slice(&[top_center, a, b]);
    }

    // Bottom cap
    let bot_center = push_vert(&mut verts, [0.0, -h, 0.0], [0.0, -1.0, 0.0], color);
    for i in 0..segments {
        let theta = 2.0 * std::f32::consts::PI * i as f32 / segments as f32;
        push_vert(
            &mut verts,
            [theta.cos() * r, -h, theta.sin() * r],
            [0.0, -1.0, 0.0],
            color,
        );
    }
    for i in 0..segments {
        let a = bot_center + 1 + i;
        let b = bot_center + 1 + (i + 1) % segments;
        idxs.extend_from_slice(&[bot_center, b, a]);
    }

    finalize_mesh("Cylinder", verts, idxs)
}

// ── Cone ──────────────────────────────────────────────────────────────

fn generate_cone(segments: u32, color: [f32; 4]) -> PrimitiveMesh {
    let mut verts = Vec::new();
    let mut idxs = Vec::new();
    let r = 0.5_f32;
    let h = 1.0_f32;

    // Apex
    let apex = push_vert(&mut verts, [0.0, h * 0.5, 0.0], [0.0, 1.0, 0.0], color);

    // Base ring
    for i in 0..=segments {
        let theta = 2.0 * std::f32::consts::PI * i as f32 / segments as f32;
        let nx = theta.cos();
        let nz = theta.sin();
        let slope = r / h;
        let ny = (1.0 / (1.0 + slope * slope)).sqrt();
        push_vert(
            &mut verts,
            [nx * r, -h * 0.5, nz * r],
            [nx * ny, slope * ny, nz * ny],
            color,
        );
    }

    for i in 0..segments {
        idxs.extend_from_slice(&[apex, apex + 1 + i, apex + 2 + i]);
    }

    // Base cap
    let base_center = push_vert(&mut verts, [0.0, -h * 0.5, 0.0], [0.0, -1.0, 0.0], color);
    for i in 0..segments {
        let theta = 2.0 * std::f32::consts::PI * i as f32 / segments as f32;
        push_vert(
            &mut verts,
            [theta.cos() * r, -h * 0.5, theta.sin() * r],
            [0.0, -1.0, 0.0],
            color,
        );
    }
    for i in 0..segments {
        let a = base_center + 1 + i;
        let b = base_center + 1 + (i + 1) % segments;
        idxs.extend_from_slice(&[base_center, b, a]);
    }

    finalize_mesh("Cone", verts, idxs)
}

// ── Plane ─────────────────────────────────────────────────────────────

fn generate_plane(color: [f32; 4]) -> PrimitiveMesh {
    let s = 0.5_f32;
    let verts = vec![
        PrimitiveVertex {
            position: [-s, 0.0, -s],
            normal: [0.0, 1.0, 0.0],
            color,
        },
        PrimitiveVertex {
            position: [s, 0.0, -s],
            normal: [0.0, 1.0, 0.0],
            color,
        },
        PrimitiveVertex {
            position: [s, 0.0, s],
            normal: [0.0, 1.0, 0.0],
            color,
        },
        PrimitiveVertex {
            position: [-s, 0.0, s],
            normal: [0.0, 1.0, 0.0],
            color,
        },
    ];
    let idxs = vec![0, 1, 2, 0, 2, 3];
    finalize_mesh("Plane", verts, idxs)
}

// ── Torus ─────────────────────────────────────────────────────────────

fn generate_torus(ring_segments: u32, tube_segments: u32, color: [f32; 4]) -> PrimitiveMesh {
    let mut verts = Vec::new();
    let mut idxs = Vec::new();
    let ring_r = 0.35_f32;
    let tube_r = 0.15_f32;

    for i in 0..=ring_segments {
        let theta = 2.0 * std::f32::consts::PI * i as f32 / ring_segments as f32;
        let ct = theta.cos();
        let st = theta.sin();

        for j in 0..=tube_segments {
            let phi = 2.0 * std::f32::consts::PI * j as f32 / tube_segments as f32;
            let cp = phi.cos();
            let sp = phi.sin();

            let x = (ring_r + tube_r * cp) * ct;
            let y = tube_r * sp;
            let z = (ring_r + tube_r * cp) * st;

            let nx = cp * ct;
            let ny = sp;
            let nz = cp * st;

            push_vert(&mut verts, [x, y, z], [nx, ny, nz], color);
        }
    }

    for i in 0..ring_segments {
        for j in 0..tube_segments {
            let a = i * (tube_segments + 1) + j;
            let b = a + tube_segments + 1;
            idxs.extend_from_slice(&[a, b, a + 1, a + 1, b, b + 1]);
        }
    }

    finalize_mesh("Torus", verts, idxs)
}

// ── Capsule ───────────────────────────────────────────────────────────

fn generate_capsule(segments: u32, rings: u32, color: [f32; 4]) -> PrimitiveMesh {
    // Capsule = cylinder body + two hemisphere caps.
    let mut verts = Vec::new();
    let mut idxs = Vec::new();
    let r = 0.25_f32;
    let h = 0.5_f32; // half-height of cylinder part

    let half_rings = rings / 2;

    // Top hemisphere
    for j in 0..=half_rings {
        let phi = std::f32::consts::FRAC_PI_2 * j as f32 / half_rings as f32;
        for i in 0..=segments {
            let theta = 2.0 * std::f32::consts::PI * i as f32 / segments as f32;
            let x = phi.cos() * theta.cos();
            let y = phi.sin();
            let z = phi.cos() * theta.sin();
            push_vert(&mut verts, [x * r, y * r + h, z * r], [x, y, z], color);
        }
    }

    // Bottom hemisphere
    for j in 0..=half_rings {
        let phi = std::f32::consts::FRAC_PI_2 + std::f32::consts::FRAC_PI_2 * j as f32 / half_rings as f32;
        for i in 0..=segments {
            let theta = 2.0 * std::f32::consts::PI * i as f32 / segments as f32;
            let x = phi.cos() * theta.cos();
            let y = phi.sin();
            let z = phi.cos() * theta.sin();
            push_vert(&mut verts, [x * r, y * r - h, z * r], [x, y, z], color);
        }
    }

    let total_rows = (half_rings + 1) * 2;
    for j in 0..total_rows - 1 {
        for i in 0..segments {
            let a = j * (segments + 1) + i;
            let b = a + segments + 1;
            idxs.extend_from_slice(&[a, b, a + 1, a + 1, b, b + 1]);
        }
    }

    finalize_mesh("Capsule", verts, idxs)
}

// ── Pyramid ───────────────────────────────────────────────────────────

fn generate_pyramid(color: [f32; 4]) -> PrimitiveMesh {
    let mut verts = Vec::new();
    let mut idxs = Vec::new();
    let s = 0.5_f32;
    let h = 0.7_f32;

    // Apex
    let apex = push_vert(&mut verts, [0.0, h, 0.0], [0.0, 1.0, 0.0], color);

    // Base corners
    let corners = [[-s, 0.0, -s], [s, 0.0, -s], [s, 0.0, s], [-s, 0.0, s]];

    // 4 side faces
    for i in 0..4 {
        let j = (i + 1) % 4;
        let a = push_vert(&mut verts, corners[i], [0.0, 0.5, 0.0], color);
        let b = push_vert(&mut verts, corners[j], [0.0, 0.5, 0.0], color);
        idxs.extend_from_slice(&[apex, a, b]);
    }

    // Base (2 triangles)
    let b0 = push_vert(&mut verts, corners[0], [0.0, -1.0, 0.0], color);
    let b1 = push_vert(&mut verts, corners[1], [0.0, -1.0, 0.0], color);
    let b2 = push_vert(&mut verts, corners[2], [0.0, -1.0, 0.0], color);
    let b3 = push_vert(&mut verts, corners[3], [0.0, -1.0, 0.0], color);
    idxs.extend_from_slice(&[b0, b2, b1, b0, b3, b2]);

    finalize_mesh("Pyramid", verts, idxs)
}

// ── Wedge ─────────────────────────────────────────────────────────────

fn generate_wedge(color: [f32; 4]) -> PrimitiveMesh {
    let mut verts = Vec::new();
    let mut idxs = Vec::new();
    let s = 0.5_f32;

    // 6 unique positions: bottom quad + top edge
    let positions = [
        [-s, 0.0, -s],
        [s, 0.0, -s],
        [s, 0.0, s],
        [-s, 0.0, s], // bottom
        [-s, s, -s],
        [s, s, -s], // top edge (back)
    ];

    // Bottom face
    let b0 = push_vert(&mut verts, positions[0], [0.0, -1.0, 0.0], color);
    let b1 = push_vert(&mut verts, positions[1], [0.0, -1.0, 0.0], color);
    let b2 = push_vert(&mut verts, positions[2], [0.0, -1.0, 0.0], color);
    let b3 = push_vert(&mut verts, positions[3], [0.0, -1.0, 0.0], color);
    idxs.extend_from_slice(&[b0, b2, b1, b0, b3, b2]);

    // Slope face (front)
    let s0 = push_vert(&mut verts, positions[3], [0.0, 0.707, 0.707], color);
    let s1 = push_vert(&mut verts, positions[2], [0.0, 0.707, 0.707], color);
    let s2 = push_vert(&mut verts, positions[5], [0.0, 0.707, 0.707], color);
    let s3 = push_vert(&mut verts, positions[4], [0.0, 0.707, 0.707], color);
    idxs.extend_from_slice(&[s0, s1, s2, s0, s2, s3]);

    // Back face
    let k0 = push_vert(&mut verts, positions[4], [0.0, 0.0, -1.0], color);
    let k1 = push_vert(&mut verts, positions[5], [0.0, 0.0, -1.0], color);
    let k2 = push_vert(&mut verts, positions[1], [0.0, 0.0, -1.0], color);
    let k3 = push_vert(&mut verts, positions[0], [0.0, 0.0, -1.0], color);
    idxs.extend_from_slice(&[k0, k1, k2, k0, k2, k3]);

    // Left face (triangle)
    let l0 = push_vert(&mut verts, positions[0], [-1.0, 0.0, 0.0], color);
    let l1 = push_vert(&mut verts, positions[3], [-1.0, 0.0, 0.0], color);
    let l2 = push_vert(&mut verts, positions[4], [-1.0, 0.0, 0.0], color);
    idxs.extend_from_slice(&[l0, l1, l2]);

    // Right face (triangle)
    let r0 = push_vert(&mut verts, positions[2], [1.0, 0.0, 0.0], color);
    let r1 = push_vert(&mut verts, positions[1], [1.0, 0.0, 0.0], color);
    let r2 = push_vert(&mut verts, positions[5], [1.0, 0.0, 0.0], color);
    idxs.extend_from_slice(&[r0, r1, r2]);

    finalize_mesh("Wedge", verts, idxs)
}

/// Vertex for the transform gizmo (position + color, no normal needed).
#[repr(C)]
#[derive(Debug, Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
pub struct GizmoVertex {
    pub position: [f32; 3],
    pub color: [f32; 4],
}

/// Generate translate gizmo mesh data — 3 axis arrows (cylinders + cones).
/// Returns (vertices, indices) for each axis: X (red), Y (green), Z (blue).
pub fn generate_gizmo_arrows() -> (Vec<GizmoVertex>, Vec<u32>) {
    let mut vertices = Vec::new();
    let mut indices = Vec::new();

    let axis_length = 1.0f32;
    let shaft_radius = 0.02f32;
    let cone_radius = 0.06f32;
    let cone_length = 0.2f32;
    let segments = 8u32;

    let colors = [
        [1.0, 0.2, 0.2, 1.0], // X = red
        [0.2, 1.0, 0.2, 1.0], // Y = green
        [0.3, 0.3, 1.0, 1.0], // Z = blue
    ];

    let directions = [
        [1.0, 0.0, 0.0f32], // X
        [0.0, 1.0, 0.0],    // Y
        [0.0, 0.0, 1.0],    // Z
    ];

    for axis in 0..3 {
        let color = colors[axis];
        let dir = directions[axis];
        let base_idx = vertices.len() as u32;

        // Create rotation basis for this axis
        let (up, right) = match axis {
            0 => ([0.0, 1.0, 0.0f32], [0.0, 0.0, 1.0f32]), // X axis
            1 => ([0.0, 0.0, 1.0], [1.0, 0.0, 0.0]),       // Y axis
            _ => ([1.0, 0.0, 0.0], [0.0, 1.0, 0.0]),       // Z axis
        };

        // Cylinder shaft: origin to (axis_length - cone_length)
        let shaft_end = axis_length - cone_length;

        // Bottom ring
        for i in 0..segments {
            let angle = (i as f32 / segments as f32) * std::f32::consts::TAU;
            let (sin_a, cos_a) = angle.sin_cos();
            let px = up[0] * cos_a * shaft_radius + right[0] * sin_a * shaft_radius;
            let py = up[1] * cos_a * shaft_radius + right[1] * sin_a * shaft_radius;
            let pz = up[2] * cos_a * shaft_radius + right[2] * sin_a * shaft_radius;
            vertices.push(GizmoVertex {
                position: [px, py, pz],
                color,
            });
        }
        // Top ring
        for i in 0..segments {
            let angle = (i as f32 / segments as f32) * std::f32::consts::TAU;
            let (sin_a, cos_a) = angle.sin_cos();
            let px = dir[0] * shaft_end + up[0] * cos_a * shaft_radius + right[0] * sin_a * shaft_radius;
            let py = dir[1] * shaft_end + up[1] * cos_a * shaft_radius + right[1] * sin_a * shaft_radius;
            let pz = dir[2] * shaft_end + up[2] * cos_a * shaft_radius + right[2] * sin_a * shaft_radius;
            vertices.push(GizmoVertex {
                position: [px, py, pz],
                color,
            });
        }

        // Shaft indices (triangle strip as quads)
        for i in 0..segments {
            let next = (i + 1) % segments;
            let b0 = base_idx + i;
            let b1 = base_idx + next;
            let t0 = base_idx + segments + i;
            let t1 = base_idx + segments + next;
            indices.extend_from_slice(&[b0, t0, b1, b1, t0, t1]);
        }

        // Cone: cone_base to tip
        let cone_base_idx = vertices.len() as u32;
        let cone_base_pos = [dir[0] * shaft_end, dir[1] * shaft_end, dir[2] * shaft_end];
        let tip_pos = [dir[0] * axis_length, dir[1] * axis_length, dir[2] * axis_length];

        // Cone base ring
        for i in 0..segments {
            let angle = (i as f32 / segments as f32) * std::f32::consts::TAU;
            let (sin_a, cos_a) = angle.sin_cos();
            let px = cone_base_pos[0] + up[0] * cos_a * cone_radius + right[0] * sin_a * cone_radius;
            let py = cone_base_pos[1] + up[1] * cos_a * cone_radius + right[1] * sin_a * cone_radius;
            let pz = cone_base_pos[2] + up[2] * cos_a * cone_radius + right[2] * sin_a * cone_radius;
            vertices.push(GizmoVertex {
                position: [px, py, pz],
                color,
            });
        }

        // Cone tip
        let tip_idx = vertices.len() as u32;
        vertices.push(GizmoVertex {
            position: tip_pos,
            color,
        });

        // Cone side indices
        for i in 0..segments {
            let next = (i + 1) % segments;
            indices.extend_from_slice(&[cone_base_idx + i, tip_idx, cone_base_idx + next]);
        }

        // Cone base cap (center + fan)
        let cap_center_idx = vertices.len() as u32;
        vertices.push(GizmoVertex {
            position: cone_base_pos,
            color,
        });
        for i in 0..segments {
            let next = (i + 1) % segments;
            indices.extend_from_slice(&[cap_center_idx, cone_base_idx + next, cone_base_idx + i]);
        }
    }

    (vertices, indices)
}

/// Catalog of all built-in primitive shapes with pre-generated meshes.
pub struct PrimitiveCatalog {
    pub meshes: Vec<PrimitiveMesh>,
}

impl PrimitiveCatalog {
    /// Generate all primitive shapes with default color.
    pub fn with_defaults() -> Self {
        let meshes = PrimitiveShape::ALL
            .iter()
            .map(|shape| shape.generate(DEFAULT_PRIMITIVE_COLOR))
            .collect();
        Self { meshes }
    }

    /// Find a mesh by name.
    pub fn find(&self, name: &str) -> Option<&PrimitiveMesh> {
        self.meshes.iter().find(|m| m.name == name)
    }

    /// Total vertex count across all primitives.
    pub fn total_vertices(&self) -> usize {
        self.meshes.iter().map(|m| m.vertices.len()).sum()
    }

    /// Total meshlet count across all primitives.
    pub fn total_meshlets(&self) -> usize {
        self.meshes.iter().map(|m| m.meshlets.len()).sum()
    }
}

// ── PBR mesh generation ──────────────────────────────────────────────

/// A generated mesh using PbrVertex (64 bytes) for the PBR pipeline.
#[derive(Debug, Clone)]
pub struct PbrMesh {
    pub name: &'static str,
    pub vertices: Vec<PbrVertex>,
    pub indices: Vec<u32>,
}

impl PrimitiveShape {
    /// Generate a PBR mesh with custom tessellation for LOD variants.
    pub fn generate_pbr_with_tessellation(self, color: [f32; 4], slices: u32, stacks: u32) -> PbrMesh {
        let base = self.generate_with_tessellation(color, slices, stacks);
        Self::base_to_pbr(self, base)
    }

    /// Generate a PBR-ready mesh with tangent vectors and UV coordinates.
    /// Tangents are computed from vertex positions and UVs (MikkTSpace-compatible).
    /// UVs use planar projection for flat faces and spherical/cylindrical for curved shapes.
    pub fn generate_pbr(self, color: [f32; 4]) -> PbrMesh {
        let base = self.generate(color);
        Self::base_to_pbr(self, base)
    }

    /// Convert a PrimitiveMesh to a PbrMesh with UVs and tangents.
    fn base_to_pbr(shape: PrimitiveShape, base: PrimitiveMesh) -> PbrMesh {
        let mut pbr_verts: Vec<PbrVertex> = base
            .vertices
            .iter()
            .map(|v| {
                // Default planar UV from XZ position (scaled to 0..1 range)
                let u = v.position[0] + 0.5;
                let w = v.position[2] + 0.5;
                PbrVertex {
                    position: v.position,
                    normal: v.normal,
                    tangent: [1.0, 0.0, 0.0, 1.0], // placeholder, computed below
                    tex_coord: [u, w],
                    color: v.color,
                }
            })
            .collect();

        // Compute proper UVs per shape type
        match shape {
            PrimitiveShape::Sphere | PrimitiveShape::Capsule => {
                for v in &mut pbr_verts {
                    let [x, y, z] = v.normal;
                    let u = 0.5 + z.atan2(x) / std::f32::consts::TAU;
                    let w = 0.5 - y.asin() / std::f32::consts::PI;
                    v.tex_coord = [u, w];
                }
            }
            PrimitiveShape::Cylinder | PrimitiveShape::Cone => {
                for v in &mut pbr_verts {
                    let [px, py, pz] = v.position;
                    if v.normal[1].abs() > 0.9 {
                        // Cap face: planar from XZ
                        v.tex_coord = [px + 0.5, pz + 0.5];
                    } else {
                        // Side: cylindrical
                        let u = 0.5 + pz.atan2(px) / std::f32::consts::TAU;
                        let w = py + 0.5;
                        v.tex_coord = [u, w];
                    }
                }
            }
            PrimitiveShape::Torus => {
                // Already has natural parameterization from ring/tube angles
                // Use normal-based spherical mapping as approximation
                for v in &mut pbr_verts {
                    let [x, y, z] = v.normal;
                    let u = 0.5 + z.atan2(x) / std::f32::consts::TAU;
                    let w = 0.5 + y * 0.5;
                    v.tex_coord = [u, w];
                }
            }
            _ => {
                // Cube, Plane, Pyramid, Wedge: per-face planar UVs
                // Process per-triangle to assign UVs based on dominant normal axis
                for tri in base.indices.chunks(3) {
                    if tri.len() < 3 {
                        continue;
                    }
                    let n = pbr_verts[tri[0] as usize].normal;
                    let ax = n[0].abs();
                    let ay = n[1].abs();
                    let az = n[2].abs();
                    for &idx in tri {
                        let v = &mut pbr_verts[idx as usize];
                        let [px, py, pz] = v.position;
                        if ay >= ax && ay >= az {
                            // Y-dominant (top/bottom): project XZ
                            v.tex_coord = [px + 0.5, pz + 0.5];
                        } else if ax >= az {
                            // X-dominant (left/right): project YZ
                            v.tex_coord = [pz + 0.5, py + 0.5];
                        } else {
                            // Z-dominant (front/back): project XY
                            v.tex_coord = [px + 0.5, py + 0.5];
                        }
                    }
                }
            }
        }

        // Compute tangent vectors from UV + normal (MikkTSpace approximation)
        compute_tangents(&mut pbr_verts, &base.indices);

        PbrMesh {
            name: base.name,
            vertices: pbr_verts,
            indices: base.indices,
        }
    }
}

/// Compute tangent vectors for PBR vertices using triangle edge/UV deltas.
/// Sets tangent.w to the bitangent sign (+1 or -1).
fn compute_tangents(vertices: &mut [PbrVertex], indices: &[u32]) {
    // Accumulate tangent contributions per vertex
    let mut tan_accum: Vec<[f32; 3]> = vec![[0.0; 3]; vertices.len()];
    let mut bitan_accum: Vec<[f32; 3]> = vec![[0.0; 3]; vertices.len()];

    for tri in indices.chunks(3) {
        if tri.len() < 3 {
            continue;
        }
        let i0 = tri[0] as usize;
        let i1 = tri[1] as usize;
        let i2 = tri[2] as usize;

        let p0 = vertices[i0].position;
        let p1 = vertices[i1].position;
        let p2 = vertices[i2].position;

        let uv0 = vertices[i0].tex_coord;
        let uv1 = vertices[i1].tex_coord;
        let uv2 = vertices[i2].tex_coord;

        let dp1 = [p1[0] - p0[0], p1[1] - p0[1], p1[2] - p0[2]];
        let dp2 = [p2[0] - p0[0], p2[1] - p0[1], p2[2] - p0[2]];

        let duv1 = [uv1[0] - uv0[0], uv1[1] - uv0[1]];
        let duv2 = [uv2[0] - uv0[0], uv2[1] - uv0[1]];

        let det = duv1[0] * duv2[1] - duv1[1] * duv2[0];
        let r = if det.abs() > 1e-8 { 1.0 / det } else { 0.0 };

        let t = [
            (duv2[1] * dp1[0] - duv1[1] * dp2[0]) * r,
            (duv2[1] * dp1[1] - duv1[1] * dp2[1]) * r,
            (duv2[1] * dp1[2] - duv1[1] * dp2[2]) * r,
        ];
        let b = [
            (duv1[0] * dp2[0] - duv2[0] * dp1[0]) * r,
            (duv1[0] * dp2[1] - duv2[0] * dp1[1]) * r,
            (duv1[0] * dp2[2] - duv2[0] * dp1[2]) * r,
        ];

        for &idx in &[i0, i1, i2] {
            tan_accum[idx][0] += t[0];
            tan_accum[idx][1] += t[1];
            tan_accum[idx][2] += t[2];
            bitan_accum[idx][0] += b[0];
            bitan_accum[idx][1] += b[1];
            bitan_accum[idx][2] += b[2];
        }
    }

    // Orthogonalize and store
    for (i, v) in vertices.iter_mut().enumerate() {
        let n = v.normal;
        let t = tan_accum[i];

        // Gram-Schmidt orthogonalize: t' = normalize(t - n * dot(n, t))
        let dot_nt = n[0] * t[0] + n[1] * t[1] + n[2] * t[2];
        let orth = [t[0] - n[0] * dot_nt, t[1] - n[1] * dot_nt, t[2] - n[2] * dot_nt];
        let len = (orth[0] * orth[0] + orth[1] * orth[1] + orth[2] * orth[2]).sqrt();

        let tangent = if len > 1e-8 {
            [orth[0] / len, orth[1] / len, orth[2] / len]
        } else {
            // Fallback tangent from normal
            if n[1].abs() > 0.9 {
                [1.0, 0.0, 0.0]
            } else {
                let cr = [0.0 - n[2] * 0.0, n[2] * 0.0 - 0.0, 0.0 * 0.0 - 0.0];
                let cl = (cr[0] * cr[0] + cr[1] * cr[1] + cr[2] * cr[2]).sqrt();
                if cl > 1e-8 {
                    [cr[0] / cl, cr[1] / cl, cr[2] / cl]
                } else {
                    [1.0, 0.0, 0.0]
                }
            }
        };

        // Bitangent sign: sign of dot(cross(n, t), b)
        let b = bitan_accum[i];
        let cross = [
            n[1] * tangent[2] - n[2] * tangent[1],
            n[2] * tangent[0] - n[0] * tangent[2],
            n[0] * tangent[1] - n[1] * tangent[0],
        ];
        let sign = if (cross[0] * b[0] + cross[1] * b[1] + cross[2] * b[2]) < 0.0 {
            -1.0
        } else {
            1.0
        };

        v.tangent = [tangent[0], tangent[1], tangent[2], sign];
    }
}

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

    #[test]
    fn all_shapes_generate() {
        for shape in PrimitiveShape::ALL {
            let mesh = shape.generate(DEFAULT_PRIMITIVE_COLOR);
            assert!(!mesh.vertices.is_empty(), "{} has no vertices", shape.name());
            assert!(!mesh.indices.is_empty(), "{} has no indices", shape.name());
            assert_eq!(mesh.indices.len() % 3, 0, "{} indices not divisible by 3", shape.name());
            assert!(!mesh.meshlets.is_empty(), "{} has no meshlets", shape.name());
        }
    }

    #[test]
    fn cube_has_24_verts_36_indices() {
        let mesh = PrimitiveShape::Cube.generate(DEFAULT_PRIMITIVE_COLOR);
        assert_eq!(mesh.vertices.len(), 24);
        assert_eq!(mesh.indices.len(), 36);
    }

    #[test]
    fn plane_has_4_verts() {
        let mesh = PrimitiveShape::Plane.generate(DEFAULT_PRIMITIVE_COLOR);
        assert_eq!(mesh.vertices.len(), 4);
        assert_eq!(mesh.indices.len(), 6);
    }

    #[test]
    fn catalog_has_all_shapes() {
        let catalog = PrimitiveCatalog::with_defaults();
        assert_eq!(catalog.meshes.len(), PrimitiveShape::ALL.len());
    }

    #[test]
    fn catalog_find_works() {
        let catalog = PrimitiveCatalog::with_defaults();
        assert!(catalog.find("Cube").is_some());
        assert!(catalog.find("Nonexistent").is_none());
    }

    #[test]
    fn vertex_layout_size() {
        assert_eq!(std::mem::size_of::<PrimitiveVertex>(), 40);
    }

    #[test]
    fn pbr_vertex_layout_size() {
        assert_eq!(std::mem::size_of::<PbrVertex>(), 64);
    }

    #[test]
    fn shape_names_unique() {
        let names: Vec<_> = PrimitiveShape::ALL.iter().map(|s| s.name()).collect();
        let mut deduped = names.clone();
        deduped.sort();
        deduped.dedup();
        assert_eq!(names.len(), deduped.len());
    }

    #[test]
    fn meshlets_respect_limits() {
        for shape in PrimitiveShape::ALL {
            let mesh = shape.generate(DEFAULT_PRIMITIVE_COLOR);
            for meshlet in &mesh.meshlets {
                assert!(
                    meshlet.vertex_count <= MESHLET_MAX_VERTS,
                    "{}: meshlet has {} verts",
                    shape.name(),
                    meshlet.vertex_count
                );
                assert!(
                    meshlet.triangle_count <= MESHLET_MAX_TRIS,
                    "{}: meshlet has {} tris",
                    shape.name(),
                    meshlet.triangle_count
                );
            }
        }
    }

    #[test]
    fn all_shapes_generate_pbr() {
        for shape in PrimitiveShape::ALL {
            let mesh = shape.generate_pbr(DEFAULT_PRIMITIVE_COLOR);
            assert!(!mesh.vertices.is_empty(), "{} PBR has no vertices", shape.name());
            assert!(!mesh.indices.is_empty(), "{} PBR has no indices", shape.name());
            assert_eq!(
                mesh.indices.len() % 3,
                0,
                "{} PBR indices not divisible by 3",
                shape.name()
            );
        }
    }

    #[test]
    fn pbr_tangents_are_normalized() {
        for shape in PrimitiveShape::ALL {
            let mesh = shape.generate_pbr(DEFAULT_PRIMITIVE_COLOR);
            for v in &mesh.vertices {
                let t = v.tangent;
                let len = (t[0] * t[0] + t[1] * t[1] + t[2] * t[2]).sqrt();
                assert!(
                    (len - 1.0).abs() < 0.01,
                    "{}: tangent not normalized (len={})",
                    shape.name(),
                    len
                );
                assert!(
                    t[3] == 1.0 || t[3] == -1.0,
                    "{}: tangent.w must be +1 or -1, got {}",
                    shape.name(),
                    t[3]
                );
            }
        }
    }

    #[test]
    fn pbr_uvs_in_valid_range() {
        for shape in PrimitiveShape::ALL {
            let mesh = shape.generate_pbr(DEFAULT_PRIMITIVE_COLOR);
            for v in &mesh.vertices {
                assert!(
                    v.tex_coord[0].is_finite() && v.tex_coord[1].is_finite(),
                    "{}: UV contains NaN/Inf",
                    shape.name()
                );
            }
        }
    }

    #[test]
    fn pbr_mesh_vertex_count_matches_primitive() {
        for shape in PrimitiveShape::ALL {
            let prim = shape.generate(DEFAULT_PRIMITIVE_COLOR);
            let pbr = shape.generate_pbr(DEFAULT_PRIMITIVE_COLOR);
            assert_eq!(
                prim.vertices.len(),
                pbr.vertices.len(),
                "{}: PBR vertex count should match primitive",
                shape.name()
            );
            assert_eq!(
                prim.indices.len(),
                pbr.indices.len(),
                "{}: PBR index count should match primitive",
                shape.name()
            );
        }
    }
}
