Source code for simvx.core.surface_tool

"""Procedural geometry construction tool and convenience mesh generators.

SurfaceTool provides a vertex-by-vertex mesh building API inspired by Godot's
SurfaceTool. ImmediateGeometry3D rebuilds its mesh every frame from a user
callback, suitable for debug visualisations, trails, and dynamic geometry.

Usage::

    from simvx.core import SurfaceTool, PrimitiveType

    st = SurfaceTool()
    st.begin(PrimitiveType.TRIANGLES)
    st.set_normal((0, 1, 0))
    st.set_uv((0, 0))
    st.add_vertex((0, 0, 0))
    st.set_uv((1, 0))
    st.add_vertex((1, 0, 0))
    st.set_uv((0.5, 1))
    st.add_vertex((0.5, 1, 0))
    mesh = st.commit()
"""


from __future__ import annotations

import logging
import math
from collections.abc import Sequence
from enum import IntEnum

import numpy as np

from .graphics.mesh import Mesh
from .math.types import Vec2, Vec3
from .nodes_3d.node3d import Node3D

log = logging.getLogger(__name__)

__all__ = [
    "SurfaceTool",
    "PrimitiveType",
    "ImmediateGeometry3D",
    "create_box",
    "create_sphere",
    "create_cylinder",
    "create_plane",
    "create_capsule",
]




class PrimitiveType(IntEnum):
    """Primitive topology for SurfaceTool."""

    TRIANGLES = 0
    LINES = 1
    POINTS = 2
    TRIANGLE_STRIP = 3





class SurfaceTool:
    """Vertex-by-vertex mesh construction tool.

    Accumulates vertex attributes then emits a ``Mesh`` on ``commit()``.
    Set per-vertex attributes (normal, uv, colour) *before* calling
    ``add_vertex()`` -- the current values are latched when the vertex is added.
    """

    __slots__ = (
        "_primitive", "_positions", "_normals", "_uvs", "_colours", "_indices",
        "_cur_normal", "_cur_uv", "_cur_colour", "_tangents",
    )

    def __init__(self) -> None:
        self._primitive: PrimitiveType | None = None
        self._positions: list[tuple[float, float, float]] = []
        self._normals: list[tuple[float, float, float]] = []
        self._uvs: list[tuple[float, float]] = []
        self._colours: list[tuple[float, float, float, float]] = []
        self._indices: list[int] = []
        self._cur_normal: tuple[float, float, float] = (0.0, 0.0, 1.0)
        self._cur_uv: tuple[float, float] = (0.0, 0.0)
        self._cur_colour: tuple[float, float, float, float] = (1.0, 1.0, 1.0, 1.0)

    # -- Configuration -------------------------------------------------------



    def begin(self, primitive: PrimitiveType = PrimitiveType.TRIANGLES) -> None:
        """Start a new surface. Clears any previously accumulated data."""
        self._primitive = PrimitiveType(primitive)
        self._positions.clear()
        self._normals.clear()
        self._uvs.clear()
        self._colours.clear()
        self._indices.clear()
        self._cur_normal = (0.0, 0.0, 1.0)
        self._cur_uv = (0.0, 0.0)
        self._cur_colour = (1.0, 1.0, 1.0, 1.0)


    # -- Per-vertex attribute setters ----------------------------------------



    def set_normal(self, normal: Sequence[float] | Vec3) -> None:
        """Set the normal for subsequent vertices."""
        n = normal
        self._cur_normal = (float(n[0]), float(n[1]), float(n[2]))




    def set_uv(self, uv: Sequence[float] | Vec2) -> None:
        """Set the UV coordinate for subsequent vertices."""
        self._cur_uv = (float(uv[0]), float(uv[1]))




    def set_colour(self, colour: Sequence[float]) -> None:
        """Set the vertex colour (RGBA) for subsequent vertices."""
        r, g, b = float(colour[0]), float(colour[1]), float(colour[2])
        a = float(colour[3]) if len(colour) > 3 else 1.0
        self._cur_colour = (r, g, b, a)


    # -- Vertex / index accumulation -----------------------------------------



    def add_vertex(self, position: Sequence[float] | Vec3) -> None:
        """Add a vertex with the currently set attributes."""
        if self._primitive is None:
            raise RuntimeError("Call begin() before add_vertex()")
        p = position
        self._positions.append((float(p[0]), float(p[1]), float(p[2])))
        self._normals.append(self._cur_normal)
        self._uvs.append(self._cur_uv)
        self._colours.append(self._cur_colour)




    def add_index(self, index: int) -> None:
        """Add an index referencing a previously added vertex."""
        self._indices.append(int(index))




    def add_triangle_fan(
        self,
        vertices: Sequence[Sequence[float]],
        uvs: Sequence[Sequence[float]] | None = None,
        normals: Sequence[Sequence[float]] | None = None,
    ) -> None:
        """Add a triangle fan as indexed triangles.

        The first vertex is the hub; triangles are formed with consecutive
        pairs of the remaining vertices.
        """
        if self._primitive is None:
            raise RuntimeError("Call begin() before add_triangle_fan()")
        base = len(self._positions)
        for i, v in enumerate(vertices):
            if normals is not None:
                self.set_normal(normals[i])
            if uvs is not None:
                self.set_uv(uvs[i])
            self.add_vertex(v)
        for i in range(1, len(vertices) - 1):
            self._indices.extend([base, base + i, base + i + 1])


    # -- Post-processing -----------------------------------------------------



    def generate_normals(self) -> None:
        """Auto-compute smooth normals from triangle geometry.

        Only works when primitive type is TRIANGLES. Uses the index list if
        present, otherwise assumes every three consecutive vertices form a
        triangle.
        """
        if self._primitive not in (PrimitiveType.TRIANGLES, None):
            return
        n = len(self._positions)
        if n == 0:
            return

        positions = np.array(self._positions, dtype=np.float32)
        normals = np.zeros((n, 3), dtype=np.float32)

        if self._indices:
            idx = np.array(self._indices, dtype=np.int32).reshape(-1, 3)
        else:
            idx = np.arange(n, dtype=np.int32).reshape(-1, 3)

        v0 = positions[idx[:, 0]]
        v1 = positions[idx[:, 1]]
        v2 = positions[idx[:, 2]]
        face_normals = np.cross(v1 - v0, v2 - v0)

        for i, tri in enumerate(idx):
            normals[tri] += face_normals[i]

        lengths = np.linalg.norm(normals, axis=1, keepdims=True)
        np.divide(normals, lengths, where=lengths > 1e-10, out=normals)
        self._normals = [tuple(normals[i]) for i in range(n)]




    def generate_tangents(self) -> None:
        """Auto-compute tangents using Lengyel's method (MikkTSpace-lite).

        Stores tangent data on the resulting mesh as a ``tangents`` attribute
        (Nx4 float32 array with w = handedness). This is a best-effort
        implementation suitable for normal mapping.
        """
        n = len(self._positions)
        if n == 0:
            return

        positions = np.array(self._positions, dtype=np.float32)
        uvs = np.array(self._uvs, dtype=np.float32)
        normals_arr = np.array(self._normals, dtype=np.float32)

        tan1 = np.zeros((n, 3), dtype=np.float32)
        tan2 = np.zeros((n, 3), dtype=np.float32)

        if self._indices:
            idx = np.array(self._indices, dtype=np.int32).reshape(-1, 3)
        else:
            idx = np.arange(n, dtype=np.int32).reshape(-1, 3)

        i0, i1, i2 = idx[:, 0], idx[:, 1], idx[:, 2]
        dp1 = positions[i1] - positions[i0]
        dp2 = positions[i2] - positions[i0]
        duv1 = uvs[i1] - uvs[i0]
        duv2 = uvs[i2] - uvs[i0]

        denom = duv1[:, 0] * duv2[:, 1] - duv1[:, 1] * duv2[:, 0]
        r = np.where(np.abs(denom) > 1e-10, 1.0 / denom, 0.0)[:, None]

        sdir = (dp1 * duv2[:, 1:2] - dp2 * duv1[:, 1:2]) * r
        tdir = (dp2 * duv1[:, 0:1] - dp1 * duv2[:, 0:1]) * r

        for i in range(len(idx)):
            tan1[i0[i]] += sdir[i]
            tan1[i1[i]] += sdir[i]
            tan1[i2[i]] += sdir[i]
            tan2[i0[i]] += tdir[i]
            tan2[i1[i]] += tdir[i]
            tan2[i2[i]] += tdir[i]

        # Gram-Schmidt orthogonalise
        dot_tn = np.sum(normals_arr * tan1, axis=1, keepdims=True)
        tangent3 = tan1 - normals_arr * dot_tn
        lengths = np.linalg.norm(tangent3, axis=1, keepdims=True)
        np.divide(tangent3, lengths, where=lengths > 1e-10, out=tangent3)

        # Handedness
        w = np.sign(np.sum(np.cross(normals_arr, tan1) * tan2, axis=1))
        w[w == 0] = 1.0

        self._tangents = np.column_stack([tangent3, w]).astype(np.float32)


    # -- Output --------------------------------------------------------------



    def commit(self) -> Mesh:
        """Build and return a ``Mesh`` from the accumulated vertex data."""
        if self._primitive is None:
            raise RuntimeError("Call begin() before commit()")
        if not self._positions:
            raise RuntimeError("No vertices added")

        topology_map = {
            PrimitiveType.TRIANGLES: "triangles",
            PrimitiveType.LINES: "lines",
            PrimitiveType.POINTS: "points",
            PrimitiveType.TRIANGLE_STRIP: "triangle_strip",
        }

        positions = np.array(self._positions, dtype=np.float32)
        normals = np.array(self._normals, dtype=np.float32)
        texcoords = np.array(self._uvs, dtype=np.float32)
        if self._indices:
            indices = np.array(self._indices, dtype=np.uint32)
        else:
            # No explicit indices — generate sequential (non-indexed draw as indexed)
            indices = np.arange(len(self._positions), dtype=np.uint32)

        mesh = Mesh(positions, indices, normals, texcoords, topology=topology_map[self._primitive])

        # Attach optional tangent data
        if hasattr(self, "_tangents"):
            mesh.tangents = self._tangents

        # Attach vertex colours if any non-white colour was set
        colours = np.array(self._colours, dtype=np.float32)
        if not np.allclose(colours, 1.0):
            mesh.colours = colours

        return mesh




    def clear(self) -> None:
        """Reset the tool for reuse (same as constructing a new instance)."""
        self._primitive = None
        self._positions.clear()
        self._normals.clear()
        self._uvs.clear()
        self._colours.clear()
        self._indices.clear()
        self._cur_normal = (0.0, 0.0, 1.0)
        self._cur_uv = (0.0, 0.0)
        self._cur_colour = (1.0, 1.0, 1.0, 1.0)
        if hasattr(self, "_tangents"):
            del self._tangents


    # -- Introspection -------------------------------------------------------

    @property
    def vertex_count(self) -> int:
        """Number of vertices accumulated so far."""
        return len(self._positions)

    @property
    def index_count(self) -> int:
        """Number of indices accumulated so far."""
        return len(self._indices)



# ============================================================================
# ImmediateGeometry3D — rebuilt every frame
# ============================================================================




class ImmediateGeometry3D(Node3D):
    """A 3D node whose mesh is rebuilt every frame via a callback.

    Override ``_draw_geometry`` or connect to the ``draw`` signal to provide
    geometry each frame. The node creates a fresh ``SurfaceTool``, passes it
    to your callback, and commits the result as the renderable mesh.

    Usage::

        class MyTrail(ImmediateGeometry3D):
            def _draw_geometry(self, st: SurfaceTool) -> None:
                st.begin(PrimitiveType.TRIANGLES)
                # ... add vertices ...

    Or with signals::

        def draw_fn(st: SurfaceTool) -> None:
            st.begin(PrimitiveType.TRIANGLES)
            # ...

        ig = ImmediateGeometry3D()
        ig.draw.connect(draw_fn)
    """

    def __init__(self, **kwargs):
        from .descriptors import Signal
        super().__init__(**kwargs)
        self.draw = Signal()
        self.mesh = None
        self.material = None
        self._tool = SurfaceTool()



    def process(self, dt: float) -> None:
        tool = self._tool
        tool.clear()
        self._draw_geometry(tool)
        self.draw.emit(tool)
        if tool.vertex_count > 0:
            self.mesh = tool.commit()
        else:
            self.mesh = None


    def _draw_geometry(self, tool: SurfaceTool) -> None:
        """Override to provide geometry. Called every frame."""
        pass



# ============================================================================
# Convenience primitive generators
# ============================================================================




def create_box(size: float | Sequence[float] = 1.0) -> Mesh:
    """Create a box mesh centred at the origin.

    Args:
        size: Uniform size (float) or per-axis extents (x, y, z).

    Returns:
        A ``Mesh`` with 24 vertices and 36 indices.
    """
    if isinstance(size, int | float):
        sx = sy = sz = float(size) / 2
    else:
        sx, sy, sz = float(size[0]) / 2, float(size[1]) / 2, float(size[2]) / 2

    # Faces: +Z, -Z, +Y, -Y, +X, -X
    face_data = [
        (( 0,  0,  1), [(-sx, -sy,  sz), ( sx, -sy,  sz), ( sx,  sy,  sz), (-sx,  sy,  sz)]),
        (( 0,  0, -1), [( sx, -sy, -sz), (-sx, -sy, -sz), (-sx,  sy, -sz), ( sx,  sy, -sz)]),
        (( 0,  1,  0), [(-sx,  sy,  sz), ( sx,  sy,  sz), ( sx,  sy, -sz), (-sx,  sy, -sz)]),
        (( 0, -1,  0), [(-sx, -sy, -sz), ( sx, -sy, -sz), ( sx, -sy,  sz), (-sx, -sy,  sz)]),
        (( 1,  0,  0), [( sx, -sy,  sz), ( sx, -sy, -sz), ( sx,  sy, -sz), ( sx,  sy,  sz)]),
        ((-1,  0,  0), [(-sx, -sy, -sz), (-sx, -sy,  sz), (-sx,  sy,  sz), (-sx,  sy, -sz)]),
    ]
    face_uvs = [(0, 0), (1, 0), (1, 1), (0, 1)]

    st = SurfaceTool()
    st.begin(PrimitiveType.TRIANGLES)
    for normal, corners in face_data:
        base = st.vertex_count
        st.set_normal(normal)
        for i, pos in enumerate(corners):
            st.set_uv(face_uvs[i])
            st.add_vertex(pos)
        st.add_index(base)
        st.add_index(base + 1)
        st.add_index(base + 2)
        st.add_index(base)
        st.add_index(base + 2)
        st.add_index(base + 3)
    return st.commit()





def create_sphere(radius: float = 1.0, rings: int = 16, sectors: int = 16) -> Mesh:
    """Create a UV sphere mesh centred at the origin.

    Args:
        radius: Sphere radius.
        rings: Number of horizontal rings (latitude divisions).
        sectors: Number of vertical sectors (longitude divisions).

    Returns:
        A ``Mesh`` with smooth normals and UVs.
    """
    st = SurfaceTool()
    st.begin(PrimitiveType.TRIANGLES)

    for ring in range(rings + 1):
        phi = math.pi * ring / rings
        sp, cp = math.sin(phi), math.cos(phi)
        for seg in range(sectors + 1):
            theta = math.tau * seg / sectors
            x = sp * math.cos(theta)
            z = sp * math.sin(theta)
            y = cp
            st.set_normal((x, y, z))
            st.set_uv((seg / sectors, ring / rings))
            st.add_vertex((x * radius, y * radius, z * radius))

    row = sectors + 1
    for ring in range(rings):
        for seg in range(sectors):
            a = ring * row + seg
            b = a + row
            st.add_index(a)
            st.add_index(a + 1)
            st.add_index(b)
            st.add_index(a + 1)
            st.add_index(b + 1)
            st.add_index(b)

    return st.commit()





def create_cylinder(radius: float = 0.5, height: float = 1.0, segments: int = 16) -> Mesh:
    """Create a cylinder mesh along the Y axis, centred at the origin.

    Args:
        radius: Cylinder radius.
        height: Total height.
        segments: Number of radial segments.

    Returns:
        A ``Mesh`` with side, top cap, and bottom cap geometry.
    """
    st = SurfaceTool()
    st.begin(PrimitiveType.TRIANGLES)
    half_h = height / 2

    # Side vertices: two rings
    for ring in range(2):
        y = half_h if ring == 0 else -half_h
        v = float(ring)
        for seg in range(segments + 1):
            theta = math.tau * seg / segments
            cx, cz = math.cos(theta), math.sin(theta)
            st.set_normal((cx, 0, cz))
            st.set_uv((seg / segments, v))
            st.add_vertex((cx * radius, y, cz * radius))

    row = segments + 1
    for seg in range(segments):
        a, b = seg, seg + row
        st.add_index(a)
        st.add_index(a + 1)
        st.add_index(b)
        st.add_index(a + 1)
        st.add_index(b + 1)
        st.add_index(b)

    # Top cap
    tc = st.vertex_count
    st.set_normal((0, 1, 0))
    st.set_uv((0.5, 0.5))
    st.add_vertex((0, half_h, 0))
    for i in range(segments):
        theta = math.tau * i / segments
        st.set_uv((0.5 + 0.5 * math.cos(theta), 0.5 + 0.5 * math.sin(theta)))
        st.add_vertex((math.cos(theta) * radius, half_h, math.sin(theta) * radius))
    for i in range(segments):
        st.add_index(tc)
        st.add_index(tc + 1 + (i + 1) % segments)
        st.add_index(tc + 1 + i)

    # Bottom cap
    bc = st.vertex_count
    st.set_normal((0, -1, 0))
    st.set_uv((0.5, 0.5))
    st.add_vertex((0, -half_h, 0))
    for i in range(segments):
        theta = math.tau * i / segments
        st.set_uv((0.5 + 0.5 * math.cos(theta), 0.5 + 0.5 * math.sin(theta)))
        st.add_vertex((math.cos(theta) * radius, -half_h, math.sin(theta) * radius))
    for i in range(segments):
        st.add_index(bc)
        st.add_index(bc + 1 + i)
        st.add_index(bc + 1 + (i + 1) % segments)

    return st.commit()





def create_plane(size: float | Sequence[float] = 1.0, subdivisions: int = 1) -> Mesh:
    """Create a flat plane on the XZ plane, centred at the origin.

    Args:
        size: Uniform size (float) or (width, depth) tuple.
        subdivisions: Number of subdivisions per axis (1 = single quad).

    Returns:
        A ``Mesh`` with upward-facing normals.
    """
    if isinstance(size, int | float):
        w = d = float(size)
    else:
        w, d = float(size[0]), float(size[1])

    st = SurfaceTool()
    st.begin(PrimitiveType.TRIANGLES)
    st.set_normal((0, 1, 0))

    hw, hd = w / 2, d / 2
    rows = cols = subdivisions
    for row in range(rows + 1):
        for col in range(cols + 1):
            u = col / cols
            v = row / rows
            x = -hw + u * w
            z = -hd + v * d
            st.set_uv((u, v))
            st.add_vertex((x, 0, z))

    stride = cols + 1
    for row in range(rows):
        for col in range(cols):
            a = row * stride + col
            b = a + stride
            st.add_index(a)
            st.add_index(a + 1)
            st.add_index(b)
            st.add_index(a + 1)
            st.add_index(b + 1)
            st.add_index(b)

    return st.commit()





def create_capsule(radius: float = 0.5, height: float = 1.0, rings: int = 8, sectors: int = 16) -> Mesh:
    """Create a capsule mesh (cylinder with hemispherical caps) along the Y axis.

    Args:
        radius: Capsule radius.
        height: Total height including the hemispherical caps.
        rings: Number of rings per hemisphere.
        sectors: Number of radial sectors.

    Returns:
        A ``Mesh`` with smooth normals.
    """
    st = SurfaceTool()
    st.begin(PrimitiveType.TRIANGLES)

    # Cylinder body height (excluding caps)
    cyl_h = max(0.0, height - 2 * radius)
    half_cyl = cyl_h / 2

    # Top hemisphere
    for ring in range(rings + 1):
        phi = (math.pi / 2) * ring / rings  # 0 .. pi/2
        sp, cp = math.sin(phi), math.cos(phi)
        y_offset = half_cyl + cp * radius
        for seg in range(sectors + 1):
            theta = math.tau * seg / sectors
            x = sp * math.cos(theta)
            z = sp * math.sin(theta)
            st.set_normal((x, cp, z))
            st.set_uv((seg / sectors, ring / (2 * rings + 1)))
            st.add_vertex((x * radius, y_offset, z * radius))

    # Cylinder body (two rings)
    cyl_base = st.vertex_count
    for ring_i in range(2):
        y = half_cyl if ring_i == 0 else -half_cyl
        v_coord = (rings + ring_i) / (2 * rings + 1)
        for seg in range(sectors + 1):
            theta = math.tau * seg / sectors
            cx, cz = math.cos(theta), math.sin(theta)
            st.set_normal((cx, 0, cz))
            st.set_uv((seg / sectors, v_coord))
            st.add_vertex((cx * radius, y, cz * radius))

    # Bottom hemisphere
    bot_base = st.vertex_count
    for ring in range(rings + 1):
        phi = math.pi / 2 + (math.pi / 2) * ring / rings  # pi/2 .. pi
        sp, cp = math.sin(phi), math.cos(phi)
        y_offset = -half_cyl + cp * radius
        for seg in range(sectors + 1):
            theta = math.tau * seg / sectors
            x = sp * math.cos(theta)
            z = sp * math.sin(theta)
            st.set_normal((x, cp, z))
            st.set_uv((seg / sectors, (rings + 1 + ring) / (2 * rings + 1)))
            st.add_vertex((x * radius, y_offset, z * radius))

    row = sectors + 1

    # Top hemisphere indices
    for ring in range(rings):
        for seg in range(sectors):
            a = ring * row + seg
            b = a + row
            st.add_index(a)
            st.add_index(a + 1)
            st.add_index(b)
            st.add_index(a + 1)
            st.add_index(b + 1)
            st.add_index(b)

    # Connect top hemisphere to cylinder body top ring
    top_last_ring = rings * row
    for seg in range(sectors):
        a = top_last_ring + seg
        b = cyl_base + seg
        st.add_index(a)
        st.add_index(a + 1)
        st.add_index(b)
        st.add_index(a + 1)
        st.add_index(b + 1)
        st.add_index(b)

    # Cylinder body indices
    for seg in range(sectors):
        a = cyl_base + seg
        b = a + row
        st.add_index(a)
        st.add_index(a + 1)
        st.add_index(b)
        st.add_index(a + 1)
        st.add_index(b + 1)
        st.add_index(b)

    # Connect cylinder body bottom ring to bottom hemisphere
    cyl_bot_ring = cyl_base + row
    for seg in range(sectors):
        a = cyl_bot_ring + seg
        b = bot_base + seg
        st.add_index(a)
        st.add_index(a + 1)
        st.add_index(b)
        st.add_index(a + 1)
        st.add_index(b + 1)
        st.add_index(b)

    # Bottom hemisphere indices
    for ring in range(rings):
        for seg in range(sectors):
            a = bot_base + ring * row + seg
            b = a + row
            st.add_index(a)
            st.add_index(a + 1)
            st.add_index(b)
            st.add_index(a + 1)
            st.add_index(b + 1)
            st.add_index(b)

    return st.commit()