hydra-rs 0.0.4

#include "hydra_bridge.h"

#include <pxr/base/gf/frustum.h>
#include <pxr/base/gf/half.h>
#include <pxr/base/gf/vec3d.h>
#include <pxr/base/gf/vec4d.h>
#include <pxr/base/vt/value.h>
#include <pxr/imaging/hd/aov.h>
#include <pxr/imaging/hd/driver.h>
#include <pxr/imaging/hd/renderBuffer.h>
#include <pxr/imaging/hd/rendererPluginRegistry.h>
#include <pxr/imaging/hd/types.h>
#include <pxr/imaging/hdx/tokens.h>
#include <pxr/imaging/hf/pluginDesc.h>
#include <pxr/imaging/hgi/tokens.h>
#include <pxr/usd/sdf/assetPath.h>
#include <pxr/usd/sdf/layer.h>
#include <pxr/usd/usd/editContext.h>
#include <pxr/usd/usd/primRange.h>
#include <pxr/usd/usd/timeCode.h>
#include <pxr/usd/usdGeom/mesh.h>
#include <pxr/usd/usdGeom/xformable.h>
#include <pxr/usd/usdGeom/xformCommonAPI.h>
#include <pxr/base/gf/rotation.h>
#include <pxr/usd/usdLux/distantLight.h>
#include <pxr/usd/usdLux/domeLight.h>
#include <pxr/usd/usdLux/shapingAPI.h>
#include <pxr/usd/usdLux/sphereLight.h>
#include <pxr/usd/usdShade/connectableAPI.h>
#include <pxr/usd/usdShade/material.h>
#include <pxr/usd/usdShade/materialBindingAPI.h>
#include <pxr/usd/usdShade/shader.h>
#include <pxr/usdImaging/usdImagingGL/renderParams.h>

#include <algorithm>
#include <cstring>
#include <functional>
#include <stdexcept>

namespace hydra_rs {

std::unique_ptr<SceneIndex> populate_from_path(rust::Str usd_path) {
    std::string p(usd_path.data(), usd_path.size());
    pxr::UsdStageRefPtr stage = pxr::UsdStage::Open(p);
    if (!stage) {
        throw std::runtime_error("UsdStage::Open returned null for: " + p);
    }

    auto si = pxr::UsdImagingStageSceneIndex::New();
    si->SetStage(stage);
    si->SetTime(pxr::UsdTimeCode::Default());

    auto wrapped = std::make_unique<SceneIndex>();
    wrapped->stage_owner = stage;
    wrapped->scene_index = si;
    return wrapped;
}

rust::String SceneIndex::stage_root() const {
    if (!stage_owner) return rust::String();
    auto root = stage_owner->GetRootLayer();
    if (!root) return rust::String();
    return rust::String(root->GetIdentifier());
}

size_t SceneIndex::prim_count() const {
    if (!scene_index) return 0;
    size_t count = 0;
    std::function<void(const pxr::SdfPath&)> walk = [&](const pxr::SdfPath& p) {
        ++count;
        for (const auto& child : scene_index->GetChildPrimPaths(p)) {
            walk(child);
        }
    };
    walk(pxr::SdfPath::AbsoluteRootPath());
    return count;
}

std::unique_ptr<std::vector<std::string>> SceneIndex::prim_paths() const {
    auto vec = std::make_unique<std::vector<std::string>>();
    if (!scene_index) return vec;
    std::function<void(const pxr::SdfPath&)> walk = [&](const pxr::SdfPath& p) {
        vec->push_back(p.GetString());
        for (const auto& child : scene_index->GetChildPrimPaths(p)) {
            walk(child);
        }
    };
    walk(pxr::SdfPath::AbsoluteRootPath());
    return vec;
}

std::unique_ptr<std::vector<std::string>> list_render_delegate_ids() {
    auto vec = std::make_unique<std::vector<std::string>>();
    auto& reg = pxr::HdRendererPluginRegistry::GetInstance();
    pxr::HfPluginDescVector descs;
    reg.GetPluginDescs(&descs);
    vec->reserve(descs.size());
    for (const auto& d : descs) {
        vec->push_back(d.id.GetString());
    }
    return vec;
}

namespace {

// Slice<const float> with 16 row-major entries into a GfMatrix4d.
pxr::GfMatrix4d matrix_from_slice(rust::Slice<const float> s) {
    if (s.size() != 16) {
        throw std::runtime_error(
            "matrix slice must have exactly 16 elements (row-major 4x4)");
    }
    double m[4][4];
    for (size_t r = 0; r < 4; ++r) {
        for (size_t c = 0; c < 4; ++c) {
            m[r][c] = static_cast<double>(s[r * 4 + c]);
        }
    }
    return pxr::GfMatrix4d(m);
}

void apply_camera_default(pxr::GfMatrix4d& view, pxr::GfMatrix4d& proj,
                          uint32_t w, uint32_t h) {
    view.SetLookAt(
        pxr::GfVec3d(5.0, 5.0, 5.0),
        pxr::GfVec3d(0.0, 0.0, 0.0),
        pxr::GfVec3d(0.0, 1.0, 0.0));
    pxr::GfFrustum frustum;
    const double aspect = static_cast<double>(w) / static_cast<double>(h);
    frustum.SetPerspective(45.0, aspect, 0.1, 1000.0);
    proj = frustum.ComputeProjectionMatrix();
}

}  // namespace

std::unique_ptr<Renderer> create_renderer(rust::Str usd_path,
                                          rust::Str render_delegate_id) {
    auto r = std::make_unique<Renderer>();

    std::string p(usd_path.data(), usd_path.size());
    r->stage = pxr::UsdStage::Open(p);
    if (!r->stage) {
        throw std::runtime_error("UsdStage::Open returned null for: " + p);
    }

    r->hgi = pxr::Hgi::CreatePlatformDefaultHgi();
    if (!r->hgi) {
        throw std::runtime_error(
            "Hgi::CreatePlatformDefaultHgi returned null — no usable GPU "
            "backend. On macOS, expected HgiMetal; on Linux, HgiGL/Vulkan.");
    }

    pxr::HdDriver hd_driver;
    hd_driver.name = pxr::HgiTokens->renderDriver;
    hd_driver.driver = pxr::VtValue(r->hgi.get());

    r->engine = std::make_unique<pxr::UsdImagingGLEngine>(hd_driver);
    r->engine->SetEnablePresentation(false);

    std::string delegate_str(render_delegate_id.data(), render_delegate_id.size());
    if (!delegate_str.empty()) {
        if (!r->engine->SetRendererPlugin(pxr::TfToken(delegate_str))) {
            throw std::runtime_error("SetRendererPlugin failed: " + delegate_str);
        }
        // SetRendererPlugin re-creates the engine's task graph, which
        // re-enables HdxPresentTask (the GL-interop blit task). Same
        // reason as Renderer::set_renderer_plugin — re-assert here so
        // delegates picked at construction get the same headless
        // safety as delegates swapped in later.
        r->engine->SetEnablePresentation(false);
    }

    // Default camera: a 45 degree FOV looking at origin from (5, 5, 5).
    apply_camera_default(r->view_matrix, r->proj_matrix, r->width, r->height);

    // Default material is white diffuse with low ambient — Storm reads this
    // when the scene has no UsdShade material bound.
    r->material.SetAmbient(pxr::GfVec4f(0.2f, 0.2f, 0.2f, 1.0f));
    r->material.SetDiffuse(pxr::GfVec4f(1.0f, 1.0f, 1.0f, 1.0f));
    r->material.SetSpecular(pxr::GfVec4f(0.5f, 0.5f, 0.5f, 1.0f));
    r->material.SetShininess(32.0f);

    return r;
}

void Renderer::set_size(uint32_t w, uint32_t h) {
    width = w;
    height = h;
    // Recompute the default-camera projection if the user hasn't set their own
    // view yet — keeps the aspect ratio sane on resize.
    pxr::GfFrustum frustum;
    const double aspect = static_cast<double>(w) / static_cast<double>(h);
    frustum.SetPerspective(45.0, aspect, 0.1, 1000.0);
    proj_matrix = frustum.ComputeProjectionMatrix();
}

void Renderer::set_camera_matrices(rust::Slice<const float> view,
                                   rust::Slice<const float> projection) {
    view_matrix = matrix_from_slice(view);
    proj_matrix = matrix_from_slice(projection);
}

void Renderer::set_time(double time) {
    frame = time;
    default_frame = false;
}

void Renderer::use_default_time() {
    default_frame = true;
}

void Renderer::clear_lights() {
    explicit_lights.clear();
    use_default_lighting = false;
}

void Renderer::use_default_light() {
    explicit_lights.clear();
    use_default_lighting = true;
}

void Renderer::add_distant_light(float dx, float dy, float dz,
                                 float r, float g, float b, float intensity) {
    pxr::GlfSimpleLight light;
    // GlfSimpleLight position with w == 0 is a directional light with the
    // direction encoded as the x/y/z components of position.
    light.SetPosition(pxr::GfVec4f(dx, dy, dz, 0.0f));
    light.SetDiffuse(pxr::GfVec4f(r * intensity, g * intensity, b * intensity, 1.0f));
    light.SetSpecular(pxr::GfVec4f(r * intensity, g * intensity, b * intensity, 1.0f));
    light.SetAmbient(pxr::GfVec4f(0.0f, 0.0f, 0.0f, 1.0f));
    explicit_lights.push_back(light);
    use_default_lighting = false;
}

void Renderer::add_positional_light(float px, float py, float pz,
                                    float r, float g, float b, float intensity) {
    pxr::GlfSimpleLight light;
    light.SetPosition(pxr::GfVec4f(px, py, pz, 1.0f));  // w == 1 is positional
    light.SetDiffuse(pxr::GfVec4f(r * intensity, g * intensity, b * intensity, 1.0f));
    light.SetSpecular(pxr::GfVec4f(r * intensity, g * intensity, b * intensity, 1.0f));
    light.SetAmbient(pxr::GfVec4f(0.0f, 0.0f, 0.0f, 1.0f));
    explicit_lights.push_back(light);
    use_default_lighting = false;
}

void Renderer::set_clear_color(float r, float g, float b, float a) {
    clear_color = pxr::GfVec4f(r, g, b, a);
}

void Renderer::set_scene_ambient(float r, float g, float b, float a) {
    scene_ambient = pxr::GfVec4f(r, g, b, a);
}

namespace {

// Stage path for the dome we author into the session layer. Leading
// underscore mirrors the convention USD uses for editor-managed prims
// (UsdSkel rest-pose helpers, Solaris stage variants) and keeps the
// dome out of the way if a consumer enumerates user-authored prims.
const pxr::SdfPath kDomeLightPath("/_hydraDomeLight");

}  // namespace

void Renderer::set_dome_light(rust::Str hdri_path,
                              float intensity,
                              float exposure,
                              float rotation_y_degrees) {
    if (!stage) {
        throw std::runtime_error(
            "Renderer::set_dome_light called before stage was opened");
    }
    // Author into the session layer rather than the root layer — the
    // dome is a viewport convenience, not a property of the asset, so
    // we never want it to follow the stage when the consumer saves
    // back out. The EditContext goes out of scope at the end of the
    // function and restores whatever edit target was set before.
    pxr::UsdEditContext ctx(stage, stage->GetSessionLayer());

    pxr::UsdLuxDomeLight dome =
        pxr::UsdLuxDomeLight::Define(stage, kDomeLightPath);
    if (!dome) {
        throw std::runtime_error(
            "UsdLuxDomeLight::Define returned an invalid prim at " +
            kDomeLightPath.GetString());
    }

    std::string p(hdri_path.data(), hdri_path.size());
    dome.CreateTextureFileAttr().Set(pxr::SdfAssetPath(p));
    dome.CreateIntensityAttr().Set(intensity);
    dome.CreateExposureAttr().Set(exposure);

    // forge-paint hands us equirectangular HDRIs; this is also the USD
    // default, but spell it out so a delegate's auto-detect can't pick
    // the wrong projection on an HDR with a square aspect ratio.
    dome.CreateTextureFormatAttr().Set(pxr::TfToken("latlong"));

    // UsdGeomXformCommonAPI gives us a single rotateY op we can set
    // by value on every call — Storm reads the local-to-world xform
    // of the dome to orient the env. Using the common-API instead of
    // raw xformOps keeps us from accidentally stacking rotation ops.
    pxr::UsdGeomXformCommonAPI xform_api(dome.GetPrim());
    xform_api.SetRotate(pxr::GfVec3f(0.0f, rotation_y_degrees, 0.0f));

    dome_light_active = true;
}

void Renderer::clear_dome_light() {
    if (!stage) return;
    pxr::UsdEditContext ctx(stage, stage->GetSessionLayer());
    stage->RemovePrim(kDomeLightPath);
    dome_light_active = false;
}

namespace {

// Where the painted UsdPreviewSurface gets authored. Keeping it
// under a single root prim makes `clear_painted_material()` a clean
// `RemovePrim` on the parent — no orphaned shader nodes left behind.
const pxr::SdfPath kPaintMaterialRoot("/_hydraPaintMaterial");
const pxr::TfToken kPaintMaterialBindingName("_hydraPaintMaterial");

// Wire a single `UsdUVTexture` into one of the UsdPreviewSurface
// inputs. `output_name` is the channel selector on the texture
// node ("rgb" for colour, "r" for scalar/single-channel reads).
// `input_type` lets us connect Float / Color3f / Normal3f inputs
// without each call site re-declaring the type.
void connect_texture_input(
    const pxr::UsdStageRefPtr& stage,
    const pxr::SdfPath& material_root,
    pxr::UsdShadeShader& surface,
    pxr::UsdShadeOutput& st_output,
    const std::string& asset_path,
    const char* tex_prim_name,
    const pxr::TfToken& surface_input,
    const pxr::SdfValueTypeName& surface_input_type,
    const pxr::TfToken& tex_output_name,
    const pxr::SdfValueTypeName& tex_output_type,
    // sourceColorType: "sRGB" for diffuse (gamma-encoded), "raw" for
    // linear data (roughness, metallic, normal). Drives the inverse
    // gamma curve UsdUVTexture applies on sample.
    const char* source_color_type)
{
    if (asset_path.empty()) return;

    pxr::SdfPath tex_path = material_root.AppendChild(pxr::TfToken(tex_prim_name));
    auto tex = pxr::UsdShadeShader::Define(stage, tex_path);
    tex.CreateIdAttr().Set(pxr::TfToken("UsdUVTexture"));
    tex.CreateInput(pxr::TfToken("file"), pxr::SdfValueTypeNames->Asset)
        .Set(pxr::SdfAssetPath(asset_path));
    tex.CreateInput(pxr::TfToken("sourceColorSpace"), pxr::SdfValueTypeNames->Token)
        .Set(pxr::TfToken(source_color_type));
    tex.CreateInput(pxr::TfToken("st"), pxr::SdfValueTypeNames->Float2)
        .ConnectToSource(st_output);

    auto tex_out = tex.CreateOutput(tex_output_name, tex_output_type);
    surface.CreateInput(surface_input, surface_input_type).ConnectToSource(tex_out);
}

}  // namespace

void Renderer::set_painted_material(rust::Str base_color_asset_path,
                                    rust::Str roughness_asset_path,
                                    rust::Str metallic_asset_path,
                                    rust::Str normal_asset_path)
{
    if (!stage) {
        throw std::runtime_error(
            "Renderer::set_painted_material called before stage was opened");
    }

    pxr::UsdEditContext ctx(stage, stage->GetSessionLayer());

    // Wipe any prior painted material so re-syncing doesn't stack
    // shader nodes from previous calls. Cheap: a single RemovePrim
    // takes the whole subtree (material + shaders + texture nodes).
    if (stage->GetPrimAtPath(kPaintMaterialRoot)) {
        stage->RemovePrim(kPaintMaterialRoot);
    }

    // Material container + UsdPreviewSurface + a shared
    // UsdPrimvarReader_float2 for the `st` primvar. Every texture
    // node connects its `inputs:st` to this reader, so the four
    // texture lookups all use the same UV set.
    auto mat = pxr::UsdShadeMaterial::Define(stage, kPaintMaterialRoot);

    auto surface = pxr::UsdShadeShader::Define(
        stage, kPaintMaterialRoot.AppendChild(pxr::TfToken("Surface")));
    surface.CreateIdAttr().Set(pxr::TfToken("UsdPreviewSurface"));
    auto surface_out = surface.CreateOutput(
        pxr::TfToken("surface"), pxr::SdfValueTypeNames->Token);

    auto st_reader = pxr::UsdShadeShader::Define(
        stage, kPaintMaterialRoot.AppendChild(pxr::TfToken("StReader")));
    st_reader.CreateIdAttr().Set(pxr::TfToken("UsdPrimvarReader_float2"));
    st_reader.CreateInput(pxr::TfToken("varname"), pxr::SdfValueTypeNames->Token)
        .Set(pxr::TfToken("st"));
    auto st_out = st_reader.CreateOutput(
        pxr::TfToken("result"), pxr::SdfValueTypeNames->Float2);

    auto str_to_string = [](rust::Str s) {
        return std::string(s.data(), s.size());
    };
    std::string bc = str_to_string(base_color_asset_path);
    std::string ro = str_to_string(roughness_asset_path);
    std::string me = str_to_string(metallic_asset_path);
    std::string nm = str_to_string(normal_asset_path);

    // diffuseColor reads gamma-encoded data → sRGB sourceColorSpace.
    connect_texture_input(
        stage, kPaintMaterialRoot, surface, st_out, bc,
        "BaseColorTex",
        pxr::TfToken("diffuseColor"), pxr::SdfValueTypeNames->Color3f,
        pxr::TfToken("rgb"), pxr::SdfValueTypeNames->Float3,
        "sRGB");
    // Linear scalar reads — roughness/metallic ship as 8-bit data
    // but represent linear values, so sourceColorSpace = raw.
    connect_texture_input(
        stage, kPaintMaterialRoot, surface, st_out, ro,
        "RoughnessTex",
        pxr::TfToken("roughness"), pxr::SdfValueTypeNames->Float,
        pxr::TfToken("r"), pxr::SdfValueTypeNames->Float,
        "raw");
    connect_texture_input(
        stage, kPaintMaterialRoot, surface, st_out, me,
        "MetallicTex",
        pxr::TfToken("metallic"), pxr::SdfValueTypeNames->Float,
        pxr::TfToken("r"), pxr::SdfValueTypeNames->Float,
        "raw");
    connect_texture_input(
        stage, kPaintMaterialRoot, surface, st_out, nm,
        "NormalTex",
        pxr::TfToken("normal"), pxr::SdfValueTypeNames->Normal3f,
        pxr::TfToken("rgb"), pxr::SdfValueTypeNames->Float3,
        "raw");

    // Surface output of the material connects to the surface shader.
    mat.CreateSurfaceOutput().ConnectToSource(
        surface.ConnectableAPI(), pxr::TfToken("surface"));

    // Bind to every mesh in the stage. Use the `MaterialBindingAPI`
    // applied to each prim so the override sticks even if the asset
    // already had its own binding — UsdShade's bind strength rules
    // give the strongest binding priority, and a session-layer-
    // authored binding wins over root-layer authoring.
    pxr::UsdPrimRange range(stage->GetPseudoRoot());
    for (const auto& prim : range) {
        if (prim.IsA<pxr::UsdGeomMesh>()) {
            auto bind = pxr::UsdShadeMaterialBindingAPI::Apply(prim);
            bind.Bind(mat);
        }
    }
}

void Renderer::set_show_render(bool show) { show_render_purpose = show; }
void Renderer::set_show_proxy(bool show) { show_proxy_purpose = show; }
void Renderer::set_show_guides(bool show) { show_guides_purpose = show; }

namespace {

// Stage-path helper: `/_hydraLight<index>`. Same `_` prefix as the
// dome path so consumers walking authored prims can skip our bridge-
// owned scaffolding cleanly.
pxr::SdfPath user_light_path(uint32_t index) {
    return pxr::SdfPath("/_hydraLight" + std::to_string(index));
}

// Place a light prim at `position` with its local -Z aimed along
// `dir`. UsdLuxDistantLight and UsdLuxSphereLight both emit along
// their local -Z (the spot's cone is centered there), so this is
// the canonical "point this light here" helper.
//
// We use a single `TransformOp` carrying a full matrix instead of
// Euler decomposition into `UsdGeomXformCommonAPI::SetRotate(XYZ)` —
// `GfRotation::Decompose` returns angles in ZYX multiplication order
// (R = R(z) * R(y) * R(x)) but the common API applies them as
// `R = R(X) * R(Y) * R(Z)`, which silently shipped a wrong rotation
// for non-trivial directions. Building the matrix directly avoids
// the convention mismatch entirely.
//
// `dir` is normalised on entry — zero-length falls back to (0,-1,0).
// `position` only matters for spot lights; directional ignores it
// (passed as zero from the caller).
void place_light_at(const pxr::UsdPrim& prim,
                    pxr::GfVec3f dir,
                    pxr::GfVec3f position) {
    const float mag = dir.GetLength();
    if (mag < 1e-6f) {
        dir = pxr::GfVec3f(0.0f, -1.0f, 0.0f);
    } else {
        dir /= mag;
    }
    const pxr::GfRotation rot(pxr::GfVec3d(0.0, 0.0, -1.0),
                              pxr::GfVec3d(dir[0], dir[1], dir[2]));
    pxr::GfMatrix4d xform;
    xform.SetRotate(rot);
    // GfMatrix4d uses row-vector convention; `SetTranslateOnly`
    // patches the translation row without touching the rotation
    // sub-matrix, so the final matrix encodes (rotate) then
    // (translate) when applied as `v * M`. Local origin (0,0,0)
    // maps to world `position`; local -Z maps to world `dir`.
    xform.SetTranslateOnly(pxr::GfVec3d(position[0], position[1], position[2]));

    pxr::UsdGeomXformable xformable(prim);
    // Wipe any previously-authored ops so re-authoring at the same
    // prim path doesn't accumulate translate + rotate + transform
    // layered on top of each other.
    xformable.ClearXformOpOrder();
    pxr::UsdGeomXformOp op = xformable.AddTransformOp();
    op.Set(xform);
}

}  // namespace

namespace {
const pxr::SdfPath kExternalMaterialPath("/_hydraExternalMaterial");
const pxr::TfToken kExternalMatBindingName("_hydraExternalMaterial");
}  // namespace

void Renderer::set_external_material(rust::Str source_usd_path,
                                     rust::Str prim_path) {
    if (!stage) {
        throw std::runtime_error(
            "Renderer::set_external_material called before stage was opened");
    }
    const std::string src(source_usd_path.data(), source_usd_path.size());
    const std::string sub_prim(prim_path.data(), prim_path.size());

    pxr::UsdEditContext ctx(stage, stage->GetSessionLayer());

    // Wipe any prior external material binding so the new reference
    // doesn't accumulate. Empty source path = "just clear."
    if (stage->GetPrimAtPath(kExternalMaterialPath)) {
        stage->RemovePrim(kExternalMaterialPath);
    }
    if (src.empty()) {
        // Caller asked for a clear; unbind from every mesh and bail.
        pxr::UsdPrimRange range(stage->GetPseudoRoot());
        for (const auto& prim : range) {
            if (prim.IsA<pxr::UsdGeomMesh>()) {
                auto bind = pxr::UsdShadeMaterialBindingAPI(prim);
                if (bind) bind.UnbindAllBindings();
            }
        }
        return;
    }

    // Create a `Material` prim and add a reference to the external
    // source. `prim_path` may be empty — then the default prim of
    // the source layer is used. Most pipeline material libraries
    // mark their root material as the default prim, so this is the
    // common case.
    auto mat = pxr::UsdShadeMaterial::Define(stage, kExternalMaterialPath);
    pxr::SdfPath ref_prim_path = sub_prim.empty()
        ? pxr::SdfPath()
        : pxr::SdfPath(sub_prim);
    mat.GetPrim().GetReferences().AddReference(src, ref_prim_path);

    // Bind to every mesh — same pattern as the painted material.
    pxr::UsdPrimRange range(stage->GetPseudoRoot());
    for (const auto& prim : range) {
        if (prim.IsA<pxr::UsdGeomMesh>()) {
            auto bind = pxr::UsdShadeMaterialBindingAPI::Apply(prim);
            bind.Bind(mat);
        }
    }
}

void Renderer::clear_external_material() {
    if (!stage) return;
    pxr::UsdEditContext ctx(stage, stage->GetSessionLayer());
    pxr::UsdPrimRange range(stage->GetPseudoRoot());
    for (const auto& prim : range) {
        if (prim.IsA<pxr::UsdGeomMesh>()) {
            auto bind = pxr::UsdShadeMaterialBindingAPI(prim);
            if (bind) bind.UnbindAllBindings();
        }
    }
    if (stage->GetPrimAtPath(kExternalMaterialPath)) {
        stage->RemovePrim(kExternalMaterialPath);
    }
}

void Renderer::set_user_lights(rust::Slice<const float> data) {
    if (!stage) {
        throw std::runtime_error(
            "Renderer::set_user_lights called before stage was opened");
    }
    if (data.size() % 16 != 0) {
        throw std::runtime_error(
            "Renderer::set_user_lights: payload length must be a multiple of 16");
    }

    pxr::UsdEditContext ctx(stage, stage->GetSessionLayer());

    const uint32_t new_count = static_cast<uint32_t>(data.size() / 16);

    for (uint32_t i = 0; i < new_count; ++i) {
        const float* p = data.data() + i * 16;
        const pxr::GfVec3f direction(p[0], p[1], p[2]);
        const float type_tag = p[3];
        const pxr::GfVec3f position(p[4], p[5], p[6]);
        const float enabled = p[7];
        const pxr::GfVec3f color(p[8], p[9], p[10]);
        const float intensity = p[11];
        const float cos_inner = p[12];
        const float cos_outer = p[13];

        const pxr::SdfPath path = user_light_path(i);

        // Disabled lights take the form of an unauthored slot —
        // remove the prim if it was authored last time. Avoids
        // having UsdLux prims with intensity = 0 lingering in the
        // session layer (Storm/hdNSI still has to consider them).
        if (enabled < 0.5f) {
            if (stage->GetPrimAtPath(path)) {
                stage->RemovePrim(path);
            }
            continue;
        }

        const bool is_spot = type_tag > 0.5f;

        // Drop any prim of the wrong type at this slot before
        // (re-)defining — switching a directional → spot in-place
        // would otherwise leave stale attrs from the previous type.
        if (stage->GetPrimAtPath(path)) {
            const auto& existing = stage->GetPrimAtPath(path);
            const bool was_spot = existing.IsA<pxr::UsdLuxSphereLight>();
            if (was_spot != is_spot) {
                stage->RemovePrim(path);
            }
        }

        if (is_spot) {
            auto light = pxr::UsdLuxSphereLight::Define(stage, path);
            // Small radius → the light reads as a point source. A
            // larger radius would soften shadows but we're keeping
            // the wgpu side's spot definition (point + cone), so
            // mirror that on the Hydra side.
            light.CreateRadiusAttr().Set(0.01f);
            light.CreateColorAttr().Set(color);
            light.CreateIntensityAttr().Set(intensity);

            // Place at world position + aim the cone axis along the
            // travel direction. One matrix transform op handles both
            // translation and rotation without the Euler-order
            // mismatch the old common-API path had.
            place_light_at(light.GetPrim(), direction, position);

            // Shaping cone — `shaping:cone:angle` is in degrees from
            // axis to outer edge; `shaping:cone:softness` is the
            // 0..1 falloff. We get cos(inner) and cos(outer) from
            // the consumer; convert back to angles and derive
            // softness as (outer - inner) / outer.
            auto shaping = pxr::UsdLuxShapingAPI::Apply(light.GetPrim());
            const auto clamp_cos = [](float v) {
                if (v < -1.0f) return -1.0f;
                if (v > 1.0f) return 1.0f;
                return v;
            };
            const float inner_deg = std::acos(clamp_cos(cos_inner)) * 180.0f / 3.14159265f;
            const float outer_deg = std::acos(clamp_cos(cos_outer)) * 180.0f / 3.14159265f;
            const float angle = outer_deg;
            const float softness = (outer_deg > 1e-3f)
                ? ((outer_deg - inner_deg) / outer_deg)
                : 0.0f;
            shaping.CreateShapingConeAngleAttr().Set(angle);
            shaping.CreateShapingConeSoftnessAttr().Set(softness);
        } else {
            auto light = pxr::UsdLuxDistantLight::Define(stage, path);
            light.CreateColorAttr().Set(color);
            light.CreateIntensityAttr().Set(intensity);
            // DistantLight `inputs:angle` is the apparent solar
            // disc size — we don't expose it; UsdLux default
            // (0.53° solar disc) is fine for crisp shadows.
            // Distant lights ignore position (they're at infinity)
            // so pass zero; only orientation matters.
            place_light_at(light.GetPrim(), direction, pxr::GfVec3f(0.0f));
        }
    }

    // Remove any previously-authored slots that fell off the end of
    // the new payload (e.g. user deleted a light in the panel).
    for (uint32_t i = new_count; i < user_light_count; ++i) {
        const pxr::SdfPath path = user_light_path(i);
        if (stage->GetPrimAtPath(path)) {
            stage->RemovePrim(path);
        }
    }
    user_light_count = new_count;
}

void Renderer::clear_user_lights() {
    if (!stage) return;
    pxr::UsdEditContext ctx(stage, stage->GetSessionLayer());
    for (uint32_t i = 0; i < user_light_count; ++i) {
        const pxr::SdfPath path = user_light_path(i);
        if (stage->GetPrimAtPath(path)) {
            stage->RemovePrim(path);
        }
    }
    user_light_count = 0;
}

void Renderer::clear_painted_material() {
    if (!stage) return;
    pxr::UsdEditContext ctx(stage, stage->GetSessionLayer());

    // Unbind from every mesh first, so the meshes fall back to
    // whatever was authored in the root layer. `UnbindAllBindings`
    // removes all `material:binding` rels authored at the current
    // edit target (the session layer); root-layer bindings are
    // unaffected and re-take precedence.
    pxr::UsdPrimRange range(stage->GetPseudoRoot());
    for (const auto& prim : range) {
        if (prim.IsA<pxr::UsdGeomMesh>()) {
            auto bind = pxr::UsdShadeMaterialBindingAPI(prim);
            if (bind) bind.UnbindAllBindings();
        }
    }
    if (stage->GetPrimAtPath(kPaintMaterialRoot)) {
        stage->RemovePrim(kPaintMaterialRoot);
    }
}

rust::String Renderer::current_renderer() const {
    if (!engine) return rust::String();
    return rust::String(engine->GetCurrentRendererId().GetString());
}

bool Renderer::is_converged() const {
    if (!engine) return false;
    return engine->IsConverged();
}

bool Renderer::set_renderer_plugin(rust::Str plugin_id) const {
    if (!engine) return false;
    std::string id(plugin_id.data(), plugin_id.size());
    bool ok = engine->SetRendererPlugin(pxr::TfToken(id));
    if (ok) {
        // `SetRendererPlugin` rebuilds the engine's internal task
        // graph from scratch, which silently re-enables HdxPresentTask
        // — the GL-interop "blit to window" task we explicitly
        // disabled at construction. Storm survives because its HGI
        // output short-circuits the task; sampling delegates like
        // hdNSI route through HgiInteropMetal → GarchGLApiLoad →
        // glGetString, which segfaults headlessly on macOS (no
        // active GL context). Re-asserting the flag right after the
        // switch keeps the present task disabled for the new
        // delegate's task graph too.
        engine->SetEnablePresentation(false);
    }
    return ok;
}

std::unique_ptr<std::vector<uint8_t>> Renderer::render_color() const {
    if (!engine || !stage) {
        throw std::runtime_error("Renderer not initialized");
    }
    if (width == 0 || height == 0) {
        throw std::runtime_error("Renderer width and height must be > 0");
    }

    engine->SetCameraState(view_matrix, proj_matrix);
    engine->SetRenderViewport(pxr::GfVec4d(0.0, 0.0,
                                           static_cast<double>(width),
                                           static_cast<double>(height)));
    engine->SetRendererAov(pxr::HdAovTokens->color);

    // Lighting: explicit_lights wins; otherwise a sensible default light
    // when use_default_lighting is on; otherwise nothing (renders unlit).
    pxr::GlfSimpleLightVector lights = explicit_lights;
    if (lights.empty() && use_default_lighting) {
        pxr::GlfSimpleLight default_light;
        default_light.SetPosition(pxr::GfVec4f(2.5f, 4.0f, 5.0f, 1.0f));
        default_light.SetDiffuse(pxr::GfVec4f(1.0f, 1.0f, 1.0f, 1.0f));
        default_light.SetSpecular(pxr::GfVec4f(1.0f, 1.0f, 1.0f, 1.0f));
        default_light.SetAmbient(pxr::GfVec4f(0.0f, 0.0f, 0.0f, 1.0f));
        lights.push_back(default_light);
    }
    engine->SetLightingState(lights, material, scene_ambient);

    pxr::UsdImagingGLRenderParams params;
    params.frame = default_frame ? pxr::UsdTimeCode::Default()
                                  : pxr::UsdTimeCode(frame);
    params.complexity = 1.0f;
    const bool any_user_lux = dome_light_active || user_light_count > 0;
    params.enableLighting = !lights.empty()
                            || use_default_lighting
                            || any_user_lux;
    // Scene lights flip on once any UsdLux prim has been authored
    // by the bridge — dome via `set_dome_light` or analytic
    // distant/spot via `set_user_lights`. Without one of those the
    // bool stays off so callers that only ever use the legacy
    // GlfSimpleLight path don't have their existing behaviour
    // changed.
    params.enableSceneLights = any_user_lux;
    params.enableSceneMaterials = true;
    // Purpose filters mirror the consumer's per-frame toggles —
    // `set_show_render` / `_proxy` / `_guides`. Default-purpose prims
    // always draw; the three optional purposes opt in via these
    // flags. Pipeline assets that wrap their detail in
    // `Scope { purpose = "render" }` only show when render is on.
    params.showRender = show_render_purpose;
    params.showProxy = show_proxy_purpose;
    params.showGuides = show_guides_purpose;
    params.clearColor = clear_color;

    // Single Render() per call: rasterising delegates (Storm) finish
    // in one pass and `IsConverged()` flips to true immediately, so
    // there's nothing for a convergence loop to do; sampling
    // delegates (hdNSI, Arnold, Embree, Cycles, …) advance their
    // image by a budgeted amount and return — the caller is expected
    // to re-issue render_color() each frame to progressively converge
    // (which is exactly what an interactive viewport's continuous-
    // repaint loop already does).
    //
    // The previous implementation looped until IsConverged() or a
    // 256-iteration safety cap. That blocks the calling thread for
    // minutes on a path tracer with a high sample budget, which froze
    // forge-paint's UI when the dual-panel Hydra preview ran on
    // hdNSI. For one-shot headless renders that DO need a converged
    // image, the caller can drive the same loop in user code:
    //     while (!IsConverged()) renderer.render();
    // once `IsConverged` is exposed through the bridge.
    engine->Render(stage->GetPseudoRoot(), params);

    pxr::HdRenderBuffer* color = engine->GetAovRenderBuffer(pxr::HdAovTokens->color);
    if (!color) {
        throw std::runtime_error("UsdImagingGLEngine has no color render buffer");
    }
    color->Resolve();

    const uint32_t bw = color->GetWidth();
    const uint32_t bh = color->GetHeight();
    const pxr::HdFormat fmt = color->GetFormat();

    void* mapped = color->Map();
    if (!mapped) {
        throw std::runtime_error("HdRenderBuffer::Map returned null");
    }

    auto out = std::make_unique<std::vector<uint8_t>>();
    out->resize(static_cast<size_t>(bw) * bh * 4);

    if (fmt == pxr::HdFormatFloat32Vec4) {
        const float* src = static_cast<const float*>(mapped);
        for (size_t i = 0; i < static_cast<size_t>(bw) * bh; ++i) {
            for (int c = 0; c < 4; ++c) {
                float v = src[i * 4 + c];
                v = std::max(0.0f, std::min(1.0f, v));
                (*out)[i * 4 + c] = static_cast<uint8_t>(v * 255.0f + 0.5f);
            }
        }
    } else if (fmt == pxr::HdFormatFloat16Vec4) {
        const pxr::GfHalf* src = static_cast<const pxr::GfHalf*>(mapped);
        for (size_t i = 0; i < static_cast<size_t>(bw) * bh; ++i) {
            for (int c = 0; c < 4; ++c) {
                float v = static_cast<float>(src[i * 4 + c]);
                v = std::max(0.0f, std::min(1.0f, v));
                (*out)[i * 4 + c] = static_cast<uint8_t>(v * 255.0f + 0.5f);
            }
        }
    } else if (fmt == pxr::HdFormatUNorm8Vec4) {
        std::memcpy(out->data(), mapped, out->size());
    } else {
        color->Unmap();
        throw std::runtime_error(
            "Unsupported render buffer format from delegate (HdFormat=" +
            std::to_string(static_cast<int>(fmt)) + ")");
    }

    color->Unmap();

    // Storm writes the AOV in OpenGL convention: origin at the bottom-
    // left of the image. Every modern consumer (egui textures, wgpu
    // sampled textures, PNG writers, the macOS NSImage/CGImage paths,
    // …) expects origin at the top-left — without this flip the
    // returned buffer renders vertically inverted everywhere except
    // legacy OpenGL textures, which forge-paint and most callers
    // aren't using. Flipping here means consumers don't each have to
    // remember the convention.
    const size_t row_bytes = static_cast<size_t>(bw) * 4;
    if (bh > 1 && out->size() >= row_bytes * bh) {
        std::vector<uint8_t> tmp(row_bytes);
        for (uint32_t y = 0; y < bh / 2; ++y) {
            uint8_t* top = out->data() + static_cast<size_t>(y) * row_bytes;
            uint8_t* bot = out->data() + static_cast<size_t>(bh - 1 - y) * row_bytes;
            std::memcpy(tmp.data(), top, row_bytes);
            std::memcpy(top, bot, row_bytes);
            std::memcpy(bot, tmp.data(), row_bytes);
        }
    }
    return out;
}

std::unique_ptr<std::vector<uint8_t>> render_to_rgba(
    rust::Str usd_path,
    rust::Str render_delegate_id,
    uint32_t width,
    uint32_t height)
{
    auto r = create_renderer(usd_path, render_delegate_id);
    r->set_size(width, height);
    // Drive convergence ourselves now that `render_color()` does a
    // single pass and leaves the loop to the caller. Storm finishes
    // in one iteration; sampling delegates hit `is_converged()` once
    // their sample budget is exhausted. The 256-iteration cap is the
    // same safety bound the old in-engine loop carried — any sampler
    // that needs more than that should expose its budget through the
    // delegate-specific render settings (not yet wired) rather than
    // relying on this fallback.
    std::unique_ptr<std::vector<uint8_t>> out;
    for (int i = 0; i < 256; ++i) {
        out = r->render_color();
        if (r->is_converged()) break;
    }
    return out;
}

}  // namespace hydra_rs