WebGPU Classes

class WebGPUWidget(QWidget, metaclass=QWidgetABCMeta):
    """An abstract base class for offscreen WebGPU widgets.

    The rendered content is presented via QPainter.

    The scene is rendered to an offscreen texture, read back to a numpy
    buffer, and blitted to the widget with QPainter. This keeps the whole
    data flow visible and allows QPainter text overlays via render_text().

    The readback is pipelined through a small ring of buffers: each frame
    the current image is copied into one buffer while the *previous*
    frame's buffer is mapped and read. Mapping a buffer whose copy was
    submitted a frame ago returns almost immediately, so the CPU never
    stalls waiting for the GPU to drain. The presented image therefore
    lags the simulation by one frame, which is imperceptible at
    interactive rates.

    Subclasses must implement paintWebGPU() and resizeWebGPU(), and must
    set self.device to a wgpu device before any rendering can occur. The
    base class owns the render target textures and the readback machinery;
    resizeWebGPU() should only update subclass state (projection matrix,
    per-pipeline sizes etc.) and must NOT recreate render buffers itself.
    """

    def __init__(self) -> None:
        """Initialize the widget.

        Note the wgpu device is not created here; subclasses are
        responsible for creating it (self.device) before the first paint.
        """
        super().__init__()
        self.device: Optional[wgpu.GPUDevice] = None
        self.msaa_sample_count = 4
        self.text_buffer: List[Tuple[int, int, str, int, str, QColor]] = []
        self.frame_buffer: Optional[np.ndarray] = None
        self._update_timer = QTimer(self)
        self._update_timer.timeout.connect(self.update)
        # Device pixel ratio for high-DPI displays. This is re-queried on
        # resize and paint, as it changes when the window moves between
        # screens with different scale factors.
        self.ratio = self.devicePixelRatioF()
        self._initialize_buffer()

    # ------------------------------------------------------------------
    # Timer control
    # ------------------------------------------------------------------
    def start_update_timer(self, interval_ms: int) -> None:
        """Start the update timer with the given interval.

        Args:
            interval_ms (int): The interval in milliseconds.
        """
        self._update_timer.start(interval_ms)

    def stop_update_timer(self) -> None:
        """Stop the update timer."""
        self._update_timer.stop()

    # ------------------------------------------------------------------
    # Abstract interface
    # ------------------------------------------------------------------
    @abstractmethod
    def resizeWebGPU(self, width: int, height: int) -> None:
        """Handle a resize, after the buffers are recreated at the new size.

        Subclasses should update projection matrices and any per-pipeline
        sizes here. Width and height are in device pixels.

        Args:
            width (int): New width in device pixels.
            height (int): New height in device pixels.
        """
        pass

    @abstractmethod
    def paintWebGPU(self) -> None:
        """Render the WebGPU content.

        Called on every paint event; all the main rendering code goes
        here. Implementations should end by calling
        _update_colour_buffer() so the rendered frame is read back for
        presentation.
        """
        pass

    # ------------------------------------------------------------------
    # Qt event handlers
    # ------------------------------------------------------------------
    def resizeEvent(self, event: QResizeEvent) -> None:
        """Handle window resize.

        Recreates the render targets and the readback frame buffer at the
        new size (in device pixels), then notifies the subclass via
        resizeWebGPU().
        """
        self.ratio = self.devicePixelRatioF()
        width = int(event.size().width() * self.ratio)
        height = int(event.size().height() * self.ratio)

        # A minimised or degenerate window gives a zero dimension, which
        # is a wgpu validation error - skip until we have a real size.
        if width <= 0 or height <= 0:
            return super().resizeEvent(event)

        self.texture_size = (width, height)
        self.frame_buffer = np.zeros([height, width, 4], dtype=np.uint8)

        if self.device is not None:
            self._create_render_buffer()
            self.resizeWebGPU(width, height)

        return super().resizeEvent(event)

    def paintEvent(self, event: QPaintEvent) -> None:
        """Handle the paint event to render and present the WebGPU content."""
        # The device pixel ratio can change without a resize when the
        # window moves between screens - rebuild render targets if it has.
        current_ratio = self.devicePixelRatioF()
        if current_ratio != self.ratio:
            self.ratio = current_ratio
            width = int(self.width() * self.ratio)
            height = int(self.height() * self.ratio)
            if width > 0 and height > 0:
                self.texture_size = (width, height)
                self.frame_buffer = np.zeros([height, width, 4], dtype=np.uint8)
                if self.device is not None:
                    self._create_render_buffer()
                    self.resizeWebGPU(width, height)

        if self.device is not None:
            self.paintWebGPU()

        painter = QPainter(self)
        if self.frame_buffer is not None:
            self._present_image(painter, self.frame_buffer)

        # Draw any queued text. Font size is scaled relative to a base
        # window height so text keeps its proportions when resizing.
        base_height = 600.0
        scale_factor = self.height() / base_height
        for x, y, text, size, font, colour in self.text_buffer:
            scaled_size = int(size * scale_factor)
            painter.setPen(colour)
            painter.setFont(QFont(font, scaled_size))
            draw_y = y
            if y < 0:
                draw_y = self.height() + y
            painter.drawText(x, draw_y, text)
        self.text_buffer.clear()
        painter.end()

    # ------------------------------------------------------------------
    # Buffer / texture management
    # ------------------------------------------------------------------
    def _initialize_buffer(self) -> None:
        """Initialize the numpy buffer used for the final framebuffer render."""
        width = max(1, int(self.width() * self.ratio))
        height = max(1, int(self.height() * self.ratio))
        self.frame_buffer = np.zeros([height, width, 4], dtype=np.uint8)
        self.texture_size = (width, height)

    def _create_render_buffer(self) -> None:
        """Create the render target textures and the readback buffer ring.

        Requires self.device to be set. Called automatically on resize;
        subclasses should call it once after creating the device.
        Recreating the ring resets the pending flags, which conveniently
        discards any in-flight copy at the old size.
        """
        if self.device is None:
            raise RuntimeError("self.device must be set before creating render buffers")

        # The single-sample texture the MSAA render is resolved into, and
        # which is copied back to the CPU for presentation.
        self.colour_buffer_texture = self.device.create_texture(
            size=self.texture_size,
            sample_count=1,
            format=wgpu.TextureFormat.rgba8unorm,
            usage=wgpu.TextureUsage.RENDER_ATTACHMENT | wgpu.TextureUsage.COPY_SRC,
        )
        self.colour_buffer_texture_view = self.colour_buffer_texture.create_view()

        # The multisampled texture that is actually rendered to.
        self.multisample_texture = self.device.create_texture(
            size=self.texture_size,
            sample_count=self.msaa_sample_count,
            format=wgpu.TextureFormat.rgba8unorm,
            usage=wgpu.TextureUsage.RENDER_ATTACHMENT,
        )
        self.multisample_texture_view = self.multisample_texture.create_view()

        # Depth buffer (multisampled to match). Keep an explicit reference
        # to the texture as well as the view so ownership is obvious.
        self.depth_buffer_texture = self.device.create_texture(
            size=self.texture_size,
            format=wgpu.TextureFormat.depth24plus,
            usage=wgpu.TextureUsage.RENDER_ATTACHMENT,
            sample_count=self.msaa_sample_count,
        )
        self.depth_buffer_view = self.depth_buffer_texture.create_view()

        # Ring of buffers for reading the resolved texture back to the
        # CPU. Rows are padded to the spec-mandated 256 byte alignment.
        buffer_size = self._calculate_aligned_buffer_size()
        self.readback_buffers = [
            self.device.create_buffer(
                size=buffer_size,
                usage=wgpu.BufferUsage.COPY_DST | wgpu.BufferUsage.MAP_READ,
                label=f"readback_buffer_{i}",
            )
            for i in range(NUM_READBACK_BUFFERS)
        ]
        # Index of the buffer the next copy will go into, and per-buffer
        # flags recording whether it holds a frame we haven't read yet.
        self._readback_index = 0
        self._readback_pending = [False] * NUM_READBACK_BUFFERS

    def render_text(
        self,
        x: int,
        y: int,
        text: str,
        size: int = 10,
        font: str = "Arial",
        colour: Union[QColor, Qt.GlobalColor] = Qt.black,
    ) -> None:
        """Queue text to be drawn over the rendered frame.

        The text buffer is cleared each frame, so text must be re-added
        every frame (typically from paintWebGPU or a timer callback).
        Note the rendered image lags the simulation by one frame due to
        the pipelined readback, so text meant to track a moving 3D
        object's screen position will swim slightly during fast motion;
        HUD-style labels are unaffected.

        Args:
            x (int): The x-coordinate of the text.
            y (int): The y-coordinate of the text. A negative value
                positions the text relative to the bottom of the window.
            text (str): The text to render.
            size (int, optional): Base font size, scaled with window
                height. Defaults to 10.
            font (str, optional): Font family. Defaults to "Arial".
            colour: Text colour. Defaults to Qt.black.
        """
        self.text_buffer.append((x, y, text, size, font, QColor(colour)))

    def _calculate_aligned_row_size(self) -> int:
        """Calculate the row stride for texture copy operations.

        The stride is padded to the WebGPU spec's
        COPY_BYTES_PER_ROW_ALIGNMENT (256 bytes).
        """
        bytes_per_pixel = 4  # rgba8unorm
        raw_row_size = self.texture_size[0] * bytes_per_pixel
        alignment = COPY_BYTES_PER_ROW_ALIGNMENT
        return ((raw_row_size + alignment - 1) // alignment) * alignment

    def _calculate_aligned_buffer_size(self) -> int:
        """Calculate the total readback buffer size.

        This is the aligned row stride times the number of rows.
        """
        return self._calculate_aligned_row_size() * self.texture_size[1]

    def _update_colour_buffer(self) -> None:
        """Copy this frame out, and read the previous frame back in.

        A buffer is only ever mapped when it is not the copy target, and
        is unmapped before it becomes the copy target again - the
        alternation guarantees this. The first frame after startup or a
        resize has nothing pending, so the previous frame buffer contents
        are shown once and everything flows from the next frame on.
        """
        # A subclass that overrides _create_render_buffer must still create the
        # read-back ring - easiest by calling super()._create_render_buffer().
        # Without it the copy-back below would raise AttributeError on the ring
        # attributes, which the try/except would swallow every single frame,
        # leaving a grey window and no obvious cause. Fail loudly instead.
        if not getattr(self, "readback_buffers", None):
            raise RuntimeError(
                "WebGPUWidget read-back ring is not initialised. A subclass that "
                "overrides _create_render_buffer() must call "
                "super()._create_render_buffer() (or otherwise create the "
                "readback_buffers ring) before rendering."
            )

        bytes_per_row = self._calculate_aligned_row_size()
        width, height = self.texture_size
        try:
            write_idx = self._readback_index
            read_idx = (write_idx + 1) % len(self.readback_buffers)

            # Kick off the copy of this frame. We do not wait for it.
            command_encoder = self.device.create_command_encoder()
            command_encoder.copy_texture_to_buffer(
                {"texture": self.colour_buffer_texture},
                {
                    "buffer": self.readback_buffers[write_idx],
                    "bytes_per_row": bytes_per_row,
                    "rows_per_image": height,
                },
                (width, height, 1),
            )
            self.device.queue.submit([command_encoder.finish()])
            self._readback_pending[write_idx] = True

            # Read back last frame's buffer, if it holds one. The map
            # returns almost immediately because the GPU has had a full
            # frame to complete that copy.
            if self._readback_pending[read_idx]:
                buf = self.readback_buffers[read_idx]
                buf.map_sync(mode=wgpu.MapMode.READ)
                raw_data = buf.read_mapped()

                # The raw data includes per-row padding to meet the
                # alignment requirement, so build a strided view and copy
                # it to a contiguous array before unmapping.
                strided_view = np.lib.stride_tricks.as_strided(
                    np.frombuffer(raw_data, dtype=np.uint8),
                    shape=(height, width, 4),
                    strides=(bytes_per_row, 4, 1),
                )
                self.frame_buffer = np.copy(strided_view)
                buf.unmap()
                self._readback_pending[read_idx] = False

            # Alternate: next frame writes into the buffer just drained.
            self._readback_index = read_idx
        except Exception:
            # A genuine (usually transient) GPU/mapping error - log it with a
            # traceback so it is visible, and fall back to a grey frame rather
            # than tearing down the event loop.
            logger.exception("Failed to update colour buffer")
            if self.frame_buffer is not None:
                self.frame_buffer.fill(128)

    def _present_image(self, painter: QPainter, image_data: np.ndarray) -> None:
        """Present the frame buffer on the widget.

        The image is tagged with the device pixel ratio so Qt blits it
        1:1 to the physical pixels rather than rescaling a
        device-pixel-sized image into a logical-pixel rect.

        Args:
            painter (QPainter): The active painter.
            image_data (np.ndarray): The image data to render.
        """
        height, width, _ = image_data.shape
        image = QImage(
            image_data.data,
            width,
            height,
            image_data.strides[0],
            QImage.Format.Format_RGBA8888,
        )
        image.setDevicePixelRatio(self.ratio)
        painter.drawImage(0, 0, image)

class PipelineType(Enum):
    """Enumeration of available pipeline types."""

    MULTI_COLOURED_LINES = "multi_coloured_lines"
    SINGLE_COLOUR_LINES = "single_colour_lines"
    MULTI_COLOURED_POINTS = "multi_coloured_points"
    SINGLE_COLOUR_POINTS = "single_colour_points"
    MULTI_COLOURED_TRIANGLES = "multi_coloured_triangles"
    SINGLE_COLOUR_TRIANGLES = "single_colour_triangles"
    TRIANGLE_LIST_MULTI_COLOURED = "triangle_list_multi_coloured"
    TRIANGLE_LIST_SINGLE_COLOUR = "triangle_list_single_colour"
    TRIANGLE_STRIP_MULTI_COLOURED = "triangle_strip_multi_coloured"
    TRIANGLE_STRIP_SINGLE_COLOUR = "triangle_strip_single_colour"
    POINT_LIST_MULTI_COLOURED = "point_list_multi_coloured"
    POINT_LIST_SINGLE_COLOUR = "point_list_single_colour"
    MULTI_COLOURED_INSTANCED_GEOMETRY = "multi_coloured_instanced_geometry"
    SINGLE_COLOUR_INSTANCED_GEOMETRY = "single_colour_instanced_geometry"

class _PipelineFactory:
    """Factory for creating pipeline instances with various configurations."""

    def __init__(self) -> None:
        """Initialize the pipeline factory with default pipeline types."""
        self._pipeline_registry: Dict[PipelineType, PipelineFactoryFn] = {}
        self.register_pipeline(
            PipelineType.MULTI_COLOURED_POINTS, PointPipelineMultiColour
        )
        self.register_pipeline(
            PipelineType.SINGLE_COLOUR_POINTS, PointPipelineSingleColour
        )
        self.register_pipeline(
            PipelineType.MULTI_COLOURED_LINES, LinePipelineMultiColour
        )
        self.register_pipeline(
            PipelineType.SINGLE_COLOUR_LINES, LinePipelineSingleColour
        )
        self.register_pipeline(
            PipelineType.MULTI_COLOURED_TRIANGLES, TrianglePipelineMultiColour
        )
        self.register_pipeline(
            PipelineType.SINGLE_COLOUR_TRIANGLES, TrianglePipelineSingleColour
        )

        # Triangle pipelines pinned to a specific topology: registered as
        # factory callables (rather than classes) so the fixed `topology`
        # keyword can be bound while still forwarding any other **kwargs
        # create_pipeline receives (e.g. `colour`, `data_type`) to the
        # underlying pipeline class.
        self.register_pipeline(
            PipelineType.TRIANGLE_LIST_MULTI_COLOURED,
            lambda device, **kwargs: TrianglePipelineMultiColour(
                device, topology=wgpu.PrimitiveTopology.triangle_list, **kwargs
            ),
        )
        self.register_pipeline(
            PipelineType.TRIANGLE_LIST_SINGLE_COLOUR,
            lambda device, **kwargs: TrianglePipelineSingleColour(
                device, topology=wgpu.PrimitiveTopology.triangle_list, **kwargs
            ),
        )
        self.register_pipeline(
            PipelineType.TRIANGLE_STRIP_MULTI_COLOURED,
            lambda device, **kwargs: TrianglePipelineMultiColour(
                device, topology=wgpu.PrimitiveTopology.triangle_strip, **kwargs
            ),
        )
        self.register_pipeline(
            PipelineType.TRIANGLE_STRIP_SINGLE_COLOUR,
            lambda device, **kwargs: TrianglePipelineSingleColour(
                device, topology=wgpu.PrimitiveTopology.triangle_strip, **kwargs
            ),
        )
        self.register_pipeline(
            PipelineType.POINT_LIST_MULTI_COLOURED, PointListPipelineMultiColour
        )
        self.register_pipeline(
            PipelineType.POINT_LIST_SINGLE_COLOUR, PointListPipelineSingleColour
        )
        self.register_pipeline(
            PipelineType.MULTI_COLOURED_INSTANCED_GEOMETRY,
            InstancedGeometryPipelineMultiColour,
        )
        self.register_pipeline(
            PipelineType.SINGLE_COLOUR_INSTANCED_GEOMETRY,
            InstancedGeometryPipelineSingleColour,
        )

    def register_pipeline(
        self, pipeline_type: PipelineType, pipeline_factory: PipelineFactoryFn
    ) -> None:
        """Register a custom pipeline type.

        Args:
            pipeline_type: Enum identifier for the pipeline
            pipeline_factory: Pipeline class, or a factory callable with the
                same `(device, **kwargs) -> BaseWebGPUPipeline` signature, to
                register for this type
        """
        self._pipeline_registry[pipeline_type] = pipeline_factory

    def create_pipeline(
        self, device: wgpu.GPUDevice, pipeline_type: PipelineType, **kwargs: Any
    ) -> BaseWebGPUPipeline:
        """Create a pipeline instance.

        Args:
            device: WebGPU device
            pipeline_type: Type of pipeline to create
            **kwargs: Pipeline-specific configuration parameters

        Returns:
            Configured pipeline instance

        Raises:
            ValueError: If pipeline type is not registered
        """
        if pipeline_type not in self._pipeline_registry:
            raise ValueError(
                f"Unknown pipeline type: {pipeline_type}. Available types: {list(self._pipeline_registry.keys())}"
            )

        pipeline_factory = self._pipeline_registry[pipeline_type]
        return pipeline_factory(device, **kwargs)

class NGLToWebGPU:
    """Maps NGL type names to WebGPU strides and vertex formats."""

    _strides = {
        "vec2": 2 * FLOAT_SIZE,
        "vec3": 3 * FLOAT_SIZE,
        "vec4": 4 * FLOAT_SIZE,
        "mat2": 4 * FLOAT_SIZE,
        "mat3": 12 * FLOAT_SIZE,
        "mat4": 16 * FLOAT_SIZE,
    }
    _vertex_format = {
        "vec2": "float32x2",
        "vec3": "float32x3",
        "vec4": "float32x4",
    }

    @staticmethod
    def stride_from_type(type: str) -> int:
        """Return the byte stride for the given NGL type name."""
        return NGLToWebGPU._strides[type.lower()]

    @staticmethod
    def vertex_format(type: str) -> str:
        """Return the WebGPU vertex format for the given NGL type name."""
        return NGLToWebGPU._vertex_format[type.lower()]

class BaseWebGPUPipeline(ABC):
    """Abstract base class for all WebGPU rendering pipelines.

    Provides common functionality for:
    - Buffer management and creation
    - Pipeline configuration
    - Uniform buffer handling
    - Resource cleanup
    """

    def __init__(
        self,
        device: wgpu.GPUDevice,
        texture_format: wgpu.TextureFormat = wgpu.TextureFormat.rgba8unorm,
        depth_format: wgpu.TextureFormat = wgpu.TextureFormat.depth24plus,
        msaa_sample_count: int = 4,
        data_type: str = "Vec3",
        stride: int = 0,
    ) -> None:
        """Initialize base pipeline.

        Args:
            device: WebGPU device
            texture_format: Colour attachment format
            depth_format: Depth attachment format
            msaa_sample_count: Number of MSAA samples
            data_type: Vertex data type (e.g., "Vec3", "Vec2")
            stride: Vertex buffer stride. If 0, inferred from data_type
        """
        self.device = device
        self.texture_format = texture_format
        self.depth_format = depth_format
        self.msaa_sample_count = msaa_sample_count
        self._data_type = data_type

        if stride != 0:
            self._stride = stride
        else:
            self._stride = NGLToWebGPU.stride_from_type(self._data_type)

        # Core pipeline resources
        self.pipeline: Optional[wgpu.GPURenderPipeline] = None
        self.uniform_buffer: Optional[wgpu.GPUBuffer] = None
        self.bind_group: Optional[wgpu.GPUBindGroup] = None

        # Initialize uniform data structure
        self.uniform_data = np.zeros((), dtype=self.get_dtype())
        self._set_default_uniforms()

        # Create the pipeline
        self._create_pipeline()

    @abstractmethod
    def get_dtype(self) -> np.dtype:
        """Get the numpy dtype for the uniform buffer structure."""
        pass

    @abstractmethod
    def _get_shader_code(self) -> str:
        """Get the WGSL shader code for this pipeline."""
        pass

    @abstractmethod
    def _get_vertex_buffer_layouts(self) -> List[Dict[str, Any]]:
        """Get vertex buffer layout configurations for the pipeline."""
        pass

    @abstractmethod
    def _get_primitive_topology(self) -> wgpu.PrimitiveTopology:
        """Get the primitive topology for the pipeline."""
        pass

    @abstractmethod
    def _set_default_uniforms(self) -> None:
        """Set default values for uniform data."""
        pass

    @abstractmethod
    def _get_pipeline_label(self) -> str:
        """Get the label for the pipeline."""
        pass

    def _create_pipeline(self) -> None:
        """Create the render pipeline and associated resources."""
        # Load shader
        shader_module = self.device.create_shader_module(code=self._get_shader_code())

        # Create render pipeline
        self.pipeline = self.device.create_render_pipeline(
            label=self._get_pipeline_label(),
            layout="auto",
            vertex={
                "module": shader_module,
                "entry_point": "vertex_main",
                "buffers": self._get_vertex_buffer_layouts(),
            },
            fragment={
                "module": shader_module,
                "entry_point": "fragment_main",
                "targets": [{"format": self.texture_format}],
            },
            primitive={"topology": self._get_primitive_topology()},
            depth_stencil={
                "format": self.depth_format,
                "depth_write_enabled": True,
                "depth_compare": wgpu.CompareFunction.less,
            },
            multisample={"count": self.msaa_sample_count},
        )

        # Create uniform buffer
        self.uniform_buffer = self.device.create_buffer_with_data(
            data=self.uniform_data.tobytes(),
            usage=int(wgpu.BufferUsage.UNIFORM | wgpu.BufferUsage.COPY_DST),
            label=f"{self._get_pipeline_label()}_uniform_buffer",
        )

        # Create bind group
        bind_group_layout = self.pipeline.get_bind_group_layout(0)
        self.bind_group = self.device.create_bind_group(
            layout=bind_group_layout,
            entries=[
                {
                    "binding": 0,
                    "resource": {"buffer": self.uniform_buffer},
                }
            ],
        )

    def _create_or_update_buffer(
        self,
        current_buffer: Optional[wgpu.GPUBuffer],
        data: Union[np.ndarray, wgpu.GPUBuffer],
        usage: wgpu.BufferUsage,
        buffer_label: str,
    ) -> Tuple[Optional[wgpu.GPUBuffer], int]:
        """Create or update a GPU buffer with new data.

        Args:
            current_buffer: Existing buffer (may be None)
            data: New data (numpy array or GPU buffer)
            usage: Buffer usage flags
            buffer_label: Label for the buffer

        Returns:
            Tuple of (buffer, data_size)
        """
        if isinstance(data, wgpu.GPUBuffer):
            # Use provided buffer directly
            return data, data.size

        # Handle numpy array
        data_bytes = data.astype(np.float32).tobytes()
        data_size = len(data_bytes)

        # Create new buffer if needed or existing one is too small
        if current_buffer is None or current_buffer.size < data_size:
            if current_buffer:
                current_buffer.destroy()
            buffer = self.device.create_buffer_with_data(
                data=data_bytes,
                usage=usage,
                label=buffer_label,
            )
            return buffer, data_size
        else:
            # Update existing buffer
            self.device.queue.write_buffer(current_buffer, 0, data_bytes)
            return current_buffer, data_size

    def _process_vertex_data(
        self,
        data: Optional[Union[np.ndarray, wgpu.GPUBuffer]],
        default_value: Optional[np.ndarray] = None,
        padding_size: Optional[int] = None,
        buffer_label: str = "vertex_buffer",
    ) -> Optional[Union[wgpu.GPUBuffer, Tuple[wgpu.GPUBuffer, int]]]:
        """Process vertex data, handling numpy arrays, GPU buffers, and defaults.

        Args:
            data: Input data (numpy array, GPU buffer, or None)
            default_value: Default value if data is None
            padding_size: Size to pad arrays to (for alignment)
            buffer_label: Label for created buffers

        Returns:
            Processed buffer(s) or None
        """
        if data is None and default_value is not None:
            data = default_value

        if data is None:
            return None

        if isinstance(data, wgpu.GPUBuffer):
            return data

        # Handle numpy array
        if padding_size:
            # Pad array to specified size
            if data.ndim == 1:
                padded_data = np.zeros(padding_size, dtype=np.float32)
                padded_data[: len(data)] = data.astype(np.float32)
            else:
                padded_data = np.zeros((data.shape[0], padding_size), dtype=np.float32)
                padded_data[:, : data.shape[1]] = data.astype(np.float32)
            data = padded_data

        buffer, _ = self._create_or_update_buffer(
            None,  # Always create new for processed data
            data,
            wgpu.BufferUsage.VERTEX | wgpu.BufferUsage.COPY_DST,
            buffer_label,
        )
        return buffer

    @abstractmethod
    def set_data(self, **kwargs: Any) -> None:
        """Set rendering data (vertices, colours, etc.).

        Args:
            **kwargs: Pipeline-specific data parameters
        """
        pass

    @abstractmethod
    def update_uniforms(self, **kwargs: Any) -> None:
        """Update uniform buffer values.

        Args:
            **kwargs: Pipeline-specific uniform parameters
        """
        pass

    @abstractmethod
    def render(self, render_pass: wgpu.GPURenderPassEncoder, **kwargs: Any) -> None:
        """Render using this pipeline.

        Args:
            render_pass: Active render pass encoder
            **kwargs: Pipeline-specific render parameters
        """
        pass

    def cleanup(self) -> None:
        """Release pipeline resources. Can be overridden for additional cleanup."""
        if self.uniform_buffer:
            self.uniform_buffer.destroy()

class BasePointPipeline(BaseWebGPUPipeline):
    """Base class for point rendering pipelines.

    Provides common functionality for:
    - Point billboarding
    - Quad generation
    - Circle clipping in fragment shader
    """

    def _get_primitive_topology(self) -> wgpu.PrimitiveTopology:
        """Points are rendered as triangle strips for quad generation."""
        return wgpu.PrimitiveTopology.triangle_strip

    def _get_default_vertex_layouts(
        self, has_colour_buffer: bool = False
    ) -> List[Dict[str, Any]]:
        """Get default vertex buffer layouts for point rendering.

        Args:
            has_colour_buffer: Whether to include colour buffer layout

        Returns:
            List of vertex buffer layout configurations
        """
        layouts = [
            {
                "array_stride": self._stride,
                "step_mode": "instance",
                "attributes": [
                    {
                        "format": NGLToWebGPU.vertex_format(self._data_type),
                        "offset": 0,
                        "shader_location": 0,
                    },
                ],
            },
        ]

        if has_colour_buffer:
            layouts.append(
                {
                    "array_stride": NGLToWebGPU.stride_from_type("Vec3"),
                    "step_mode": "instance",
                    "attributes": [
                        {
                            "format": NGLToWebGPU.vertex_format("Vec3"),
                            "offset": 0,
                            "shader_location": 1,
                        },
                    ],
                }
            )

        return layouts

    def _render_points(
        self,
        render_pass: wgpu.GPURenderPassEncoder,
        position_buffer: wgpu.GPUBuffer,
        colour_buffer: Optional[wgpu.GPUBuffer] = None,
        num_points: Optional[int] = None,
    ) -> None:
        """Common point rendering implementation.

        Args:
            render_pass: Active render pass encoder
            position_buffer: Buffer containing point positions
            colour_buffer: Optional buffer containing point colours
            num_points: Number of points to render
        """
        if position_buffer is None:
            return

        count = num_points if num_points is not None else getattr(self, "num_points", 0)

        render_pass.set_pipeline(self.pipeline)
        render_pass.set_bind_group(0, self.bind_group, [], 0, 999999)
        render_pass.set_vertex_buffer(0, position_buffer)

        if colour_buffer:
            render_pass.set_vertex_buffer(1, colour_buffer)

        # 4 vertices per quad for point rendering
        render_pass.draw(4, count)

class CustomShaderPipeline(BaseWebGPUPipeline):
    """A WebGPU pipeline that uses custom shader source provided by the user.

    This pipeline allows users to provide their own WGSL shader source code
    while handling the boilerplate for buffer management, uniform updates,
    and rendering setup.
    """

    def __init__(
        self,
        device: wgpu.GPUDevice,
        shader_source: str,
        vertex_formats: Optional[List[Union[str, wgpu.VertexFormat]]] = None,
        primitive_topology: wgpu.PrimitiveTopology = wgpu.PrimitiveTopology.triangle_list,
        texture_format: wgpu.TextureFormat = wgpu.TextureFormat.rgba8unorm,
        depth_format: wgpu.TextureFormat = wgpu.TextureFormat.depth24plus,
        msaa_sample_count: int = 4,
        uniform_struct_definition: Optional[str] = None,
        pipeline_label: str = "CustomShaderPipeline",
    ) -> None:
        """Initialize custom shader pipeline.

        Args:
            device: WebGPU device
            shader_source: WGSL shader source code as string
            vertex_formats: List of vertex data formats (e.g., ["Vec3", "Vec3"] for position+colour)
            primitive_topology: Primitive topology for rendering
            texture_format: Colour attachment format
            depth_format: Depth attachment format
            msaa_sample_count: Number of MSAA samples
            uniform_struct_definition: Custom uniform struct definition (optional)
            pipeline_label: Label for debugging
        """
        self._shader_source = shader_source
        self._vertex_formats = vertex_formats or ["Vec3"]
        self._primitive_topology = primitive_topology
        self._uniform_struct_definition = uniform_struct_definition
        self._pipeline_label = pipeline_label

        # Calculate stride from vertex formats
        self._stride = sum(
            NGLToWebGPU.stride_from_type(fmt) for fmt in self._vertex_formats
        )

        # Initialize buffers storage
        self.vertex_buffers: Dict[int, wgpu.GPUBuffer] = {}
        self.num_vertices = 0

        # Call parent constructor after setting up our attributes
        super().__init__(
            device=device,
            texture_format=texture_format,
            depth_format=depth_format,
            msaa_sample_count=msaa_sample_count,
            data_type=self._vertex_formats[0],  # Use first format as default
            stride=self._stride,
        )

    def get_dtype(self) -> np.dtype:
        """Get the numpy dtype for the uniform buffer structure."""
        if self._uniform_struct_definition:
            # For custom uniforms, we need to parse the struct or use a default
            # For now, use a basic MVP + colour structure
            return np.dtype(
                [
                    ("MVP", np.float32, (4, 4)),
                    ("colour", np.float32, 4),
                ]
            )
        else:
            # Default uniform structure
            return np.dtype(
                [
                    ("MVP", np.float32, (4, 4)),
                    ("colour", np.float32, 4),
                ]
            )

    def _get_shader_code(self) -> str:
        """Return the custom shader source."""
        return self._shader_source

    def _get_vertex_buffer_layouts(self) -> List[Dict[str, any]]:
        """Get vertex buffer layouts based on provided vertex formats."""
        if len(self._vertex_formats) == 1:
            # Single interleaved buffer
            return [
                {
                    "array_stride": self._stride,
                    "step_mode": "vertex",
                    "attributes": [
                        {
                            "format": NGLToWebGPU.vertex_format(
                                self._vertex_formats[0]
                            ),
                            "offset": 0,
                            "shader_location": 0,
                        },
                    ],
                },
            ]
        else:
            # Multiple separate buffers or interleaved with multiple attributes
            layouts = []
            current_offset = 0

            for i, fmt in enumerate(self._vertex_formats):
                stride = NGLToWebGPU.stride_from_type(fmt)
                layouts.append(
                    {
                        "array_stride": stride,
                        "step_mode": "vertex",
                        "attributes": [
                            {
                                "format": NGLToWebGPU.vertex_format(fmt),
                                "offset": 0,
                                "shader_location": i,
                            },
                        ],
                    }
                )
                current_offset += stride

            return layouts

    def _get_primitive_topology(self) -> wgpu.PrimitiveTopology:
        """Return the primitive topology."""
        return self._primitive_topology

    def _set_default_uniforms(self) -> None:
        """Set default uniform values."""
        self.uniform_data["MVP"] = np.eye(4, dtype=np.float32)
        self.uniform_data["colour"] = np.array([1.0, 1.0, 1.0, 1.0], dtype=np.float32)

    def _get_pipeline_label(self) -> str:
        """Return the pipeline label."""
        return self._pipeline_label

    def set_data(
        self,
        positions: Optional[np.ndarray] = None,
        colours: Optional[np.ndarray] = None,
        interleaved_data: Optional[np.ndarray] = None,
        **kwargs: Any,
    ) -> None:
        """Set vertex data for rendering.

        Args:
            positions: Vertex position data (N, 3)
            colours: Vertex colour data (N, 3) or (N, 4)
            interleaved_data: Pre-interleaved vertex data
            **kwargs: Additional vertex data arrays (e.g., velocities, life, initial_position)
        """
        if interleaved_data is not None:
            # Use pre-interleaved data
            buffer = self._create_or_update_buffer(
                self.vertex_buffers.get(0),
                interleaved_data,
                wgpu.BufferUsage.VERTEX | wgpu.BufferUsage.COPY_DST,
                f"{self._pipeline_label}_vertex_buffer_0",
            )
            self.vertex_buffers[0] = buffer[0]
            self.num_vertices = len(interleaved_data)
        else:
            # Handle separate vertex data arrays
            binding = 0

            if positions is not None:
                buffer = self._create_or_update_buffer(
                    self.vertex_buffers.get(binding),
                    positions,
                    wgpu.BufferUsage.VERTEX | wgpu.BufferUsage.COPY_DST,
                    f"{self._pipeline_label}_vertex_buffer_{binding}",
                )
                self.vertex_buffers[binding] = buffer[0]
                self.num_vertices = len(positions)
                binding += 1

            if colours is not None:
                buffer = self._create_or_update_buffer(
                    self.vertex_buffers.get(binding),
                    colours,
                    wgpu.BufferUsage.VERTEX | wgpu.BufferUsage.COPY_DST,
                    f"{self._pipeline_label}_vertex_buffer_{binding}",
                )
                self.vertex_buffers[binding] = buffer[0]
                binding += 1

            # Handle additional vertex attributes (velocities, life, initial_position, etc.)
            for attr_name, attr_data in kwargs.items():
                if attr_data is not None:
                    buffer = self._create_or_update_buffer(
                        self.vertex_buffers.get(binding),
                        attr_data,
                        wgpu.BufferUsage.VERTEX | wgpu.BufferUsage.COPY_DST,
                        f"{self._pipeline_label}_vertex_buffer_{binding}",
                    )
                    self.vertex_buffers[binding] = buffer[0]
                    binding += 1

    def update_uniforms(self, **kwargs: Any) -> None:
        """Update uniform buffer values.

        Args:
            **kwargs: Uniform values to update (e.g., mvp=matrix, colour=array)
        """
        if "mvp" in kwargs:
            self.uniform_data["MVP"] = kwargs["mvp"]

        if "colour" in kwargs:
            colour = kwargs["colour"]
            if len(colour) == 3:
                self.uniform_data["colour"] = np.array([*colour, 1.0], dtype=np.float32)
            else:
                self.uniform_data["colour"] = np.array(colour, dtype=np.float32)

        # Update the GPU buffer
        if self.uniform_buffer:
            self.device.queue.write_buffer(
                self.uniform_buffer, 0, self.uniform_data.tobytes()
            )

    def render(self, render_pass: wgpu.GPURenderPassEncoder, **kwargs: Any) -> None:
        """Render using this pipeline.

        Args:
            render_pass: Active render pass encoder
            **kwargs: Additional render parameters
        """
        if self.pipeline is None or self.num_vertices == 0:
            return

        render_pass.set_pipeline(self.pipeline)
        render_pass.set_bind_group(0, self.bind_group, [], 0, 999999)

        # Set vertex buffers
        for binding, buffer in self.vertex_buffers.items():
            render_pass.set_vertex_buffer(binding, buffer)

        # Draw
        render_pass.draw(self.num_vertices)

    @classmethod
    def from_file(
        cls, device: wgpu.GPUDevice, shader_file: str, **kwargs: Any
    ) -> "CustomShaderPipeline":
        """Create a CustomShaderPipeline from a WGSL file.

        Args:
            device: WebGPU device
            shader_file: Path to WGSL shader file
            **kwargs: Additional arguments to pass to constructor

        Returns:
            CustomShaderPipeline instance
        """
        if not os.path.exists(shader_file):
            raise FileNotFoundError(f"Shader file not found: {shader_file}")

        with open(shader_file, "r", encoding="utf-8") as f:
            shader_source = f.read()

        # Use filename as default pipeline label if not provided
        if "pipeline_label" not in kwargs:
            kwargs["pipeline_label"] = f"Custom_{os.path.basename(shader_file)}"

        return cls(device, shader_source, **kwargs)