Source code for simvx.graphics.renderer.buffer_manager

"""BufferManager — owns the forward renderer's SSBOs and descriptor sets.

Extracted from Renderer so transform/material/light/shadow/joint
buffers and the Forward+ tile-culling placeholders live in one place. The
descriptor set layout and cubemap placeholder are owned here too so the
renderer can swap the IBL cubemap in via ``write_cubemap_descriptor``.
"""

import logging
from typing import Any

import numpy as np
import vulkan as vk

from .._types import LIGHT_DTYPE, MATERIAL_DTYPE, TRANSFORM_DTYPE
from ..gpu.descriptors import (
    allocate_descriptor_set,
    create_descriptor_pool,
    create_ssbo_layout,
    write_image_descriptor,
    write_ssbo_descriptor,
)
from ..gpu.memory import create_buffer, upload_numpy

__all__ = ["BufferManager", "SHADOW_DATA_SIZE"]

log = logging.getLogger(__name__)

# Shadow SSBO total size — must match the ShadowBuffer struct in cube_textured.frag.
# Layout: cascade_vps[3](192) + cascade_splits(16) + flags/indices(32) +
#         point_light_pos_range(16) + spot_vp(64) + spot_light_pos_range(16) + ambient_colour(16) = 352
SHADOW_DATA_SIZE = 352

# Default ambient colour (cool grey fill) written at offset 336 when no WorldEnvironment overrides it.
_DEFAULT_AMBIENT = np.array([0.15, 0.15, 0.2, 1.0], dtype=np.float32)

_HOST_FLAGS = vk.VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | vk.VK_MEMORY_PROPERTY_HOST_COHERENT_BIT
_SSBO_USAGE = vk.VK_BUFFER_USAGE_STORAGE_BUFFER_BIT



class BufferManager:
    """Owns the renderer's SSBOs and descriptor sets.

    Main descriptor set (``ssbo_set``) exposes seven bindings:
      0: transforms, 1: materials, 2: lights, 3: shadow,
      4: IBL cubemap sampler, 5: tile light indices, 6: tile info.

    Joint descriptor set (``joint_set``) is set 2 binding 0 for skinned meshes.
    """

    def __init__(self, engine: Any, max_objects: int, max_materials: int = 1024,
                 max_lights: int = 256, max_joints: int = 256) -> None:
        self._engine = engine
        self.max_objects = max_objects
        self.max_materials = max_materials
        self.max_lights = max_lights
        self.max_joints = max_joints

        # SSBO resources
        self.transform_buf: Any = None
        self.transform_mem: Any = None
        self.material_buf: Any = None
        self.material_mem: Any = None
        self.light_buf: Any = None
        self.light_mem: Any = None
        self.shadow_buf: Any = None
        self.shadow_mem: Any = None
        self.tile_light_idx_buf: Any = None
        self.tile_light_idx_mem: Any = None
        self.tile_info_buf: Any = None
        self.tile_info_mem: Any = None
        self.joint_buf: Any = None
        self.joint_mem: Any = None

        # Descriptors
        self.ssbo_layout: Any = None
        self.ssbo_pool: Any = None
        self.ssbo_set: Any = None
        self.joint_layout: Any = None
        self.joint_pool: Any = None
        self.joint_set: Any = None

        # IBL cubemap placeholder (owned for lifetime of manager)
        self.placeholder_cubemap_view: Any = None
        self.placeholder_cubemap_sampler: Any = None
        self.placeholder_cubemap_img: Any = None
        self.placeholder_cubemap_mem: Any = None

        # Dirty-tracking: skip redundant GPU uploads when data hasn't changed
        self._materials_hash: int = 0
        self._lights_hash: int = 0

    # ------------------------------------------------------------------ setup



    def setup(self) -> None:
        """Allocate all SSBOs and descriptor sets."""
        e = self._engine
        device = e.ctx.device
        phys = e.ctx.physical_device

        transform_size = self.max_objects * TRANSFORM_DTYPE.itemsize
        material_size = self.max_materials * MATERIAL_DTYPE.itemsize
        light_size = self.max_lights * LIGHT_DTYPE.itemsize
        joint_buf_size = self.max_joints * 64  # mat4 = 64 bytes

        self.transform_buf, self.transform_mem = create_buffer(device, phys, transform_size, _SSBO_USAGE, _HOST_FLAGS)
        self.material_buf, self.material_mem = create_buffer(device, phys, material_size, _SSBO_USAGE, _HOST_FLAGS)
        self.light_buf, self.light_mem = create_buffer(device, phys, light_size, _SSBO_USAGE, _HOST_FLAGS)
        self.shadow_buf, self.shadow_mem = create_buffer(device, phys, SHADOW_DATA_SIZE, _SSBO_USAGE, _HOST_FLAGS)
        self.tile_light_idx_buf, self.tile_light_idx_mem = create_buffer(device, phys, 16, _SSBO_USAGE, _HOST_FLAGS)
        self.tile_info_buf, self.tile_info_mem = create_buffer(device, phys, 16, _SSBO_USAGE, _HOST_FLAGS)
        self.joint_buf, self.joint_mem = create_buffer(device, phys, joint_buf_size, _SSBO_USAGE, _HOST_FLAGS)

        # Main SSBO set: 4 SSBOs + 1 cubemap sampler + 2 trailing SSBOs
        self.ssbo_layout = create_ssbo_layout(device, binding_count=4, extra_samplers=1, trailing_ssbos=2)
        self.ssbo_pool = create_descriptor_pool(device, max_sets=1, extra_samplers=1, ssbo_count=6)
        self.ssbo_set = allocate_descriptor_set(device, self.ssbo_pool, self.ssbo_layout)
        write_ssbo_descriptor(device, self.ssbo_set, 0, self.transform_buf, transform_size)
        write_ssbo_descriptor(device, self.ssbo_set, 1, self.material_buf, material_size)
        write_ssbo_descriptor(device, self.ssbo_set, 2, self.light_buf, light_size)
        write_ssbo_descriptor(device, self.ssbo_set, 3, self.shadow_buf, SHADOW_DATA_SIZE)
        write_ssbo_descriptor(device, self.ssbo_set, 5, self.tile_light_idx_buf, 16)
        write_ssbo_descriptor(device, self.ssbo_set, 6, self.tile_info_buf, 16)

        # Joint SSBO set (set 2, binding 0)
        self.joint_layout = create_ssbo_layout(device, binding_count=1)
        self.joint_pool = create_descriptor_pool(device, max_sets=2)
        self.joint_set = allocate_descriptor_set(device, self.joint_pool, self.joint_layout)
        write_ssbo_descriptor(device, self.joint_set, 0, self.joint_buf, joint_buf_size)

        # Shadow SSBO defaults: no-shadow sentinels + ambient colour
        init_shadow = np.zeros(SHADOW_DATA_SIZE, dtype=np.uint8)
        sentinel = np.array([0xFF, 0xFF, 0xFF, 0xFF], dtype=np.uint8)
        init_shadow[208:212] = sentinel
        init_shadow[220:224] = sentinel
        init_shadow[224:228] = sentinel
        init_shadow[336:352] = _DEFAULT_AMBIENT.view(np.uint8)
        upload_numpy(device, self.shadow_mem, init_shadow)

        # IBL cubemap placeholder (replaced by Renderer.set_skybox)
        from ..assets.cubemap_loader import load_cubemap
        (
            self.placeholder_cubemap_view, self.placeholder_cubemap_sampler,
            self.placeholder_cubemap_img, self.placeholder_cubemap_mem,
        ) = load_cubemap(device, phys, e.ctx.graphics_queue, e.ctx.command_pool, colour=(0.0, 0.0, 0.0))
        write_image_descriptor(
            device, self.ssbo_set, 4, self.placeholder_cubemap_view, self.placeholder_cubemap_sampler
        )


    # ---------------------------------------------------------------- uploads



    def upload_transforms(self, instances: list) -> None:
        """Upload all instance transforms + normal matrices + material ids to the SSBO."""
        if not instances:
            return
        count = min(len(instances), self.max_objects)
        if len(instances) > self.max_objects:
            log.warning(
                "Instance count (%d) exceeds max_objects (%d), clamping", len(instances), self.max_objects
            )
        instances = instances[:count]

        model_mats = np.empty((count, 4, 4), dtype=np.float32)
        mat_ids = np.empty(count, dtype=np.uint32)
        for i, (_mh, xform, mid, _vp) in enumerate(instances):
            model_mats[i] = xform if xform.shape == (4, 4) else xform.T
            mat_ids[i] = mid

        model_mats_T = np.ascontiguousarray(model_mats.transpose(0, 2, 1))

        # Normal matrix: transpose(inverse(M3x3)), packed into mat4 column-major.
        m3x3 = model_mats[:, :3, :3]
        try:
            inv3x3 = np.linalg.inv(m3x3)
        except np.linalg.LinAlgError:
            inv3x3 = np.empty_like(m3x3)
            for j in range(count):
                try:
                    inv3x3[j] = np.linalg.inv(m3x3[j])
                except np.linalg.LinAlgError:
                    inv3x3[j] = m3x3[j].T
        normal4x4 = np.zeros((count, 4, 4), dtype=np.float32)
        normal4x4[:, 3, 3] = 1.0
        normal4x4[:, :3, :3] = inv3x3

        transforms = np.zeros(count, dtype=TRANSFORM_DTYPE)
        transforms["model"] = model_mats_T
        transforms["normal_mat"] = normal4x4
        transforms["material_index"] = mat_ids
        upload_numpy(self._engine.ctx.device, self.transform_mem, transforms)




    def set_materials(self, materials: np.ndarray) -> np.ndarray:
        """Upload material array. Returns the (possibly clamped) array stored."""
        if len(materials) > self.max_materials:
            log.warning("Material count (%d) exceeds max (%d), clamping", len(materials), self.max_materials)
            materials = materials[: self.max_materials]
        if self.material_mem:
            h = hash(materials.tobytes())
            if h != self._materials_hash:
                self._materials_hash = h
                upload_numpy(self._engine.ctx.device, self.material_mem, materials)
        return materials




    def set_lights(self, lights: np.ndarray) -> None:
        """Upload light array prefixed with the uint32 count (GLSL LightBuffer layout)."""
        if not self.light_mem:
            return
        h = hash(lights.tobytes())
        if h == self._lights_hash:
            return
        self._lights_hash = h
        count = np.array([len(lights)], dtype=np.uint32)
        padding = np.zeros(3, dtype=np.uint32)
        header = np.concatenate([count, padding])
        buf = np.concatenate([header.view(np.uint8), lights.view(np.uint8)])
        upload_numpy(self._engine.ctx.device, self.light_mem, buf)




    def set_hdr_flag(self, enabled: bool) -> None:
        """Toggle ``hdr_output`` (byte offset 216) in the shadow SSBO."""
        flag = np.array([1 if enabled else 0], dtype=np.uint32)
        shadow_data = np.zeros(SHADOW_DATA_SIZE, dtype=np.uint8)
        ffi = vk.ffi
        device = self._engine.ctx.device
        src = vk.vkMapMemory(device, self.shadow_mem, 0, SHADOW_DATA_SIZE, 0)
        ffi.memmove(ffi.cast("void*", shadow_data.ctypes.data), src, SHADOW_DATA_SIZE)
        vk.vkUnmapMemory(device, self.shadow_mem)
        shadow_data[216:220] = flag.view(np.uint8)
        upload_numpy(device, self.shadow_mem, shadow_data)




    def write_shadow_data(self, shadow_data: np.ndarray) -> None:
        """Upload raw shadow SSBO bytes (used by shadow passes + IBL-only fallback)."""
        upload_numpy(self._engine.ctx.device, self.shadow_mem, shadow_data)




    def write_cubemap_descriptor(self, view: Any, sampler: Any) -> None:
        """Bind a cubemap view+sampler to the IBL slot (binding 4)."""
        write_image_descriptor(self._engine.ctx.device, self.ssbo_set, 4, view, sampler)


    # ---------------------------------------------------------------- cleanup



    def cleanup(self) -> None:
        """Destroy all buffers, descriptor pools/layouts, and placeholder cubemap."""
        device = self._engine.ctx.device

        for buf, mem in (
            (self.joint_buf, self.joint_mem),
            (self.transform_buf, self.transform_mem),
            (self.material_buf, self.material_mem),
            (self.light_buf, self.light_mem),
            (self.shadow_buf, self.shadow_mem),
            (self.tile_light_idx_buf, self.tile_light_idx_mem),
            (self.tile_info_buf, self.tile_info_mem),
        ):
            if buf:
                vk.vkDestroyBuffer(device, buf, None)
            if mem:
                vk.vkFreeMemory(device, mem, None)

        for layout in (self.joint_layout, self.ssbo_layout):
            if layout:
                vk.vkDestroyDescriptorSetLayout(device, layout, None)
        for pool in (self.joint_pool, self.ssbo_pool):
            if pool:
                vk.vkDestroyDescriptorPool(device, pool, None)

        if self.placeholder_cubemap_sampler:
            vk.vkDestroySampler(device, self.placeholder_cubemap_sampler, None)
        if self.placeholder_cubemap_view:
            vk.vkDestroyImageView(device, self.placeholder_cubemap_view, None)
        if self.placeholder_cubemap_img:
            vk.vkDestroyImage(device, self.placeholder_cubemap_img, None)
        if self.placeholder_cubemap_mem:
            vk.vkFreeMemory(device, self.placeholder_cubemap_mem, None)