VAO Classes

class AbstractVAO(abc.ABC):
    """Abstract base class for Vertex Array Objects (VAOs).

    Defines the interface for different VAO implementations, including
    methods for binding, drawing, setting data, and managing the VAO's
    lifecycle.
    """

    def __init__(self, mode: int = gl.GL_TRIANGLES) -> None:
        """Generate a new OpenGL VAO id and initialize default state.

        Args:
            mode: OpenGL primitive drawing mode (e.g. GL_TRIANGLES).
        """
        self.id = gl.glGenVertexArrays(1)
        self.mode = mode
        self.bound = False
        self.allocated = False
        self.indices_count = 0

    def bind(self) -> None:
        """Bind this VAO as the current OpenGL vertex array."""
        gl.glBindVertexArray(self.id)
        self.bound = True

    def unbind(self) -> None:
        """Unbind the current OpenGL vertex array."""
        gl.glBindVertexArray(0)
        self.bound = False

    def __enter__(self) -> "AbstractVAO":
        """Bind the VAO on entering a `with` block."""
        self.bind()
        return self

    def __exit__(self, exc_type: type, exc_val: BaseException, exc_tb: Any) -> None:
        """Unbind the VAO on exiting a `with` block."""
        self.unbind()

    @abc.abstractmethod
    def draw(self) -> None:
        """Draw the contents of this VAO."""
        raise NotImplementedError

    @abc.abstractmethod
    def set_data(self, data: VertexData) -> None:
        """Upload vertex data to the VAO's buffer(s)."""
        raise NotImplementedError

    @abc.abstractmethod
    def remove_vao(self) -> None:
        """Delete the VAO and its associated buffers."""
        raise NotImplementedError

    def set_vertex_attribute_pointer(
        self,
        id: int,
        size: int,
        type: int,
        stride: int,
        offset: int,
        normalize: bool = False,
    ) -> None:
        """Configure and enable a vertex attribute pointer for the bound buffer.

        Args:
            id: Attribute location index.
            size: Number of components per vertex attribute.
            type: OpenGL data type of each component (e.g. GL_FLOAT).
            stride: Byte offset between consecutive vertex attributes.
            offset: Byte offset of the first component in the buffer.
            normalize: Whether integer data should be normalized.
        """
        if not self.bound:
            logger.error("VAO not bound in set_vertex_attribute_pointer")
        gl.glVertexAttribPointer(
            id, size, type, normalize, stride, ctypes.c_void_p(offset)
        )
        gl.glEnableVertexAttribArray(id)

    def set_num_indices(self, count: int) -> None:
        """Set the number of indices/vertices to draw."""
        self.indices_count = count

    def num_indices(self) -> int:
        """Return the number of indices/vertices to draw."""
        return self.indices_count

    def get_mode(self) -> int:
        """Return the OpenGL primitive drawing mode."""
        return self.mode

    def set_mode(self, mode: int) -> None:
        """Set the OpenGL primitive drawing mode."""
        self.mode = mode

    @abc.abstractmethod
    def get_buffer_id(self, index: int = 0) -> int:
        """Return the OpenGL buffer id at the given index."""
        raise NotImplementedError

    @abc.abstractmethod
    def map_buffer(self, index: int = 0, access_mode: int = gl.GL_READ_WRITE) -> Any:
        """Map the buffer at the given index into client memory."""
        raise NotImplementedError

    def unmap_buffer(self) -> None:
        """Unmap the currently mapped GL_ARRAY_BUFFER."""
        gl.glUnmapBuffer(gl.GL_ARRAY_BUFFER)

    def get_id(self) -> int:
        """Return the OpenGL VAO id."""
        return self.id

class SimpleVAO(AbstractVAO):
    """A basic VAO implementation that uses a single buffer for non-indexed drawing."""

    def __init__(self, mode: int = gl.GL_TRIANGLES) -> None:
        """Create the VAO and generate its single vertex buffer."""
        super().__init__(mode)
        self.buffer = gl.glGenBuffers(1)

    def draw(self) -> None:
        """Draw the VAO's vertices using glDrawArrays."""
        if self.bound and self.allocated:
            gl.glDrawArrays(self.mode, 0, self.indices_count)
        else:
            logger.error("SimpleVAO not bound or not allocated")

    def set_data(self, data: VertexData) -> None:
        """Upload vertex data to the buffer.

        Raises:
            TypeError: If data is not a VertexData instance.
            RuntimeError: If the VAO is not currently bound.
        """
        if not isinstance(data, VertexData):
            logger.error("SimpleVAO: Invalid data type")
            raise TypeError("data must be of type VertexData")
        if not self.bound:
            logger.error("SimpleVAO not bound")
            raise RuntimeError("SimpleVAO not bound")
        gl.glBindBuffer(gl.GL_ARRAY_BUFFER, self.buffer)
        gl.glBufferData(gl.GL_ARRAY_BUFFER, data.data.nbytes, data.data, data.mode)
        self.allocated = True
        self.indices_count = data.size

    def num_indices(self) -> int:
        """Return the number of vertices to draw."""
        return self.indices_count

    def remove_vao(self) -> None:
        """Delete the VAO's buffer and vertex array."""
        gl.glDeleteBuffers(1, [self.buffer])
        gl.glDeleteVertexArrays(1, [self.id])

    def get_buffer_id(self, index: int = 0) -> int:
        """Return the OpenGL buffer id (index is ignored, only one buffer)."""
        return self.buffer

    def map_buffer(self, index: int = 0, access_mode: int = gl.GL_READ_WRITE) -> Any:
        """Map the buffer into client memory (index is ignored, only one buffer)."""
        gl.glBindBuffer(gl.GL_ARRAY_BUFFER, self.buffer)
        return gl.glMapBuffer(gl.GL_ARRAY_BUFFER, access_mode)

class SimpleIndexVAO(AbstractVAO):
    """A VAO implementation that uses an index buffer for indexed drawing."""

    def __init__(self, mode: int = gl.GL_TRIANGLES) -> None:
        """Create the VAO and generate its vertex and index buffers."""
        super().__init__(mode)
        self.buffer = gl.glGenBuffers(1)
        self.idx_buffer = gl.glGenBuffers(1)
        self.index_type = gl.GL_UNSIGNED_INT

    def draw(self) -> None:
        """Draw the VAO's vertices using glDrawElements."""
        if self.bound and self.allocated:
            gl.glDrawElements(self.mode, self.indices_count, self.index_type, None)
        else:
            logger.error("SimpleIndexVAO not bound or not allocated")

    def set_data(self, data: IndexVertexData) -> None:
        """Upload vertex and index data to their respective buffers.

        Raises:
            TypeError: If data is not an IndexVertexData instance.
        """
        if not isinstance(data, IndexVertexData):
            logger.error("SimpleIndexVAO: Unsupported index type")
            raise TypeError("data must be of type IndexVertexData")

        gl.glBindBuffer(gl.GL_ARRAY_BUFFER, self.buffer)
        gl.glBufferData(gl.GL_ARRAY_BUFFER, data.data.nbytes, data.data, data.mode)

        gl.glBindBuffer(gl.GL_ELEMENT_ARRAY_BUFFER, self.idx_buffer)
        gl.glBufferData(
            gl.GL_ELEMENT_ARRAY_BUFFER, data.indices.nbytes, data.indices, data.mode
        )

        self.allocated = True
        self.indices_count = len(data.indices)
        self.index_type = data.index_type

    def remove_vao(self) -> None:
        """Delete the VAO's vertex buffer, index buffer, and vertex array."""
        gl.glDeleteBuffers(1, [self.buffer])
        gl.glDeleteBuffers(1, [self.idx_buffer])
        gl.glDeleteVertexArrays(1, [self.id])

    def get_buffer_id(self, index: int = 0) -> int:
        """Return the OpenGL vertex buffer id (index is ignored)."""
        return self.buffer

    def map_buffer(self, index: int = 0, access_mode: int = gl.GL_READ_WRITE) -> Any:
        """Map the vertex buffer into client memory (index is ignored)."""
        gl.glBindBuffer(gl.GL_ARRAY_BUFFER, self.buffer)
        return gl.glMapBuffer(gl.GL_ARRAY_BUFFER, access_mode)

class MultiBufferVAO(AbstractVAO):
    """A VAO implementation that can manage multiple vertex buffers.

    Useful for separating different types of vertex attributes (e.g.
    positions, colors, normals) into different buffers.
    """

    def __init__(self, mode: int = gl.GL_TRIANGLES) -> None:
        """Create the VAO with no buffers allocated yet."""
        super().__init__(mode)
        self.vbo_ids: list[int] = []

    def draw(self) -> None:
        """Draw the VAO's vertices using glDrawArrays."""
        if self.bound and self.allocated:
            gl.glDrawArrays(self.mode, 0, self.indices_count)
        else:
            logger.error("MultiBufferVAO is not bound or not allocated")

    def set_data(self, data: VertexData, index: int | None = None) -> None:
        """Upload vertex data to the buffer at index, creating buffers as needed.

        Args:
            data: The vertex data to upload.
            index: Buffer index to upload to; appends a new buffer if None.

        Raises:
            TypeError: If data is not a VertexData instance.
        """
        if not isinstance(data, VertexData):
            logger.error("MultiBufferVAO: Invalid data type")
            raise TypeError("data must be of type VertexData")
        if index is None:
            index = len(self.vbo_ids)

        if index >= len(self.vbo_ids):
            new_buffers = index - len(self.vbo_ids) + 1
            new_ids = gl.glGenBuffers(new_buffers)
            if isinstance(new_ids, np.ndarray):
                self.vbo_ids.extend(new_ids)
            else:
                self.vbo_ids.append(new_ids)

        gl.glBindBuffer(gl.GL_ARRAY_BUFFER, self.vbo_ids[index])
        gl.glBufferData(gl.GL_ARRAY_BUFFER, data.data.nbytes, data.data, data.mode)
        self.allocated = True
        if index == 0:  # Assume first buffer determines the number of indices
            self.indices_count = data.size

    def remove_vao(self) -> None:
        """Delete all of the VAO's buffers and its vertex array."""
        gl.glDeleteBuffers(len(self.vbo_ids), self.vbo_ids)
        gl.glDeleteVertexArrays(1, [self.id])

    def get_buffer_id(self, index: int = 0) -> int:
        """Return the OpenGL buffer id at the given index."""
        return self.vbo_ids[index]

    def map_buffer(self, index: int = 0, access_mode: int = gl.GL_READ_WRITE) -> Any:
        """Map the buffer at the given index into client memory."""
        gl.glBindBuffer(gl.GL_ARRAY_BUFFER, self.vbo_ids[index])
        return gl.glMapBuffer(gl.GL_ARRAY_BUFFER, access_mode)

class VAOFactory:
    """Factory for creating VAOs, extensible with custom creator functions."""

    _creators: Dict[VAOType, Callable[[int], AbstractVAO]] = {}

    @staticmethod
    def register_vao_creator(
        name: VAOType, creator_func: Callable[[int], AbstractVAO]
    ) -> None:
        """Register a creator callable for the given VAO type name."""
        VAOFactory._creators[name] = creator_func

    @staticmethod
    def create_vao(name: VAOType, mode: int) -> AbstractVAO:
        """Create a VAO of the named type with the given draw mode.

        Raises:
            ValueError: If the VAO type is not registered.
        """
        creator = VAOFactory._creators.get(name)
        if not creator:
            logger.warning(f"VAO type '{name}' not found.")
            raise ValueError(name)
        return creator(mode)

class VAOType(enum.Enum):
    """Identifiers for the built-in VAO implementations."""

    SIMPLE = "simpleVAO"
    MULTI_BUFFER = "multiBufferVAO"
    SIMPLE_INDEX = "simpleIndexVAO"

class VertexData:
    """A simple data structure to hold vertex data for a VAO."""

    def __init__(
        self, data: np.ndarray | list[float], size: int, mode: int = gl.GL_STATIC_DRAW
    ) -> None:
        """Store vertex data as a float32 numpy array along with its size and draw mode.

        Args:
            data: Vertex data, either a numpy array or a list of floats.
            size: Number of vertices represented by this data.
            mode: OpenGL buffer usage hint (e.g. GL_STATIC_DRAW).
        """
        if isinstance(data, np.ndarray):
            self.data = data
        else:
            self.data = np.array(data, dtype=np.float32)
        self.size = size
        self.mode = mode

class IndexVertexData(VertexData):
    """Vertex data paired with an index buffer for indexed drawing."""

    def __init__(
        self,
        data: np.ndarray | list[float],
        size: int,
        indices: np.ndarray | list[int],
        index_type: int,
        mode: int = gl.GL_STATIC_DRAW,
    ) -> None:
        """Store vertex data plus an index array of the given GL index type.

        Args:
            data: Vertex data, either a numpy array or a list of floats.
            size: Number of vertices represented by this data.
            indices: Index values referencing vertices in data.
            index_type: OpenGL index type (e.g. GL_UNSIGNED_INT).
            mode: OpenGL buffer usage hint (e.g. GL_STATIC_DRAW).

        Raises:
            TypeError: If index_type is not a supported GL index type.
        """
        super().__init__(data, size, mode)
        gl.GL_to_numpy_type = {
            gl.GL_UNSIGNED_INT: np.uint32,
            gl.GL_UNSIGNED_SHORT: np.uint16,
            gl.GL_UNSIGNED_BYTE: np.uint8,
        }
        numpy_dtype = gl.GL_to_numpy_type.get(index_type)
        if numpy_dtype is None:
            logger.error("SimpleIndexVAO: Unsupported index type")
            raise TypeError(f"Unsupported index type: {index_type}")

        self.indices = np.array(indices, dtype=numpy_dtype)
        self.index_type = index_type