hydra-rs 0.0.2

#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/usd/timeCode.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);
        }
    }

    // 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);
}

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

bool Renderer::set_renderer_plugin(rust::Str plugin_id) const {
    if (!engine) return false;
    std::string id(plugin_id.data(), plugin_id.size());
    return engine->SetRendererPlugin(pxr::TfToken(id));
}

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;
    params.enableLighting = !lights.empty() || use_default_lighting;
    params.enableSceneLights = false;
    params.enableSceneMaterials = true;
    params.clearColor = clear_color;

    engine->Render(stage->GetPseudoRoot(), params);
    int safety = 0;
    while (!engine->IsConverged() && safety < 256) {
        engine->Render(stage->GetPseudoRoot(), params);
        ++safety;
    }

    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();
    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);
    return r->render_color();
}

}  // namespace hydra_rs