"""Export mesh data (vertices + triangles) to glTF binary (.glb) format."""

from __future__ import annotations

import numpy as np
import pygltflib


def _compute_normals(vertices: np.ndarray, triangles: np.ndarray) -> np.ndarray:
    """Compute smooth per-vertex normals by averaging face normals."""
    normals = np.zeros_like(vertices, dtype=np.float32)

    v0 = vertices[triangles[:, 0]]
    v1 = vertices[triangles[:, 1]]
    v2 = vertices[triangles[:, 2]]

    face_normals = np.cross(v1 - v0, v2 - v0)

    # Accumulate face normals to each vertex
    for i in range(3):
        np.add.at(normals, triangles[:, i], face_normals)

    # Normalize
    lengths = np.linalg.norm(normals, axis=1, keepdims=True)
    lengths = np.maximum(lengths, 1e-10)
    normals /= lengths

    return normals.astype(np.float32)


def mesh_to_glb(vertices: np.ndarray, triangles: np.ndarray) -> bytes:
    """Convert vertices and triangles to a binary glTF (.glb) buffer.

    Includes computed vertex normals for proper lighting/shading.

    Args:
        vertices: Mesh vertices with shape (N, 3), float64.
        triangles: Triangle indices with shape (M, 3), int32.

    Returns:
        Bytes of the .glb file.
    """
    vertices_f32 = vertices.astype(np.float32)
    triangles_u32 = triangles.astype(np.uint32)
    normals_f32 = _compute_normals(vertices_f32, triangles_u32)

    # Byte buffers
    vert_bytes = vertices_f32.tobytes()
    norm_bytes = normals_f32.tobytes()
    tri_bytes = triangles_u32.tobytes()

    # Align each buffer to 4 bytes
    def pad4(b: bytes) -> bytes:
        pad = (4 - len(b) % 4) % 4
        return b + b"\x00" * pad

    vert_padded = pad4(vert_bytes)
    norm_padded = pad4(norm_bytes)

    blob = vert_padded + norm_padded + tri_bytes
    total_bytes = len(blob)

    norm_offset = len(vert_padded)
    tri_offset = norm_offset + len(norm_padded)

    v_min = vertices_f32.min(axis=0).tolist()
    v_max = vertices_f32.max(axis=0).tolist()

    gltf = pygltflib.GLTF2(
        scene=0,
        scenes=[pygltflib.Scene(nodes=[0])],
        nodes=[pygltflib.Node(mesh=0)],
        meshes=[
            pygltflib.Mesh(
                primitives=[
                    pygltflib.Primitive(
                        attributes=pygltflib.Attributes(POSITION=0, NORMAL=1),
                        indices=2,
                    )
                ]
            )
        ],
        accessors=[
            # Accessor 0: vertex positions
            pygltflib.Accessor(
                bufferView=0,
                componentType=pygltflib.FLOAT,
                count=len(vertices_f32),
                type=pygltflib.VEC3,
                max=v_max,
                min=v_min,
            ),
            # Accessor 1: vertex normals
            pygltflib.Accessor(
                bufferView=1,
                componentType=pygltflib.FLOAT,
                count=len(normals_f32),
                type=pygltflib.VEC3,
            ),
            # Accessor 2: triangle indices
            pygltflib.Accessor(
                bufferView=2,
                componentType=pygltflib.UNSIGNED_INT,
                count=triangles_u32.size,
                type=pygltflib.SCALAR,
                max=[int(triangles_u32.max())],
                min=[int(triangles_u32.min())],
            ),
        ],
        bufferViews=[
            # BufferView 0: positions
            pygltflib.BufferView(
                buffer=0,
                byteOffset=0,
                byteLength=len(vert_bytes),
                target=pygltflib.ARRAY_BUFFER,
            ),
            # BufferView 1: normals
            pygltflib.BufferView(
                buffer=0,
                byteOffset=norm_offset,
                byteLength=len(norm_bytes),
                target=pygltflib.ARRAY_BUFFER,
            ),
            # BufferView 2: indices
            pygltflib.BufferView(
                buffer=0,
                byteOffset=tri_offset,
                byteLength=len(tri_bytes),
                target=pygltflib.ELEMENT_ARRAY_BUFFER,
            ),
        ],
        buffers=[pygltflib.Buffer(byteLength=total_bytes)],
    )

    gltf.set_binary_blob(blob)
    return b"".join(gltf.save_to_bytes())