#include "UsdSceneRenderer.h"
#include "../utils/Logger.h"
#include "../utils/GLExt.h"

#include <pxr/usd/usd/prim.h>
#include <pxr/usd/usd/primRange.h>
#include <pxr/usd/usdGeom/metrics.h>
#include <pxr/usd/usdGeom/tokens.h>
#include <pxr/usd/usdGeom/bboxCache.h>
#include <pxr/usd/usdGeom/camera.h>
#include <pxr/usd/usdGeom/xformCache.h>
#include <pxr/usd/usdLux/lightAPI.h>
#include <pxr/usd/usdLux/sphereLight.h>
#include <pxr/usd/usdLux/rectLight.h>
#include <pxr/usd/usdLux/diskLight.h>
#include <pxr/usd/usdLux/distantLight.h>
#include <pxr/usd/usdLux/domeLight.h>
#include <pxr/usd/usdLux/cylinderLight.h>
#include <pxr/imaging/glf/contextCaps.h>
#include <pxr/imaging/cameraUtil/conformWindow.h>
#include <pxr/imaging/cameraUtil/framing.h>
#include <pxr/base/gf/range2f.h>
#include <pxr/base/gf/rect2i.h>
#include <pxr/base/gf/frustum.h>
#include <pxr/base/gf/rotation.h>
#include <pxr/base/tf/diagnosticMgr.h>
#include <pxr/base/tf/error.h>
#include <pxr/base/tf/warning.h>
#include <pxr/base/tf/status.h>

#include <OpenColorIO/OpenColorIO.h>

#include <algorithm>
#include <cmath>
#include <cstdlib>
#include <exception>
#include <unordered_set>

namespace OCIO = OCIO_NAMESPACE;

namespace UsdLayerManager {

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

static std::string Vec3fStr(const pxr::GfVec3f& v) {
    return "(" + std::to_string(v[0]) + ", " + std::to_string(v[1]) + ", "
               + std::to_string(v[2]) + ")";
}

// Forwards USD's TfError / TF_WARN / TF_STATUS messages into the app log.
// Without this, Hydra diagnostics (e.g. the reason HdxColorCorrectionTask
// fails to build an OCIO processor) only go to a stderr console the GUI
// never shows.  Installed once, process-wide, from InitRenderer().
namespace {
class UsdDiagnosticLogger : public pxr::TfDiagnosticMgr::Delegate {
public:
    void IssueError(pxr::TfError const& err) override {
        LOG_ERROR("[USD] " + err.GetCommentary());
    }
    void IssueFatalError(pxr::TfCallContext const&,
                         std::string const& msg) override {
        LOG_ERROR("[USD fatal] " + msg);
    }
    void IssueStatus(pxr::TfStatus const& status) override {
        LOG_INFO("[USD] " + status.GetCommentary());
    }
    void IssueWarning(pxr::TfWarning const& warning) override {
        LOG_WARNING("[USD] " + warning.GetCommentary());
    }
};

void InstallUsdDiagnosticLogger() {
    static UsdDiagnosticLogger s_logger;   // process-lifetime; never removed
    static bool s_installed = false;
    if (s_installed) return;
    s_installed = true;
    pxr::TfDiagnosticMgr::GetInstance().AddDelegate(&s_logger);
}
} // namespace

/// Port of stageView._ComputeCameraFraming():
/// Converts a Y-up integer viewport rect into a CameraUtilFraming whose
/// display/data windows are expressed in the Y-down coordinate system that
/// CameraUtilFraming / OpenEXR use.
static pxr::CameraUtilFraming ComputeCameraFraming(
    int x, int y, int w, int h,
    int renderBufferWidth, int renderBufferHeight)
{
    // Flip Y: viewport is Y-up, display/data windows are Y-down.
    float dy  = static_cast<float>(renderBufferHeight - y - h);
    float dy2 = static_cast<float>(renderBufferHeight - y);

    pxr::GfRange2f displayWindow(
        pxr::GfVec2f(static_cast<float>(x), dy),
        pxr::GfVec2f(static_cast<float>(x + w), dy2));

    // dataWindow: integer rect, same area but Y-flipped.
    pxr::GfRect2i renderBufferRect(
        pxr::GfVec2i(0, 0),
        renderBufferWidth, renderBufferHeight);
    pxr::GfRect2i dataWindow = renderBufferRect.GetIntersection(
        pxr::GfRect2i(
            pxr::GfVec2i(x, static_cast<int>(dy)),
            w, h));

    return pxr::CameraUtilFraming(displayWindow, dataWindow);
}

static GLuint CompileShader(const char* source, GLenum type) {
    GLuint shader = glCreateShader(type);
    glShaderSource(shader, 1, &source, nullptr);
    glCompileShader(shader);
    GLint ok = 0;
    glGetShaderiv(shader, GL_COMPILE_STATUS, &ok);
    if (!ok) {
        char log[512];
        glGetShaderInfoLog(shader, sizeof(log), nullptr, log);
        LOG_ERROR("Shader compile error: " + std::string(log));
        glDeleteShader(shader);
        return 0;
    }
    return shader;
}

static GLuint LinkProgram(GLuint vs, GLuint fs) {
    GLuint prog = glCreateProgram();
    glAttachShader(prog, vs);
    glAttachShader(prog, fs);
    glLinkProgram(prog);
    glDeleteShader(vs);
    glDeleteShader(fs);
    GLint ok = 0;
    glGetProgramiv(prog, GL_LINK_STATUS, &ok);
    if (!ok) {
        char log[512];
        glGetProgramInfoLog(prog, sizeof(log), nullptr, log);
        LOG_ERROR("Program link error: " + std::string(log));
        glDeleteProgram(prog);
        return 0;
    }
    return prog;
}

// ===========================================================================
// ViewportColorCorrector — custom GL color correction (sRGB / OCIO)
// ===========================================================================
//
// Replaces HdxColorCorrectionTask.  The scene is rendered linear; Apply()
// samples that linear texture and draws a corrected fullscreen triangle into
// the currently-bound FBO.  For OCIO it uses the OCIO GPU shader API directly
// (GpuShaderDesc) so the exact transform — including 1D/3D LUTs — is built and
// bound under our control, independent of Hydra.
class ViewportColorCorrector {
public:
    enum class Mode { sRGB, OpenColorIO };

    ~ViewportColorCorrector() { DestroyGLResources(); }

    /// Draw a corrected fullscreen quad into the bound FBO sampling srcTex
    /// (single-sample, linear RGBA).  Caller has set the GL viewport.
    /// Returns true if a pass was drawn.  On OCIO build failure it falls back
    /// to the sRGB encode so the viewport degrades gracefully (never black).
    bool Apply(GLuint srcTex, Mode mode,
               const std::string& disp, const std::string& view,
               const std::string& cs,   const std::string& look)
    {
        if (!EnsureCommon()) return false;

        GLuint prog = m_srgbProgram;
        std::vector<LutTex>* luts = nullptr;
        if (mode == Mode::OpenColorIO && EnsureOcioProgram(disp, view, cs, look)) {
            prog = m_ocioProgram;
            luts = &m_ocioLuts;
        }
        if (!prog) return false;

        // Save the GL state we touch.
        GLboolean depthTest = glIsEnabled(GL_DEPTH_TEST);
        GLboolean blend     = glIsEnabled(GL_BLEND);
        GLboolean depthMask = GL_TRUE;
        glGetBooleanv(GL_DEPTH_WRITEMASK, &depthMask);
        glDisable(GL_DEPTH_TEST);
        glDisable(GL_BLEND);
        glDepthMask(GL_FALSE);

        glUseProgram(prog);
        glActiveTexture(GL_TEXTURE0);
        glBindTexture(GL_TEXTURE_2D, srcTex);
        glUniform1i(glGetUniformLocation(prog, "uTex"), 0);

        if (luts) {
            int unit = 1;
            for (const auto& l : *luts) {
                glActiveTexture(GL_TEXTURE0 + unit);
                glBindTexture(l.target, l.id);
                GLint loc = glGetUniformLocation(prog, l.sampler.c_str());
                if (loc >= 0) glUniform1i(loc, unit);
                ++unit;
            }
            glActiveTexture(GL_TEXTURE0);
        }

        glBindVertexArray(m_vao);
        glDrawArrays(GL_TRIANGLES, 0, 3);
        glBindVertexArray(0);
        glUseProgram(0);

        // Restore state.
        if (depthTest) glEnable(GL_DEPTH_TEST);
        if (blend)     glEnable(GL_BLEND);
        glDepthMask(depthMask);
        return true;
    }

    void DestroyGLResources() {
        DestroyOcioResources();
        if (m_srgbProgram) { glDeleteProgram(m_srgbProgram); m_srgbProgram = 0; }
        if (m_vao)         { glDeleteVertexArrays(1, &m_vao); m_vao = 0; }
    }

private:
    struct LutTex { GLuint id; GLenum target; std::string sampler; };

    bool EnsureCommon() {
        if (m_vao == 0) glGenVertexArrays(1, &m_vao);
        if (m_srgbProgram == 0) {
            GLuint vs = CompileShader(kFullscreenVS, GL_VERTEX_SHADER);
            GLuint fs = CompileShader(kSrgbFS,       GL_FRAGMENT_SHADER);
            if (vs && fs) m_srgbProgram = LinkProgram(vs, fs);
        }
        return m_vao != 0 && m_srgbProgram != 0;
    }

    bool EnsureOcioProgram(const std::string& disp, const std::string& view,
                           const std::string& cs,   const std::string& look)
    {
        const std::string key = disp + "|" + view + "|" + cs + "|" + look;
        if (m_ocioProgram && key == m_ocioKey)      return true;
        if (key == m_ocioFailedKey)                 return false;

        DestroyOcioResources();
        m_ocioKey.clear();

        std::string fragText;
        OCIO::GpuShaderDescRcPtr desc;
        try {
            OCIO::ConstConfigRcPtr config = OCIO::GetCurrentConfig();
            if (!config) throw std::runtime_error("no current OCIO config");
            const char* srcCS = cs.empty() ? OCIO::ROLE_SCENE_LINEAR : cs.c_str();

            OCIO::ConstProcessorRcPtr proc = config->getProcessor(
                srcCS, disp.c_str(), view.c_str(), OCIO::TRANSFORM_DIR_FORWARD);
            OCIO::ConstGPUProcessorRcPtr gpu = proc->getDefaultGPUProcessor();

            desc = OCIO::GpuShaderDesc::CreateShaderDesc();
            desc->setLanguage(OCIO::GPU_LANGUAGE_GLSL_1_3);
            desc->setFunctionName("OCIODisplay");
            desc->setResourcePrefix("ocio_");
            gpu->extractGpuShaderInfo(desc);
            fragText = desc->getShaderText();
        } catch (const std::exception& e) {
            LOG_WARNING("Custom OCIO build failed (" + key
                        + "), using sRGB: " + std::string(e.what()));
            m_ocioFailedKey = key;
            return false;
        }

        // Assemble the fragment shader: OCIO declares its samplers + the
        // OCIODisplay(vec4) function; our main() samples the linear input and
        // runs it through.
        std::string fs = "#version 130\n"
                         "uniform sampler2D uTex;\n"
                         "in vec2 vUv;\n"
                         "out vec4 outColor;\n"
                       + fragText +
                         "\nvoid main(){ outColor = OCIODisplay(texture(uTex, vUv)); }\n";

        GLuint vsh = CompileShader(kFullscreenVS, GL_VERTEX_SHADER);
        GLuint fsh = CompileShader(fs.c_str(),    GL_FRAGMENT_SHADER);
        GLuint prog = (vsh && fsh) ? LinkProgram(vsh, fsh) : 0;
        if (!prog) {
            LOG_WARNING("Custom OCIO shader compile/link failed for " + key
                        + " — using sRGB");
            m_ocioFailedKey = key;
            return false;
        }

        // Upload the LUT textures OCIO requested.
        if (!CreateOcioTextures(desc)) {
            glDeleteProgram(prog);
            DestroyOcioResources();
            m_ocioFailedKey = key;
            return false;
        }

        m_ocioProgram = prog;
        m_ocioKey     = key;
        LOG_INFO("Custom OCIO program built: " + key
                 + " (" + std::to_string(m_ocioLuts.size()) + " LUTs)");
        return true;
    }

    bool CreateOcioTextures(const OCIO::GpuShaderDescRcPtr& desc) {
        // 3D LUTs (RGB).
        for (unsigned i = 0; i < desc->getNum3DTextures(); ++i) {
            const char* texName = nullptr; const char* samplerName = nullptr;
            unsigned edgelen = 0; OCIO::Interpolation interp = OCIO::INTERP_LINEAR;
            desc->get3DTexture(i, texName, samplerName, edgelen, interp);
            const float* values = nullptr;
            desc->get3DTextureValues(i, values);
            if (!values || edgelen == 0 || !samplerName) return false;

            GLuint id = 0;
            glGenTextures(1, &id);
            glBindTexture(GL_TEXTURE_3D, id);
            glTexImage3D(GL_TEXTURE_3D, 0, GL_RGB32F, edgelen, edgelen, edgelen,
                         0, GL_RGB, GL_FLOAT, values);
            GLint filt = (interp == OCIO::INTERP_NEAREST) ? GL_NEAREST : GL_LINEAR;
            glTexParameteri(GL_TEXTURE_3D, GL_TEXTURE_MIN_FILTER, filt);
            glTexParameteri(GL_TEXTURE_3D, GL_TEXTURE_MAG_FILTER, filt);
            glTexParameteri(GL_TEXTURE_3D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
            glTexParameteri(GL_TEXTURE_3D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
            glTexParameteri(GL_TEXTURE_3D, GL_TEXTURE_WRAP_R, GL_CLAMP_TO_EDGE);
            m_ocioLuts.push_back({ id, GL_TEXTURE_3D, samplerName });
        }

        // 1D / 2D LUTs.
        for (unsigned i = 0; i < desc->getNumTextures(); ++i) {
            const char* texName = nullptr; const char* samplerName = nullptr;
            unsigned width = 0, height = 0;
            OCIO::GpuShaderDesc::TextureType channel = OCIO::GpuShaderDesc::TEXTURE_RGB_CHANNEL;
            OCIO::Interpolation interp = OCIO::INTERP_LINEAR;
            desc->getTexture(i, texName, samplerName, width, height, channel, interp);
            const float* values = nullptr;
            desc->getTextureValues(i, values);
            if (!values || width == 0 || !samplerName) return false;

            const bool   isRed   = (channel == OCIO::GpuShaderDesc::TEXTURE_RED_CHANNEL);
            const GLint  internal = isRed ? GL_R32F : GL_RGB32F;
            const GLenum format   = isRed ? GL_RED  : GL_RGB;
            const GLint  filt = (interp == OCIO::INTERP_NEAREST) ? GL_NEAREST : GL_LINEAR;
            const GLenum target = (height > 1) ? GL_TEXTURE_2D : GL_TEXTURE_1D;

            GLuint id = 0;
            glGenTextures(1, &id);
            glBindTexture(target, id);
            if (target == GL_TEXTURE_2D) {
                glTexImage2D(GL_TEXTURE_2D, 0, internal, width, height, 0,
                             format, GL_FLOAT, values);
                glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
            } else {
                glTexImage1D(GL_TEXTURE_1D, 0, internal, width, 0,
                             format, GL_FLOAT, values);
            }
            glTexParameteri(target, GL_TEXTURE_MIN_FILTER, filt);
            glTexParameteri(target, GL_TEXTURE_MAG_FILTER, filt);
            glTexParameteri(target, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
            m_ocioLuts.push_back({ id, target, samplerName });
        }
        return true;
    }

    void DestroyOcioResources() {
        for (auto& l : m_ocioLuts) glDeleteTextures(1, &l.id);
        m_ocioLuts.clear();
        if (m_ocioProgram) { glDeleteProgram(m_ocioProgram); m_ocioProgram = 0; }
        m_ocioKey.clear();
        m_ocioFailedKey.clear();
    }

    static const char* kFullscreenVS;
    static const char* kSrgbFS;

    GLuint m_vao         = 0;
    GLuint m_srgbProgram = 0;
    GLuint m_ocioProgram = 0;
    std::string m_ocioKey;
    std::string m_ocioFailedKey;
    std::vector<LutTex> m_ocioLuts;
};

// Attribute-less fullscreen triangle; UV in [0,2] covers the [0,1] screen.
const char* ViewportColorCorrector::kFullscreenVS = R"(#version 130
out vec2 vUv;
void main() {
    vec2 p = vec2(float((gl_VertexID << 1) & 2), float(gl_VertexID & 2));
    vUv = p;
    gl_Position = vec4(p * 2.0 - 1.0, 0.0, 1.0);
}
)";

// Linear → sRGB encode (matches HdxColorCorrectionTask's sRGB path).
const char* ViewportColorCorrector::kSrgbFS = R"(#version 130
uniform sampler2D uTex;
in vec2 vUv;
out vec4 outColor;
vec3 lin2srgb(vec3 c) {
    vec3 lo = c * 12.92;
    vec3 hi = 1.055 * pow(max(c, vec3(0.0)), vec3(1.0/2.4)) - 0.055;
    bvec3 cut = lessThanEqual(c, vec3(0.0031308));
    return mix(hi, lo, vec3(cut));
}
void main() {
    vec4 c = texture(uTex, vUv);
    outColor = vec4(lin2srgb(c.rgb), c.a);
}
)";

// ---------------------------------------------------------------------------
// GLSL sources
// ---------------------------------------------------------------------------

// The grid reuses the same axis shader (kAxisVS / kAxisFS) — no separate
// grid program is needed.  Line geometry is built in RebuildGridVBO().

// Axis shader (stageView.DrawAxis port)
// VS receives 3D position; uniform MVP scales and projects.
static const char* kAxisVS = R"(#version 130
in vec3 position;
uniform mat4 mvpMatrix;
void main() { gl_Position = vec4(position, 1.0) * mvpMatrix; }
)";

static const char* kAxisFS = R"(#version 130
uniform vec4 color;
out vec4 outColor;
void main() { outColor = color; }
)";

// ===========================================================================
// Constructor / Destructor
// ===========================================================================

UsdSceneRenderer::UsdSceneRenderer()
    : m_showGrid(true)
    , m_aaEnabled(false)
    , m_backgroundColor(0.15f, 0.15f, 0.15f)
    , m_shadingMode(ShadingMode::SmoothShaded)
    , m_colorCorrectionMode(ColorCorrectionMode::sRGB)
    , m_ambientLightOnly(true)
    , m_domeLightEnabled(false)
    , m_stageIsZup(false)
    , m_defaultMaterialAmbient(0.2f)
    , m_defaultMaterialSpecular(0.1f)
    , m_rendererInitialized(false)
    , m_forceRefresh(false)
    , m_useCameraPath(false)
    , m_diagFrameCount(0)
    , m_gridVAO(0)
    , m_gridVBO(0)
    , m_gridMinorFirst(0),  m_gridMinorCount(0)
    , m_gridMajorFirst(0),  m_gridMajorCount(0)
    , m_gridAxisAFirst(0),  m_gridAxisACount(0)
    , m_gridAxisBFirst(0),  m_gridAxisBCount(0)
    , m_gridHalfSize(50.0f)
    , m_axisVAO(0)
    , m_axisVBO(0)
    , m_axisProgram(0)
    , m_axisUniformMVP(-1)
    , m_axisUniformColor(-1)
    , m_bboxMode(BBoxMode::None)
    , m_bboxColor(1.0f, 1.0f, 1.0f, 1.0f)
    , m_bboxVAO(0)
    , m_bboxVBO(0)
    , m_bboxProgram(0)
    , m_bboxUniformMVP(-1)
    , m_bboxUniformColor(-1)
{
    m_viewMatrix.SetIdentity();
    m_projMatrix.SetIdentity();
}

UsdSceneRenderer::~UsdSceneRenderer() {
    DestroyGridResources();
    DestroyAxisResources();
    DestroyBBoxResources();
    DestroyCamWireResources();
    DestroyLightWireResources();
    m_colorCorrector.reset();   // deletes its GL program / LUT textures
    if (m_ccLinearFBO) glDeleteFramebuffers(1, &m_ccLinearFBO);
    if (m_ccLinearTex) glDeleteTextures(1, &m_ccLinearTex);
}

// ===========================================================================
// Stage
// ===========================================================================

void UsdSceneRenderer::SetStage(pxr::UsdStageRefPtr stage) {
    m_stage = stage;
    m_rendererInitialized = false;
    m_useCameraPath = false;
    m_cameraPath    = pxr::SdfPath();
    m_clipPlanes.clear();
    m_cameraCacheDirty = true;
    DestroyGridResources();
    DestroyAxisResources();
    DestroyBBoxResources();
    DestroyCamWireResources();
    DestroyLightWireResources();

    // Determine stage up-axis for dome light rotation (mirrors stageView._stageIsZup)
    if (stage) {
        pxr::TfToken upAxis = pxr::UsdGeomGetStageUpAxis(stage);
        m_stageIsZup = (upAxis == pxr::UsdGeomTokens->z);
    } else {
        m_stageIsZup = false;
    }

    // Rebuild grid geometry for the new up-axis (if GL resources are ready)
    RebuildGridVBO();
}

// ===========================================================================
// Renderer Initialization (lazy, deferred to first Render())
// ===========================================================================

void UsdSceneRenderer::InitRenderer() {
    if (m_rendererInitialized || !m_stage) return;

    LOG_INFO("UsdSceneRenderer::InitRenderer - initializing...");

    // Route USD/Hydra diagnostics (incl. OCIO failures) to the app log.
    InstallUsdDiagnosticLogger();

    pxr::GlfContextCaps::InitInstance();

    pxr::UsdImagingGLEngine::Parameters params;
    params.rootPath    = m_stage->GetPseudoRoot().GetPath();
    params.excludedPaths = {};
    m_renderer = std::make_shared<pxr::UsdImagingGLEngine>(params);
    if (!m_renderer) {
        LOG_ERROR("Failed to create UsdImagingGLEngine");
        return;
    }

    // Plugin selection: prefer HdStorm / GL-based renderers
    auto plugins = pxr::UsdImagingGLEngine::GetRendererPlugins();
    LOG_INFO("Available renderer plugins: " + std::to_string(plugins.size()));
    bool hasCycles = false;
    for (const auto& p : plugins) {
        LOG_INFO("  " + std::string(p.GetText()) + " -> "
                 + pxr::UsdImagingGLEngine::GetRendererDisplayName(p));
        if (std::string(p.GetText()) == "HdCyclesPlugin") hasCycles = true;
    }
    if (!hasCycles)
        LOG_WARNING("HdCyclesPlugin not available (see startup log for DLL load errors).");

    pxr::TfToken currentPlugin = m_renderer->GetCurrentRendererId();
    if (currentPlugin.IsEmpty() && !plugins.empty()) {
        pxr::TfToken best;
        // If a plugin was previously selected, honour it
        if (!m_currentRendererPlugin.IsEmpty()) {
            best = m_currentRendererPlugin;
        } else {
            for (const auto& p : plugins) {
                std::string n(p.GetText());
                if (n.find("Storm") != std::string::npos ||
                    n.find("GL")    != std::string::npos) { best = p; break; }
            }
            if (best.IsEmpty()) best = plugins[0];
        }
        m_renderer->SetRendererPlugin(best);
        LOG_INFO("Selected renderer: " + std::string(m_renderer->GetCurrentRendererId().GetText()));
    }

    // Enable AOV "color" (matches stageView._handleRendererChanged)
    m_renderer->SetRendererAov(pxr::TfToken("color"));

    // Selection highlight color (usdview default: yellow)
    m_renderer->SetSelectionColor(pxr::GfVec4f(1.0f, 1.0f, 0.0f, 1.0f));

    // Remember which plugin is active
    m_currentRendererPlugin = m_renderer->GetCurrentRendererId();

    m_rendererInitialized = true;
    LOG_INFO("UsdSceneRenderer initialized.");

    InitGridResources();
    InitAxisResources();
    InitBBoxResources();
    InitCamWireResources();
    InitLightWireResources();
}

// ===========================================================================
// Render delegate
// ===========================================================================

/*static*/ std::vector<pxr::TfToken> UsdSceneRenderer::GetRendererPlugins() {
    return pxr::UsdImagingGLEngine::GetRendererPlugins();
}

pxr::TfToken UsdSceneRenderer::GetCurrentRendererId() const {
    if (m_renderer) return m_renderer->GetCurrentRendererId();
    return m_currentRendererPlugin;
}

/*static*/ std::string UsdSceneRenderer::GetRendererDisplayName(const pxr::TfToken& pluginId) {
    return pxr::UsdImagingGLEngine::GetRendererDisplayName(pluginId);
}

bool UsdSceneRenderer::SetRendererPlugin(const pxr::TfToken& pluginId) {
    m_currentRendererPlugin = pluginId;

    if (!m_renderer) {
        // Will be picked up on first Render() call via InitRenderer()
        return true;
    }

    bool ok = m_renderer->SetRendererPlugin(pluginId);
    if (ok) {
        // Re-enable colour AOV after delegate switch (mirrors usdtweak behaviour)
        m_renderer->SetRendererAov(pxr::TfToken("color"));
        LOG_INFO("Render delegate switched to: " + std::string(pluginId.GetText()));
    } else {
        LOG_ERROR("Failed to switch render delegate to: " + std::string(pluginId.GetText()));
    }
    return ok;
}

// ===========================================================================
// Camera State
// ===========================================================================

void UsdSceneRenderer::SetCameraState(
    const pxr::GfMatrix4d& viewMatrix,
    const pxr::GfMatrix4d& projMatrix)
{
    m_viewMatrix    = viewMatrix;
    m_projMatrix    = projMatrix;
    m_useCameraPath = false;
    m_clipPlanes.clear();
}

void UsdSceneRenderer::SetCameraStateFromGfCamera(const pxr::GfCamera& gfCamera)
{
    pxr::GfFrustum frustum = gfCamera.GetFrustum();
    m_viewMatrix    = frustum.ComputeViewMatrix();
    m_projMatrix    = frustum.ComputeProjectionMatrix();
    m_useCameraPath = false;
    m_cameraFrustum    = frustum;
    m_hasCameraFrustum = true;

    // Extract clip planes from GfCamera (same as stageView's renderParams.clipPlanes)
    m_clipPlanes.clear();
    for (const auto& p : gfCamera.GetClippingPlanes()) {
        m_clipPlanes.emplace_back(
            static_cast<double>(p[0]), static_cast<double>(p[1]),
            static_cast<double>(p[2]), static_cast<double>(p[3]));
    }
}

void UsdSceneRenderer::SetCameraPath(const pxr::SdfPath& cameraPath)
{
    m_cameraPath    = cameraPath;
    m_useCameraPath = true;
}

// ===========================================================================
// Selection
// ===========================================================================

void UsdSceneRenderer::ClearSelected()
{
    if (m_renderer) m_renderer->ClearSelected();
}

void UsdSceneRenderer::AddSelected(const pxr::SdfPath& path, int instanceIndex)
{
    if (m_renderer) m_renderer->AddSelected(path, instanceIndex);
}

void UsdSceneRenderer::SetSelectedPaths(const pxr::SdfPathVector& paths)
{
    if (m_renderer) m_renderer->SetSelected(paths);
}

// ===========================================================================
// Bounds
// ===========================================================================

pxr::GfRange3d UsdSceneRenderer::ComputeStageBounds()
{
    if (!m_stage) return {};
    pxr::TfTokenVector purposes = {
        pxr::UsdGeomTokens->default_, pxr::UsdGeomTokens->proxy };
    pxr::UsdGeomBBoxCache bboxCache(m_currentTime, purposes, true);
    return bboxCache.ComputeWorldBound(m_stage->GetPseudoRoot()).ComputeAlignedRange();
}

// ===========================================================================
// Picking (stageView.pick / computePickFrustum port)
// ===========================================================================

bool UsdSceneRenderer::PickObject(
    int mouseX, int mouseY,
    int viewWidth, int viewHeight,
    pxr::GfVec3d*  outHitPoint,
    pxr::SdfPath*  outHitPrimPath)
{
    if (!m_renderer || !m_stage) return false;
    if (!m_hasCameraFrustum) return false;

    // Normalize mouse to NDC [-1, 1]; Y is flipped (screen Y-down → NDC Y-up)
    double nx = (static_cast<double>(mouseX) / static_cast<double>(viewWidth))  * 2.0 - 1.0;
    double ny = 1.0 - (static_cast<double>(mouseY) / static_cast<double>(viewHeight)) * 2.0;
    pxr::GfVec2d point(nx, ny);
    // Pick window: one pixel in NDC
    pxr::GfVec2d size(1.0 / static_cast<double>(viewWidth),
                      1.0 / static_cast<double>(viewHeight));

    // Build narrow pick frustum from the stored camera frustum
    // (mirrors stageView.computePickFrustum / GfFrustum.ComputeNarrowedFrustum)
    pxr::GfFrustum pickFrustum = m_cameraFrustum.ComputeNarrowedFrustum(point, size);

    pxr::UsdImagingGLRenderParams pickParams;
    pickParams.drawMode   = pxr::UsdImagingGLDrawMode::DRAW_GEOM_ONLY; // no shading: faster pick pass
    pickParams.showGuides = false;
    pickParams.showProxy  = true;
    pickParams.showRender = false;
    pickParams.enableSampleAlphaToCoverage = false;
    pickParams.enableLighting              = false;

    pxr::GfVec3d hitPoint, hitNormal;
    pxr::SdfPath hitPrimPath;

    bool hit = m_renderer->TestIntersection(
        pickFrustum.ComputeViewMatrix(),
        pickFrustum.ComputeProjectionMatrix(),
        m_stage->GetPseudoRoot(),
        pickParams,
        &hitPoint, &hitNormal, &hitPrimPath);

    if (hit) {
        if (outHitPoint)    *outHitPoint    = hitPoint;
        if (outHitPrimPath) *outHitPrimPath = hitPrimPath;
        LOG_INFO("PickObject hit: " + hitPrimPath.GetString());
    }
    return hit;
}

bool UsdSceneRenderer::PickObjectsInRect(
    int x0, int y0, int x1, int y1,
    int viewWidth, int viewHeight,
    pxr::SdfPathVector* outHitPaths)
{
    if (!m_renderer || !m_stage || !m_hasCameraFrustum) return false;
    if (!outHitPaths) return false;
    outHitPaths->clear();

    // Clamp rect to viewport bounds
    x0 = std::max(0, std::min(x0, viewWidth  - 1));
    x1 = std::max(0, std::min(x1, viewWidth  - 1));
    y0 = std::max(0, std::min(y0, viewHeight - 1));
    y1 = std::max(0, std::min(y1, viewHeight - 1));
    if (x0 > x1) std::swap(x0, x1);
    if (y0 > y1) std::swap(y0, y1);
    if (x0 == x1 || y0 == y1) return false;

    // Convert rect to NDC [-1, 1]; Y is flipped (screen Y-down → NDC Y-up)
    double ndcCenterX = ((static_cast<double>(x0 + x1) * 0.5) / viewWidth)  * 2.0 - 1.0;
    double ndcCenterY = 1.0 - ((static_cast<double>(y0 + y1) * 0.5) / viewHeight) * 2.0;
    double ndcSizeX   =  static_cast<double>(x1 - x0) / viewWidth;
    double ndcSizeY   =  static_cast<double>(y1 - y0) / viewHeight;

    pxr::GfVec2d center(ndcCenterX, ndcCenterY);
    pxr::GfVec2d size  (ndcSizeX,   ndcSizeY);

    // Narrow the stored camera frustum to the selection rect
    pxr::GfFrustum pickFrustum = m_cameraFrustum.ComputeNarrowedFrustum(center, size);

    pxr::UsdImagingGLRenderParams pickParams;
    pickParams.drawMode   = pxr::UsdImagingGLDrawMode::DRAW_GEOM_ONLY; // no shading: faster pick pass
    pickParams.showGuides = false;
    pickParams.showProxy  = true;
    pickParams.showRender = false;
    pickParams.enableSampleAlphaToCoverage = false;
    pickParams.enableLighting              = false;

    // resolveUnique: single pick-buffer render pass, returns all unique VISIBLE prims.
    // Much faster than resolveDeep (which does a full deep-selection traversal).
    pxr::UsdImagingGLEngine::PickParams pp;
    pp.resolveMode = pxr::TfToken("resolveUnique");

    pxr::UsdImagingGLEngine::IntersectionResultVector results;
    bool hit = m_renderer->TestIntersection(
        pp,
        pickFrustum.ComputeViewMatrix(),
        pickFrustum.ComputeProjectionMatrix(),
        m_stage->GetPseudoRoot(),
        pickParams,
        &results);

    if (hit) {
        // Deduplicate by prim path (a prim may appear multiple times for
        // different instances or mesh subsets within the rect).
        std::unordered_set<std::string> seen;
        for (const auto& r : results) {
            const std::string& s = r.hitPrimPath.GetString();
            if (!s.empty() && seen.insert(s).second) {
                outHitPaths->push_back(r.hitPrimPath);
            }
        }
        LOG_INFO("PickObjectsInRect: " + std::to_string(outHitPaths->size()) + " prims selected");
    }
    return !outHitPaths->empty();
}

// ===========================================================================
// Render
// ===========================================================================

void UsdSceneRenderer::Render(int width, int height)
{
    if (!m_stage || width <= 0 || height <= 0) return;
    m_lastRenderWidth  = width;
    m_lastRenderHeight = height;

    InitRenderer();
    if (!m_renderer) return;

    // (Re)create draw target when it doesn't exist or when the AA/MSAA
    // preference has changed.  With AA enabled we request a multisampled
    // FBO so Hydra and all overlay geometry benefit from hardware MSAA.
    const bool wantMSAA = m_aaEnabled;
    if (m_drawTarget && m_drawTarget->HasMSAA() != wantMSAA) {
        if (m_drawTarget->IsBound()) m_drawTarget->Unbind();
        m_drawTarget.Reset();
        LOG_INFO("DrawTarget recreated: AA toggled (" +
                 std::string(wantMSAA ? "MSAA on" : "MSAA off") + ")");
    }
    if (!m_drawTarget) {
        m_drawTarget = pxr::GlfDrawTarget::New(pxr::GfVec2i(width, height), wantMSAA);
        if (!m_drawTarget) { LOG_ERROR("Failed to create GlfDrawTarget"); return; }
        m_drawTarget->Bind();
        // RGBA16F so the scene is stored linear with HDR headroom: we render
        // linear (Hydra correction disabled) and apply our own color correction
        // afterwards, which needs values outside [0,1] for OCIO/ACES.
        m_drawTarget->AddAttachment("color", GL_RGBA, GL_FLOAT, GL_RGBA16F);
        m_drawTarget->AddAttachment("depth",
            GL_DEPTH_COMPONENT, GL_FLOAT, GL_DEPTH_COMPONENT32F);
        m_drawTarget->Unbind();
        LOG_INFO("DrawTarget created: fboId="
            + std::to_string(m_drawTarget->GetFramebufferId())
            + " msaa=" + std::string(wantMSAA ? "yes" : "no")
            + " size=" + std::to_string(width) + "x" + std::to_string(height));
    }

    // Bind first, then resize if needed.
    // GlfDrawTarget::SetSize requires the FBO to be bound (asserts otherwise).
    m_drawTarget->Bind();
    pxr::GfVec2i desiredSize(width, height);
    if (m_drawTarget->GetSize() != desiredSize) {
        m_drawTarget->SetSize(desiredSize);
    }

    // Clear
    glViewport(0, 0, width, height);
    glClearColor(m_backgroundColor[0], m_backgroundColor[1],
                 m_backgroundColor[2], 1.0f);
    glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);

    // --- Camera state ---
    if (m_useCameraPath && !m_cameraPath.IsEmpty()) {
        m_renderer->SetCameraPath(m_cameraPath);
    } else {
        m_renderer->SetCameraState(m_viewMatrix, m_projMatrix);
    }

    // --- Use SetRenderBufferSize + SetFraming instead of deprecated SetRenderViewport ---
    // (mirrors stageView.paintGL: renderer.SetRenderBufferSize + renderer.SetFraming)
    m_renderer->SetRenderBufferSize(pxr::GfVec2i(width, height));
    m_renderer->SetFraming(ComputeCameraFraming(0, 0, width, height, width, height));
    m_renderer->SetOverrideWindowPolicy(pxr::CameraUtilMatchVertically);

    // --- Lighting: mirrors stageView.py paintGL lighting setup ---
    // stageView uses two optional lights controlled by viewSettings:
    //   ambientLightOnly  (default True)  → camera headlight (point at cam pos)
    //   domeLightEnabled  (default False) → dome/IBL light
    // sceneAmbient and material values from viewSettingsDataModel.py defaults.

    pxr::GfVec4f sceneAmbient(0.01f, 0.01f, 0.01f, 1.0f);

    pxr::GlfSimpleMaterial material;
    float kA = m_defaultMaterialAmbient;   // 0.2
    float kS = m_defaultMaterialSpecular;  // 0.1
    material.SetAmbient (pxr::GfVec4f(kA, kA, kA, 1.0f));
    material.SetSpecular(pxr::GfVec4f(kS, kS, kS, 1.0f));
    material.SetShininess(32.0f);

    pxr::GlfSimpleLightVector lights;

    // Camera headlight: point light (w=1) positioned at the camera world-origin,
    // transformed by the view-inverse so it tracks the camera each frame.
    // (stageView.py: l.position = cam_pos + (1,); l.transform = frustum.ComputeViewInverse())
    //
    // The headlight is the *default* fill used only while the stage has no
    // authored lights.  Once the scene contains real UsdLux lights, suppress it
    // so the scene is lit purely by those lights (Hydra evaluates them via
    // enableSceneLights) -- mirrors usdview's "use scene lights when present".
    const bool stageHasLights = StageHasAuthoredLights();
    if (m_ambientLightOnly && !stageHasLights) {
        pxr::GfMatrix4d viewInverse = m_viewMatrix.GetInverse();
        pxr::GfVec3d camPos = viewInverse.ExtractTranslation();

        pxr::GlfSimpleLight camLight;
        camLight.SetAmbient (pxr::GfVec4f(0.0f, 0.0f, 0.0f, 0.0f));
        camLight.SetDiffuse (pxr::GfVec4f(1.0f, 1.0f, 1.0f, 1.0f));
        camLight.SetSpecular(pxr::GfVec4f(1.0f, 1.0f, 1.0f, 1.0f));
        camLight.SetPosition(pxr::GfVec4f(
            static_cast<float>(camPos[0]),
            static_cast<float>(camPos[1]),
            static_cast<float>(camPos[2]),
            1.0f));                        // w=1 → point light
        camLight.SetTransform(viewInverse);
        lights.push_back(camLight);
    }

    // Dome light (IBL): isDomeLight=true, Z-up stages need a 90° X-axis rotation.
    // (stageView.py: l.isDomeLight = True; if stageIsZup: l.transform = rot90X)
    if (m_domeLightEnabled) {
        pxr::GlfSimpleLight domeLight;
        domeLight.SetIsDomeLight(true);
        if (m_stageIsZup) {
            pxr::GfMatrix4d rot;
            rot.SetRotate(pxr::GfRotation(pxr::GfVec3d::XAxis(), 90.0));
            domeLight.SetTransform(rot);
        }
        lights.push_back(domeLight);
    }

    m_renderer->SetLightingState(lights, material, sceneAmbient);

    // --- Render params (matches stageView.renderSinglePass) ---
    m_renderParams = pxr::UsdImagingGLRenderParams();
    m_renderParams.frame                    = m_currentTime;
    m_renderParams.complexity               = 1.0f;
    // Apply shading mode → drawMode + enableLighting
    switch (m_shadingMode) {
        case ShadingMode::FlatShaded:
            m_renderParams.drawMode     = pxr::UsdImagingGLDrawMode::DRAW_SHADED_FLAT;
            m_renderParams.enableLighting = true;  break;
        case ShadingMode::WireframeOnSurface:
            m_renderParams.drawMode     = pxr::UsdImagingGLDrawMode::DRAW_WIREFRAME_ON_SURFACE;
            m_renderParams.enableLighting = true;  break;
        case ShadingMode::Wireframe:
            m_renderParams.drawMode     = pxr::UsdImagingGLDrawMode::DRAW_WIREFRAME;
            m_renderParams.enableLighting = false; break;
        case ShadingMode::Unlit:
            m_renderParams.drawMode     = pxr::UsdImagingGLDrawMode::DRAW_SHADED_SMOOTH;
            m_renderParams.enableLighting = false; break;
        default: // SmoothShaded
            m_renderParams.drawMode     = pxr::UsdImagingGLDrawMode::DRAW_SHADED_SMOOTH;
            m_renderParams.enableLighting = true;  break;
    }
    m_renderParams.showGuides               = true;
    m_renderParams.showProxy                = true;
    m_renderParams.showRender               = false;
    m_renderParams.enableSampleAlphaToCoverage = true;
    m_renderParams.gammaCorrectColors       = false;
    m_renderParams.cullStyle                = pxr::UsdImagingGLCullStyle::CULL_STYLE_BACK_UNLESS_DOUBLE_SIDED;
    m_renderParams.enableSceneMaterials     = true;
    m_renderParams.enableSceneLights        = true;
    m_renderParams.highlight                = true;
    m_renderParams.clearColor               = pxr::GfVec4f(
        m_backgroundColor[0], m_backgroundColor[1], m_backgroundColor[2], 1.0f);
    m_renderParams.forceRefresh             = m_forceRefresh;
    m_renderParams.clipPlanes               = m_clipPlanes;

    // Color correction is done by our own GL post-process (ApplyViewport-
    // ColorCorrection below), not HdxColorCorrectionTask — so Hydra always
    // renders linear ("disabled").  This bypasses the hdx OCIO path entirely.
    m_renderParams.colorCorrectionMode = pxr::TfToken("disabled");
    m_renderer->SetColorCorrectionSettings(pxr::TfToken("disabled"));

    // Guard the Hydra render: if it throws, the m_drawTarget->Unbind() below
    // would be skipped, permanently unbalancing the GlfDrawTarget bind stack
    // and blacking out every later frame.
    try {
        m_renderer->Render(m_stage->GetPseudoRoot(), m_renderParams);
    } catch (const std::exception& e) {
        LOG_ERROR("Hydra render failed: " + std::string(e.what()));
    }
    m_forceRefresh = false;

    // --- Custom color-correction post-process (linear → sRGB / OCIO) ---
    if (m_colorCorrectionMode != ColorCorrectionMode::Disabled) {
        ApplyViewportColorCorrection(width, height);
    }

    // --- Optional grid overlay ---
    if (m_showGrid) {
        // When MSAA is active, render the grid into the MSAA FBO so it is
        // also multisampled and resolved together with the Hydra output.
        GLuint gridFbo = m_drawTarget->HasMSAA()
            ? m_drawTarget->GetFramebufferMSId()
            : m_drawTarget->GetFramebufferId();
        glBindFramebuffer(GL_FRAMEBUFFER, gridFbo);
        glViewport(0, 0, width, height);
        RenderGrid(width, height);
    }

    // --- Diagnostic logging ---
    if (m_diagFrameCount < 3) {
        auto att = m_drawTarget->GetAttachment("color");
        if (att) {
            LOG_INFO("Frame " + std::to_string(m_diagFrameCount)
                + ": texId=" + std::to_string(att->GetGlTextureName())
                + " fboId=" + std::to_string(m_drawTarget->GetFramebufferId())
                + " size=" + std::to_string(width) + "x" + std::to_string(height)
                + " bg=" + Vec3fStr(m_backgroundColor));
        }
        ++m_diagFrameCount;
    }

    m_drawTarget->Unbind();
}

// ===========================================================================
// Custom color-correction post-process
// ===========================================================================

void UsdSceneRenderer::ApplyViewportColorCorrection(int width, int height)
{
    if (!m_drawTarget) return;

    if (!m_colorCorrector)
        m_colorCorrector = std::make_unique<ViewportColorCorrector>();

    // (Re)create the single-sample linear copy texture + its FBO on size change.
    if (m_ccLinearTex == 0 || m_ccLinearW != width || m_ccLinearH != height) {
        if (m_ccLinearTex == 0) glGenTextures(1, &m_ccLinearTex);
        glBindTexture(GL_TEXTURE_2D, m_ccLinearTex);
        glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA16F, width, height, 0,
                     GL_RGBA, GL_FLOAT, nullptr);
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
        glBindTexture(GL_TEXTURE_2D, 0);

        if (m_ccLinearFBO == 0) glGenFramebuffers(1, &m_ccLinearFBO);
        glBindFramebuffer(GL_FRAMEBUFFER, m_ccLinearFBO);
        glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0,
                               GL_TEXTURE_2D, m_ccLinearTex, 0);
        glBindFramebuffer(GL_FRAMEBUFFER, 0);
        m_ccLinearW = width;
        m_ccLinearH = height;
    }

    // Resolve MSAA (no-op otherwise) so the color attachment holds the linear
    // single-sample image, then copy it into m_ccLinearTex — sampling and
    // writing the same texture in one pass is illegal, hence the copy.
    m_drawTarget->Resolve();
    glBindFramebuffer(GL_READ_FRAMEBUFFER, m_drawTarget->GetFramebufferId());
    glBindFramebuffer(GL_DRAW_FRAMEBUFFER, m_ccLinearFBO);
    glBlitFramebuffer(0, 0, width, height, 0, 0, width, height,
                      GL_COLOR_BUFFER_BIT, GL_NEAREST);

    // Draw the corrected result back into the draw-target render FBO (the MSAA
    // FBO when multisampling, so it resolves together with the overlays drawn
    // on top of it afterwards).
    GLuint dstFbo = m_drawTarget->HasMSAA()
        ? m_drawTarget->GetFramebufferMSId()
        : m_drawTarget->GetFramebufferId();
    glBindFramebuffer(GL_FRAMEBUFFER, dstFbo);
    glViewport(0, 0, width, height);

    ViewportColorCorrector::Mode mode =
        (m_colorCorrectionMode == ColorCorrectionMode::OpenColorIO)
            ? ViewportColorCorrector::Mode::OpenColorIO
            : ViewportColorCorrector::Mode::sRGB;

    m_colorCorrector->Apply(m_ccLinearTex, mode,
                            m_ocioDisplay, m_ocioView,
                            m_ocioColorSpace, m_ocioLook);
}

// ===========================================================================
// Output
// ===========================================================================

uint32_t UsdSceneRenderer::GetColorTextureID()
{
    if (!m_drawTarget) return 0;
    // Resolve MSAA → regular texture before the caller samples it.
    // Called after all overlay draws (axis, bboxes, camera wireframes) so
    // every layer of MSAA-rendered content is included in the resolve.
    m_drawTarget->Resolve();
    auto att = m_drawTarget->GetAttachment("color");
    return att ? static_cast<uint32_t>(att->GetGlTextureName()) : 0;
}

bool UsdSceneRenderer::CaptureFrame(std::vector<uint8_t>& outRGBA)
{
    if (!m_drawTarget || m_lastRenderWidth <= 0 || m_lastRenderHeight <= 0)
        return false;

    int w = m_lastRenderWidth;
    int h = m_lastRenderHeight;

    // Resolve MSAA so we read from the non-multisampled colour attachment.
    m_drawTarget->Resolve();

    outRGBA.resize(static_cast<size_t>(w) * h * 4);

    // Save/restore the read-framebuffer binding so we don't upset ImGui.
    GLint prevFbo = 0;
    glGetIntegerv(GL_READ_FRAMEBUFFER_BINDING, &prevFbo);
    glBindFramebuffer(GL_READ_FRAMEBUFFER, m_drawTarget->GetFramebufferId());
    glReadPixels(0, 0, w, h, GL_RGBA, GL_UNSIGNED_BYTE, outRGBA.data());
    glBindFramebuffer(GL_READ_FRAMEBUFFER, static_cast<GLuint>(prevFbo));

    // OpenGL reads bottom-up; flip to top-down.
    int stride = w * 4;
    std::vector<uint8_t> row(stride);
    for (int y = 0; y < h / 2; ++y) {
        uint8_t* top = outRGBA.data() + y * stride;
        uint8_t* bot = outRGBA.data() + (h - 1 - y) * stride;
        std::copy(top, top + stride, row.data());
        std::copy(bot, bot + stride, top);
        std::copy(row.data(), row.data() + stride, bot);
    }

    return true;
}

// ===========================================================================
// Axis Overlay (stageView.DrawAxis port)
// ===========================================================================

void UsdSceneRenderer::InitAxisResources()
{
    if (m_axisProgram != 0) return;

    GLuint vs = CompileShader(kAxisVS, GL_VERTEX_SHADER);
    GLuint fs = CompileShader(kAxisFS, GL_FRAGMENT_SHADER);
    if (!vs || !fs) { glDeleteShader(vs); glDeleteShader(fs); return; }
    m_axisProgram = LinkProgram(vs, fs);
    if (!m_axisProgram) return;

    m_axisUniformMVP   = glGetUniformLocation(m_axisProgram, "mvpMatrix");
    m_axisUniformColor = glGetUniformLocation(m_axisProgram, "color");

    // 3 line segments: X=(1,0,0), Y=(0,1,0), Z=(0,0,1) from origin (0,0,0)
    float axisVerts[] = {
        1,0,0,  0,0,0,
        0,1,0,  0,0,0,
        0,0,1,  0,0,0
    };

    glGenVertexArrays(1, &m_axisVAO);
    glGenBuffers(1, &m_axisVBO);
    glBindVertexArray(m_axisVAO);
    glBindBuffer(GL_ARRAY_BUFFER, m_axisVBO);
    glBufferData(GL_ARRAY_BUFFER, sizeof(axisVerts), axisVerts, GL_STATIC_DRAW);
    glEnableVertexAttribArray(0);
    glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, 0, nullptr);
    glBindVertexArray(0);
    LOG_INFO("Axis resources initialized.");
}

void UsdSceneRenderer::DestroyAxisResources()
{
    if (m_axisVAO)     { glDeleteVertexArrays(1, &m_axisVAO);  m_axisVAO = 0; }
    if (m_axisVBO)     { glDeleteBuffers(1, &m_axisVBO);        m_axisVBO = 0; }
    if (m_axisProgram) { glDeleteProgram(m_axisProgram);        m_axisProgram = 0; }
}

void UsdSceneRenderer::DrawAxis(
    const pxr::GfMatrix4d& viewProjMatrix, double cameraDist)
{
    if (!m_axisProgram || !m_axisVAO) return;
    if (!m_drawTarget) return;

    // Overlays are called after Render() has unbound the FBO.
    // Re-bind so we draw into the offscreen texture, not the default framebuffer.
    pxr::GfVec2i sz = m_drawTarget->GetSize();
    m_drawTarget->Bind();
    glViewport(0, 0, sz[0], sz[1]);

    // Scale the gizmo to stay roughly fixed in screen space (stageView: dist/20)
    pxr::GfMatrix4f mvp =
        pxr::GfMatrix4f(1.0f).SetScale(static_cast<float>(cameraDist / 20.0))
        * pxr::GfMatrix4f(viewProjMatrix);

    glUseProgram(m_axisProgram);
    glBindVertexArray(m_axisVAO);

    glUniformMatrix4fv(m_axisUniformMVP, 1, GL_TRUE, mvp.GetArray());

    GLboolean prevDepthMask;
    glGetBooleanv(GL_DEPTH_WRITEMASK, &prevDepthMask);
    glDepthMask(GL_FALSE);
    GLboolean prevDepthTest  = glIsEnabled(GL_DEPTH_TEST);
    GLboolean prevLineSmooth = glIsEnabled(GL_LINE_SMOOTH);
    glEnable(GL_DEPTH_TEST);
    if (m_aaEnabled) { glEnable(GL_LINE_SMOOTH); glHint(GL_LINE_SMOOTH_HINT, GL_NICEST); }

    // X axis: red
    glUniform4f(m_axisUniformColor, 1, 0, 0, 1);
    glDrawArrays(GL_LINES, 0, 2);
    // Y axis: green
    glUniform4f(m_axisUniformColor, 0, 1, 0, 1);
    glDrawArrays(GL_LINES, 2, 2);
    // Z axis: blue
    glUniform4f(m_axisUniformColor, 0, 0, 1, 1);
    glDrawArrays(GL_LINES, 4, 2);

    glDepthMask(prevDepthMask);
    if (!prevDepthTest)  glDisable(GL_DEPTH_TEST);
    if (prevLineSmooth)  glEnable(GL_LINE_SMOOTH);
    else                 glDisable(GL_LINE_SMOOTH);

    glBindVertexArray(0);
    glUseProgram(0);
    m_drawTarget->Unbind();
}

// ===========================================================================
// Grid Overlay
// ===========================================================================

void UsdSceneRenderer::InitGridResources()
{
    if (m_gridVAO != 0) return;

    // Grid reuses the axis shader program — both use vec3 position + uniform MVP/color.
    // The VAO only needs a position attribute.
    glGenVertexArrays(1, &m_gridVAO);
    glGenBuffers(1, &m_gridVBO);
    glBindVertexArray(m_gridVAO);
    glBindBuffer(GL_ARRAY_BUFFER, m_gridVBO);
    glEnableVertexAttribArray(0);
    glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, 3 * sizeof(float), nullptr);
    glBindVertexArray(0);

    LOG_INFO("Grid VAO/VBO created.");
    RebuildGridVBO();
}

void UsdSceneRenderer::RebuildGridVBO()
{
    // Can only rebuild once the VAO/VBO have been created.
    if (m_gridVAO == 0 || m_gridVBO == 0) return;

    // -------------------------------------------------------------------------
    // Build all grid line vertices grouped by category so we can draw each
    // group with a different colour in a single draw call.
    //
    //  Group A — minor lines  (1-unit spacing, excluding multiples of 10 and 0)
    //  Group B — major lines  (10-unit spacing, excluding 0)
    //  Group C — "A"-axis     (the axis along direction a, at b = 0  → red   )
    //  Group D — "B"-axis     (the axis along direction b, at a = 0  → blue/green)
    //
    // For Y-up: ground plane = XZ (Y=0), a = X axis, b = Z axis.
    // For Z-up: ground plane = XY (Z=0), a = X axis, b = Y axis.
    // -------------------------------------------------------------------------

    const float H   = m_gridHalfSize;          // half-extent (default 50)
    const int   N   = static_cast<int>(H);     // integer half-extent (e.g. 50)
    const int   maj = 10;                      // major-line interval

    // Helper: given a 2-D coordinate pair (a, b) on the ground plane,
    // return the 3-D world position based on the current up-axis.
    auto toWorld = [&](float a, float b) -> std::array<float,3> {
        if (m_stageIsZup)
            return { a, b, 0.0f };   // XY plane at Z=0
        else
            return { a, 0.0f, b };   // XZ plane at Y=0
    };

    std::vector<float> minor_verts, major_verts, axisA_verts, axisB_verts;

    // Lines parallel to the A-axis (at fixed b values):
    for (int bi = -N; bi <= N; ++bi) {
        float b = static_cast<float>(bi);
        auto  p0 = toWorld(-H, b);
        auto  p1 = toWorld( H, b);

        if (bi == 0) {
            // A-axis (red)
            axisA_verts.insert(axisA_verts.end(), p0.begin(), p0.end());
            axisA_verts.insert(axisA_verts.end(), p1.begin(), p1.end());
        } else if (bi % maj == 0) {
            // Major line
            major_verts.insert(major_verts.end(), p0.begin(), p0.end());
            major_verts.insert(major_verts.end(), p1.begin(), p1.end());
        } else {
            // Minor line
            minor_verts.insert(minor_verts.end(), p0.begin(), p0.end());
            minor_verts.insert(minor_verts.end(), p1.begin(), p1.end());
        }
    }

    // Lines parallel to the B-axis (at fixed a values):
    for (int ai = -N; ai <= N; ++ai) {
        float a = static_cast<float>(ai);
        auto  p0 = toWorld(a, -H);
        auto  p1 = toWorld(a,  H);

        if (ai == 0) {
            // B-axis (blue / green)
            axisB_verts.insert(axisB_verts.end(), p0.begin(), p0.end());
            axisB_verts.insert(axisB_verts.end(), p1.begin(), p1.end());
        } else if (ai % maj == 0) {
            major_verts.insert(major_verts.end(), p0.begin(), p0.end());
            major_verts.insert(major_verts.end(), p1.begin(), p1.end());
        } else {
            minor_verts.insert(minor_verts.end(), p0.begin(), p0.end());
            minor_verts.insert(minor_verts.end(), p1.begin(), p1.end());
        }
    }

    // Pack into one contiguous buffer: [minor | major | axisA | axisB]
    std::vector<float> all;
    all.reserve(minor_verts.size() + major_verts.size() +
                axisA_verts.size() + axisB_verts.size());

    auto floatsToVerts = [](size_t f) { return static_cast<GLint>(f / 3); };

    m_gridMinorFirst = 0;
    m_gridMinorCount = floatsToVerts(minor_verts.size());
    all.insert(all.end(), minor_verts.begin(), minor_verts.end());

    m_gridMajorFirst = m_gridMinorFirst + m_gridMinorCount;
    m_gridMajorCount = floatsToVerts(major_verts.size());
    all.insert(all.end(), major_verts.begin(), major_verts.end());

    m_gridAxisAFirst = m_gridMajorFirst + m_gridMajorCount;
    m_gridAxisACount = floatsToVerts(axisA_verts.size());
    all.insert(all.end(), axisA_verts.begin(), axisA_verts.end());

    m_gridAxisBFirst = m_gridAxisAFirst + m_gridAxisACount;
    m_gridAxisBCount = floatsToVerts(axisB_verts.size());
    all.insert(all.end(), axisB_verts.begin(), axisB_verts.end());

    glBindBuffer(GL_ARRAY_BUFFER, m_gridVBO);
    glBufferData(GL_ARRAY_BUFFER,
                 static_cast<GLsizeiptr>(all.size() * sizeof(float)),
                 all.data(), GL_STATIC_DRAW);
    glBindBuffer(GL_ARRAY_BUFFER, 0);

    LOG_INFO("Grid VBO rebuilt: "
        + std::to_string(m_gridMinorCount / 2) + " minor lines, "
        + std::to_string(m_gridMajorCount / 2) + " major lines, "
        + std::string(m_stageIsZup ? "Z-up" : "Y-up"));
}

void UsdSceneRenderer::DestroyGridResources()
{
    if (m_gridVAO) { glDeleteVertexArrays(1, &m_gridVAO); m_gridVAO = 0; }
    if (m_gridVBO) { glDeleteBuffers(1, &m_gridVBO);       m_gridVBO = 0; }
    // No separate grid program — the axis program is destroyed in DestroyAxisResources().
}

void UsdSceneRenderer::RenderGrid(int /*width*/, int /*height*/)
{
    if (!m_gridVAO || !m_axisProgram) return;

    // Grid reuses the axis shader.
    // MVP = view * proj  (no extra scale — grid is already in world units).
    pxr::GfMatrix4f mvp(m_viewMatrix * m_projMatrix);

    // Helper: enable/disable GL_LINE_SMOOTH depending on the AA setting.
    auto setLineAA = [&](bool enable) {
        if (enable) {
            glEnable(GL_LINE_SMOOTH);
            glHint(GL_LINE_SMOOTH_HINT, GL_NICEST);
        } else {
            glDisable(GL_LINE_SMOOTH);
        }
    };

    // Save relevant GL state.
    GLboolean prevDepthMask;
    GLboolean prevDepthTest = glIsEnabled(GL_DEPTH_TEST);
    GLboolean prevBlend     = glIsEnabled(GL_BLEND);
    GLboolean prevLineSmooth= glIsEnabled(GL_LINE_SMOOTH);
    GLint     prevDepthFunc, prevBlendSrc, prevBlendDst;
    glGetBooleanv(GL_DEPTH_WRITEMASK, &prevDepthMask);
    glGetIntegerv(GL_DEPTH_FUNC,      &prevDepthFunc);
    glGetIntegerv(GL_BLEND_SRC_ALPHA, &prevBlendSrc);
    glGetIntegerv(GL_BLEND_DST_ALPHA, &prevBlendDst);

    glDepthMask(GL_FALSE);
    glEnable(GL_DEPTH_TEST);
    glDepthFunc(GL_LEQUAL);
    glEnable(GL_BLEND);
    glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA);
    setLineAA(m_aaEnabled);

    glUseProgram(m_axisProgram);
    glUniformMatrix4fv(m_axisUniformMVP, 1, GL_TRUE, mvp.GetArray());
    glBindVertexArray(m_gridVAO);

    // --- Minor lines — dark grey ---
    glLineWidth(1.0f);
    glUniform4f(m_axisUniformColor, 0.35f, 0.35f, 0.35f, 1.0f);
    if (m_gridMinorCount > 0)
        glDrawArrays(GL_LINES, m_gridMinorFirst, m_gridMinorCount);

    // --- Major lines — medium grey ---
    glUniform4f(m_axisUniformColor, 0.52f, 0.52f, 0.52f, 1.0f);
    if (m_gridMajorCount > 0)
        glDrawArrays(GL_LINES, m_gridMajorFirst, m_gridMajorCount);

    // --- A-axis (at b=0): X axis — red ---
    glLineWidth(m_aaEnabled ? 1.5f : 1.0f);
    glUniform4f(m_axisUniformColor, 0.62f, 0.28f, 0.28f, 1.0f);
    if (m_gridAxisACount > 0)
        glDrawArrays(GL_LINES, m_gridAxisAFirst, m_gridAxisACount);

    // --- B-axis (at a=0): Z-axis (blue) for Y-up, Y-axis (green) for Z-up ---
    if (m_stageIsZup)
        glUniform4f(m_axisUniformColor, 0.28f, 0.55f, 0.28f, 1.0f);   // green (Y)
    else
        glUniform4f(m_axisUniformColor, 0.28f, 0.28f, 0.62f, 1.0f);   // blue  (Z)
    if (m_gridAxisBCount > 0)
        glDrawArrays(GL_LINES, m_gridAxisBFirst, m_gridAxisBCount);

    glBindVertexArray(0);
    glUseProgram(0);

    // Restore GL state.
    glLineWidth(1.0f);
    glDepthMask(prevDepthMask);
    glDepthFunc(static_cast<GLenum>(prevDepthFunc));
    glBlendFunc(static_cast<GLenum>(prevBlendSrc),
                static_cast<GLenum>(prevBlendDst));
    if (!prevDepthTest)  glDisable(GL_DEPTH_TEST);
    if (!prevBlend)      glDisable(GL_BLEND);
    if (prevLineSmooth)  glEnable(GL_LINE_SMOOTH);
    else                 glDisable(GL_LINE_SMOOTH);
}

// ===========================================================================
// Bounding Box Overlay
// ===========================================================================

// Reuse the same simple line shader as the axis gizmo.
static const char* kBBoxVS = R"(#version 130
in vec3 position;
uniform mat4 mvpMatrix;
void main() { gl_Position = vec4(position, 1.0) * mvpMatrix; }
)";

static const char* kBBoxFS = R"(#version 130
uniform vec4 color;
out vec4 outColor;
void main() { outColor = color; }
)";

void UsdSceneRenderer::InitBBoxResources()
{
    if (m_bboxProgram != 0) return;

    GLuint vs = CompileShader(kBBoxVS, GL_VERTEX_SHADER);
    GLuint fs = CompileShader(kBBoxFS, GL_FRAGMENT_SHADER);
    if (!vs || !fs) { glDeleteShader(vs); glDeleteShader(fs); return; }
    m_bboxProgram = LinkProgram(vs, fs);
    if (!m_bboxProgram) return;

    m_bboxUniformMVP   = glGetUniformLocation(m_bboxProgram, "mvpMatrix");
    m_bboxUniformColor = glGetUniformLocation(m_bboxProgram, "color");

    // Allocate a VAO/VBO for 24 vertices (12 edges × 2 endpoints) — data uploaded dynamically.
    glGenVertexArrays(1, &m_bboxVAO);
    glGenBuffers(1, &m_bboxVBO);
    glBindVertexArray(m_bboxVAO);
    glBindBuffer(GL_ARRAY_BUFFER, m_bboxVBO);
    glBufferData(GL_ARRAY_BUFFER, 24 * 3 * sizeof(float), nullptr, GL_DYNAMIC_DRAW);
    glEnableVertexAttribArray(0);
    glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, 0, nullptr);
    glBindVertexArray(0);
    LOG_INFO("BBox resources initialized.");
}

void UsdSceneRenderer::DestroyBBoxResources()
{
    if (m_bboxVAO)     { glDeleteVertexArrays(1, &m_bboxVAO);  m_bboxVAO = 0; }
    if (m_bboxVBO)     { glDeleteBuffers(1, &m_bboxVBO);        m_bboxVBO = 0; }
    if (m_bboxProgram) { glDeleteProgram(m_bboxProgram);        m_bboxProgram = 0; }
}

// ──────────────────────────────────────────────────────────────────────────────
// Camera Wireframe GL resources
// Uses the same simple line shader as the axis / bbox overlays.
// A separate VAO/VBO avoids any interaction with the bbox dynamic VBO.
// ──────────────────────────────────────────────────────────────────────────────
void UsdSceneRenderer::InitCamWireResources()
{
    if (m_camWireVAO != 0) return;
    if (!m_bboxProgram) return;  // reuse the already-compiled bbox shader program

    glGenVertexArrays(1, &m_camWireVAO);
    glGenBuffers(1, &m_camWireVBO);

    glBindVertexArray(m_camWireVAO);
    glBindBuffer(GL_ARRAY_BUFFER, m_camWireVBO);
    // Allocate an initial budget; DrawCameraWireframes reallocates per frame.
    glBufferData(GL_ARRAY_BUFFER, 256 * 3 * sizeof(float), nullptr, GL_DYNAMIC_DRAW);
    glEnableVertexAttribArray(0);
    glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, 0, nullptr);
    glBindVertexArray(0);
    LOG_INFO("Camera wireframe GL resources initialized.");
}

void UsdSceneRenderer::DestroyCamWireResources()
{
    if (m_camWireVAO) { glDeleteVertexArrays(1, &m_camWireVAO); m_camWireVAO = 0; }
    if (m_camWireVBO) { glDeleteBuffers(1, &m_camWireVBO);       m_camWireVBO = 0; }
}

void UsdSceneRenderer::InitLightWireResources()
{
    if (m_lightWireVAO != 0) return;
    if (!m_bboxProgram) return;  // reuse the already-compiled bbox shader program

    glGenVertexArrays(1, &m_lightWireVAO);
    glGenBuffers(1, &m_lightWireVBO);

    glBindVertexArray(m_lightWireVAO);
    glBindBuffer(GL_ARRAY_BUFFER, m_lightWireVBO);
    // Sphere/cylinder gizmos use several circles; budget more than the camera VBO.
    glBufferData(GL_ARRAY_BUFFER, 1024 * 3 * sizeof(float), nullptr, GL_DYNAMIC_DRAW);
    glEnableVertexAttribArray(0);
    glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, 0, nullptr);
    glBindVertexArray(0);
    LOG_INFO("Light wireframe GL resources initialized.");
}

void UsdSceneRenderer::DestroyLightWireResources()
{
    if (m_lightWireVAO) { glDeleteVertexArrays(1, &m_lightWireVAO); m_lightWireVAO = 0; }
    if (m_lightWireVBO) { glDeleteBuffers(1, &m_lightWireVBO);       m_lightWireVBO = 0; }
}

void UsdSceneRenderer::DrawBox(
    const pxr::GfRange3d&  range,
    const pxr::GfMatrix4f& mvp)
{
    if (range.IsEmpty()) return;

    pxr::GfVec3d mn = range.GetMin();
    pxr::GfVec3d mx = range.GetMax();

    // 8 corners of the AABB
    float verts[24][3] = {
        // Bottom face
        { (float)mn[0], (float)mn[1], (float)mn[2] }, { (float)mx[0], (float)mn[1], (float)mn[2] },
        { (float)mx[0], (float)mn[1], (float)mn[2] }, { (float)mx[0], (float)mn[1], (float)mx[2] },
        { (float)mx[0], (float)mn[1], (float)mx[2] }, { (float)mn[0], (float)mn[1], (float)mx[2] },
        { (float)mn[0], (float)mn[1], (float)mx[2] }, { (float)mn[0], (float)mn[1], (float)mn[2] },
        // Top face
        { (float)mn[0], (float)mx[1], (float)mn[2] }, { (float)mx[0], (float)mx[1], (float)mn[2] },
        { (float)mx[0], (float)mx[1], (float)mn[2] }, { (float)mx[0], (float)mx[1], (float)mx[2] },
        { (float)mx[0], (float)mx[1], (float)mx[2] }, { (float)mn[0], (float)mx[1], (float)mx[2] },
        { (float)mn[0], (float)mx[1], (float)mx[2] }, { (float)mn[0], (float)mx[1], (float)mn[2] },
        // Vertical pillars
        { (float)mn[0], (float)mn[1], (float)mn[2] }, { (float)mn[0], (float)mx[1], (float)mn[2] },
        { (float)mx[0], (float)mn[1], (float)mn[2] }, { (float)mx[0], (float)mx[1], (float)mn[2] },
        { (float)mx[0], (float)mn[1], (float)mx[2] }, { (float)mx[0], (float)mx[1], (float)mx[2] },
        { (float)mn[0], (float)mn[1], (float)mx[2] }, { (float)mn[0], (float)mx[1], (float)mx[2] },
    };

    glBindBuffer(GL_ARRAY_BUFFER, m_bboxVBO);
    glBufferSubData(GL_ARRAY_BUFFER, 0, sizeof(verts), verts);

    glUniformMatrix4fv(m_bboxUniformMVP, 1, GL_TRUE, mvp.GetArray());
    glBindVertexArray(m_bboxVAO);
    glDrawArrays(GL_LINES, 0, 24);
    glBindVertexArray(0);
}

void UsdSceneRenderer::DrawBoundingBoxes(
    const pxr::SdfPathVector& selectedPaths,
    const pxr::GfMatrix4d&   viewProjMatrix)
{
    if (m_bboxMode == BBoxMode::None) return;
    if (!m_bboxProgram || !m_bboxVAO) return;
    if (!m_stage || selectedPaths.empty()) return;
    if (!m_drawTarget) return;

    // Re-bind the offscreen FBO so we draw into the texture, not the default framebuffer.
    pxr::GfVec2i sz = m_drawTarget->GetSize();
    m_drawTarget->Bind();
    glViewport(0, 0, sz[0], sz[1]);

    // Compute world bboxes via UsdGeomBBoxCache
    pxr::TfTokenVector purposes = {
        pxr::UsdGeomTokens->default_, pxr::UsdGeomTokens->proxy };
    pxr::UsdGeomBBoxCache bboxCache(
        m_currentTime, purposes, /*useExtentsHint=*/true);

    // Save/restore GL state
    GLboolean prevDepthMask;
    glGetBooleanv(GL_DEPTH_WRITEMASK, &prevDepthMask);
    GLboolean prevDepthTest = glIsEnabled(GL_DEPTH_TEST);
    GLboolean prevBlend     = glIsEnabled(GL_BLEND);

    glDepthMask(GL_FALSE);
    glEnable(GL_DEPTH_TEST);
    glEnable(GL_BLEND);
    glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA);

    glUseProgram(m_bboxProgram);
    glUniform4f(m_bboxUniformColor,
        m_bboxColor[0], m_bboxColor[1], m_bboxColor[2], m_bboxColor[3]);

    pxr::GfMatrix4f vp = pxr::GfMatrix4f(viewProjMatrix);

    if (m_bboxMode == BBoxMode::PerObject) {
        for (const auto& path : selectedPaths) {
            pxr::UsdPrim prim = m_stage->GetPrimAtPath(path);
            if (!prim) continue;
            pxr::GfBBox3d bbox = bboxCache.ComputeWorldBound(prim);
            DrawBox(bbox.ComputeAlignedRange(), vp);
        }
    } else {
        // AllSelection: one combined AABB around all selected prims
        pxr::GfRange3d combined;
        for (const auto& path : selectedPaths) {
            pxr::UsdPrim prim = m_stage->GetPrimAtPath(path);
            if (!prim) continue;
            pxr::GfBBox3d bbox = bboxCache.ComputeWorldBound(prim);
            combined.UnionWith(bbox.ComputeAlignedRange());
        }
        DrawBox(combined, vp);
    }

    glDepthMask(prevDepthMask);
    if (!prevDepthTest) glDisable(GL_DEPTH_TEST);
    if (!prevBlend)     glDisable(GL_BLEND);

    glUseProgram(0);
    m_drawTarget->Unbind();
}

// ===========================================================================
// Draw-target FBO helpers
// ===========================================================================

void UsdSceneRenderer::BindDrawTarget()
{
    if (!m_drawTarget) return;
    pxr::GfVec2i sz = m_drawTarget->GetSize();
    m_drawTarget->Bind();
    glViewport(0, 0, sz[0], sz[1]);
}

void UsdSceneRenderer::UnbindDrawTarget()
{
    if (m_drawTarget) m_drawTarget->Unbind();
}

// ===========================================================================
// Camera Wireframe helpers
// ===========================================================================

// Project a world-space point to absolute screen coords (imagePosX/Y + pixel offsets).
// Returns false if the point is behind the near plane (w ≤ 0).
bool UsdSceneRenderer::WorldToScreen(const pxr::GfVec3d&    world,
                                      const pxr::GfMatrix4d& vp,
                                      int viewW, int viewH,
                                      float imagePosX, float imagePosY,
                                      float& outX, float& outY)
{
    double cx = vp[0][0]*world[0] + vp[1][0]*world[1] + vp[2][0]*world[2] + vp[3][0];
    double cy = vp[0][1]*world[0] + vp[1][1]*world[1] + vp[2][1]*world[2] + vp[3][1];
    double cw = vp[0][3]*world[0] + vp[1][3]*world[1] + vp[2][3]*world[2] + vp[3][3];
    if (cw <= 0.0) return false;
    double invW = 1.0 / cw;
    outX = imagePosX + static_cast<float>(( cx * invW + 1.0) * 0.5 * viewW);
    outY = imagePosY + static_cast<float>((1.0 - cy * invW) * 0.5 * viewH);
    return true;
}

float UsdSceneRenderer::PointToSegmentDist(float px, float py,
                                            float ax, float ay,
                                            float bx, float by)
{
    float dx = bx - ax, dy = by - ay;
    float lenSq = dx*dx + dy*dy;
    if (lenSq < 1e-6f) {
        float ex = px - ax, ey = py - ay;
        return std::sqrt(ex*ex + ey*ey);
    }
    float t  = std::max(0.f, std::min(1.f, ((px-ax)*dx + (py-ay)*dy) / lenSq));
    float cx2 = ax + t*dx - px;
    float cy2 = ay + t*dy - py;
    return std::sqrt(cx2*cx2 + cy2*cy2);
}

// Build camera wireframe line segments in world space.
// Geometry: body box (12 edges) + frustum pyramid (4 lines to near quad) + up arrow (1 line).
void UsdSceneRenderer::BuildCameraWireframeLines(const pxr::GfCamera& gfCam,
                                                  double                scale,
                                                  std::vector<float>&   outVerts)
{
    // Camera transform: rows are camera axes in world space.
    // Column convention: cameraToWorld = camMat (USD row-vector).
    pxr::GfMatrix4d camToWorld = gfCam.GetTransform();

    // Camera origin in world space
    pxr::GfVec3d origin(camToWorld[3][0], camToWorld[3][1], camToWorld[3][2]);

    // Camera axes (columns of the rotation part, row-major: row i = axis i of camera)
    // In USD GfMatrix4d row-vector convention: world = local * M
    // So camera right = row 0, up = row 1, -forward = row 2
    pxr::GfVec3d right(  camToWorld[0][0], camToWorld[0][1], camToWorld[0][2]);
    pxr::GfVec3d up(     camToWorld[1][0], camToWorld[1][1], camToWorld[1][2]);
    pxr::GfVec3d forward(camToWorld[2][0], camToWorld[2][1], camToWorld[2][2]);
    // Note: USD cameras look down -Z in local space, so forward here is the +Z local = backward in view.
    // The camera shoots along -forward (local -Z).
    pxr::GfVec3d lookDir = -forward;  // world-space look direction

    // ---- Body box ----
    double bh = scale * 0.10;  // half-size
    // 8 corners of the body box in world space
    pxr::GfVec3d c[8];
    for (int xi = -1; xi <= 1; xi += 2)
    for (int yi = -1; yi <= 1; yi += 2)
    for (int zi = -1; zi <= 1; zi += 2) {
        int idx = ((xi+1)/2) | (((yi+1)/2) << 1) | (((zi+1)/2) << 2);
        c[idx] = origin + right*bh*xi + up*bh*yi + lookDir*bh*zi;
    }
    // 12 edges of the box: connect corners whose indices differ by exactly one bit
    static const int kEdges[12][2] = {
        {0,1},{2,3},{4,5},{6,7},  // X edges
        {0,2},{1,3},{4,6},{5,7},  // Y edges
        {0,4},{1,5},{2,6},{3,7}   // Z edges
    };
    auto push = [&](const pxr::GfVec3d& a, const pxr::GfVec3d& b) {
        outVerts.push_back(static_cast<float>(a[0]));
        outVerts.push_back(static_cast<float>(a[1]));
        outVerts.push_back(static_cast<float>(a[2]));
        outVerts.push_back(static_cast<float>(b[0]));
        outVerts.push_back(static_cast<float>(b[1]));
        outVerts.push_back(static_cast<float>(b[2]));
    };
    for (auto& e : kEdges) push(c[e[0]], c[e[1]]);

    // ---- Frustum pyramid (4 lines from origin to near quad corners) ----
    // Use perspective projection: half-widths at near depth proportional to aperture/focalLen
    double nearDepth = std::max(gfCam.GetClippingRange().GetMin(), 0.01f);
    // Cap display depth at scale so the pyramid isn't enormous
    double dispDepth = std::min(nearDepth * 3.0, scale * 1.5);
    dispDepth = std::max(dispDepth, scale * 0.3);

    double hApert = static_cast<double>(gfCam.GetHorizontalAperture()) * 0.5;
    double vApert = static_cast<double>(gfCam.GetVerticalAperture())   * 0.5;
    double focalL = static_cast<double>(gfCam.GetFocalLength());
    if (focalL < 1e-6) focalL = 50.0;  // fallback

    // Scale aperture to the display depth using similar triangles
    double hw = hApert / focalL * dispDepth;
    double hv = vApert / focalL * dispDepth;

    // 4 corners of the near quad at dispDepth along lookDir
    pxr::GfVec3d nearCentre = origin + lookDir * dispDepth;
    pxr::GfVec3d nBL = nearCentre - right*hw - up*hv;
    pxr::GfVec3d nBR = nearCentre + right*hw - up*hv;
    pxr::GfVec3d nTL = nearCentre - right*hw + up*hv;
    pxr::GfVec3d nTR = nearCentre + right*hw + up*hv;

    // Lines from origin to each corner (pyramid edges)
    push(origin, nBL);
    push(origin, nBR);
    push(origin, nTL);
    push(origin, nTR);
    // Near quad rectangle
    push(nBL, nBR);
    push(nBR, nTR);
    push(nTR, nTL);
    push(nTL, nBL);

    // ---- Up arrow ----
    pxr::GfVec3d arrowBase = origin + up * bh;
    pxr::GfVec3d arrowTip  = origin + up * (bh + scale * 0.18);
    push(arrowBase, arrowTip);
}

// ===========================================================================
// DrawCameraWireframes
// ===========================================================================

void UsdSceneRenderer::DrawCameraWireframes(
    pxr::UsdStageRefPtr        stage,
    const pxr::SdfPathVector&  selectedPaths,
    const pxr::SdfPath&        activeCameraPath,
    const pxr::GfMatrix4d&     viewProjMatrix,
    double                     viewportCameraDist)
{
    if (!stage) return;
    if (!m_bboxProgram || !m_camWireVAO || !m_camWireVBO) return;
    if (!m_drawTarget) return;

    // Rebuild camera path list fresh each call.
    m_cachedCameraPaths.clear();
    for (const pxr::UsdPrim& prim : stage->Traverse()) {
        if (prim.IsA<pxr::UsdGeomCamera>())
            m_cachedCameraPaths.push_back(prim.GetPath());
    }
    if (m_cachedCameraPaths.empty()) return;

    double rawScale = viewportCameraDist * 0.12;
    double scale    = std::min(rawScale, 50.0);
    scale = std::max(scale, 0.5);    // minimum visible size

    // Build set of selected paths for O(1) lookup
    std::unordered_set<std::string> selectedSet;
    for (const auto& p : selectedPaths) selectedSet.insert(p.GetString());

    // Re-bind FBO
    pxr::GfVec2i sz = m_drawTarget->GetSize();
    m_drawTarget->Bind();
    glViewport(0, 0, sz[0], sz[1]);

    // Save GL state
    GLboolean prevDepthTest = glIsEnabled(GL_DEPTH_TEST);
    GLboolean prevDepthMask;
    glGetBooleanv(GL_DEPTH_WRITEMASK, &prevDepthMask);
    GLboolean prevBlend = glIsEnabled(GL_BLEND);

    // Camera wireframes draw on top of everything (like the axis gizmo) so they
    // are always visible regardless of scene geometry at the camera location.
    glDisable(GL_DEPTH_TEST);
    glDepthMask(GL_FALSE);
    glEnable(GL_BLEND);
    glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA);
    if (m_aaEnabled) glEnable(GL_LINE_SMOOTH);

    glUseProgram(m_bboxProgram);

    pxr::GfMatrix4f vp = pxr::GfMatrix4f(viewProjMatrix);

    for (const auto& camPath : m_cachedCameraPaths) {
        pxr::UsdPrim prim = stage->GetPrimAtPath(camPath);
        if (!prim) continue;

        pxr::UsdGeomCamera usdCam(prim);
        pxr::GfCamera gfCam = usdCam.GetCamera(m_currentTime);

        // Determine colour
        bool isSelected = (selectedSet.count(camPath.GetString()) > 0);
        bool isActive   = (camPath == activeCameraPath && !activeCameraPath.IsEmpty());

        pxr::GfVec4f col;
        if      (isSelected) col = pxr::GfVec4f(1.0f, 0.75f, 0.10f, 1.0f);
        else if (isActive)   col = pxr::GfVec4f(0.2f, 0.90f, 1.00f, 1.0f);
        else                 col = pxr::GfVec4f(0.65f, 0.85f, 1.00f, 0.90f); // light blue

        glUniform4f(m_bboxUniformColor, col[0], col[1], col[2], col[3]);

        std::vector<float> verts;
        BuildCameraWireframeLines(gfCam, scale, verts);
        if (verts.empty()) continue;

        int vertCount = static_cast<int>(verts.size() / 3);

        // Upload to the dedicated camera wireframe VBO
        glBindBuffer(GL_ARRAY_BUFFER, m_camWireVBO);
        glBufferData(GL_ARRAY_BUFFER,
                     static_cast<GLsizeiptr>(verts.size() * sizeof(float)),
                     verts.data(), GL_DYNAMIC_DRAW);

        glUniformMatrix4fv(m_bboxUniformMVP, 1, GL_TRUE, vp.GetArray());

        glBindVertexArray(m_camWireVAO);
        glDrawArrays(GL_LINES, 0, vertCount);
        glBindVertexArray(0);
    }

    // Restore GL state
    if (m_aaEnabled) glDisable(GL_LINE_SMOOTH);
    if (prevDepthTest)  glEnable(GL_DEPTH_TEST); else glDisable(GL_DEPTH_TEST);
    glDepthMask(prevDepthMask);
    if (!prevBlend) glDisable(GL_BLEND);

    glUseProgram(0);
    m_drawTarget->Unbind();
}

// ===========================================================================
// BuildLightWireframeLines
//
// Geometry is built in the light's LOCAL space (using its schema attributes in
// local units) and every vertex is pushed through localToWorld, so position,
// orientation and scale are all handled uniformly.  USD lights emit along local
// -Z (rect/disk/distant); cylinder length runs along local X.
// ===========================================================================

void UsdSceneRenderer::BuildLightWireframeLines(const pxr::UsdPrim&    lightPrim,
                                                const pxr::GfMatrix4d& l2w,
                                                double                 scale,
                                                std::vector<float>&    outVerts)
{
    using pxr::GfVec3d;
    static constexpr int    kSeg   = 24;
    static constexpr double kTwoPi = 6.283185307179586;

    auto pushLine = [&](const GfVec3d& a, const GfVec3d& b) {
        GfVec3d wa = l2w.Transform(a);
        GfVec3d wb = l2w.Transform(b);
        outVerts.push_back((float)wa[0]); outVerts.push_back((float)wa[1]); outVerts.push_back((float)wa[2]);
        outVerts.push_back((float)wb[0]); outVerts.push_back((float)wb[1]); outVerts.push_back((float)wb[2]);
    };
    // Circle centred at c, spanned by unit axes u,v, radius r (local space).
    auto circle = [&](const GfVec3d& c, const GfVec3d& u, const GfVec3d& v, double r) {
        GfVec3d prev;
        for (int i = 0; i <= kSeg; ++i) {
            double a = kTwoPi * double(i) / double(kSeg);
            GfVec3d p = c + u * (r * std::cos(a)) + v * (r * std::sin(a));
            if (i > 0) pushLine(prev, p);
            prev = p;
        }
    };

    const GfVec3d X(1,0,0), Y(0,1,0), Z(0,0,1), O(0,0,0);
    const double  marker = scale * 0.5;  // display-scaled size for direction lines / markers

    if (lightPrim.IsA<pxr::UsdLuxRectLight>()) {
        float w = 1.0f, h = 1.0f;
        pxr::UsdLuxRectLight rl(lightPrim);
        rl.GetWidthAttr().Get(&w, m_currentTime);
        rl.GetHeightAttr().Get(&h, m_currentTime);
        double hw = w * 0.5, hh = h * 0.5;
        GfVec3d bl(-hw,-hh,0), br(hw,-hh,0), tr(hw,hh,0), tl(-hw,hh,0);
        pushLine(bl,br); pushLine(br,tr); pushLine(tr,tl); pushLine(tl,bl);
        pushLine(O, GfVec3d(0,0,-marker));  // emission direction (-Z)
    }
    else if (lightPrim.IsA<pxr::UsdLuxDiskLight>()) {
        float r = 0.5f; pxr::UsdLuxDiskLight(lightPrim).GetRadiusAttr().Get(&r, m_currentTime);
        circle(O, X, Y, r);
        pushLine(O, GfVec3d(0,0,-marker));
    }
    else if (lightPrim.IsA<pxr::UsdLuxCylinderLight>()) {
        float r = 0.5f, len = 1.0f;
        pxr::UsdLuxCylinderLight cl(lightPrim);
        cl.GetRadiusAttr().Get(&r, m_currentTime);
        cl.GetLengthAttr().Get(&len, m_currentTime);
        double hx = len * 0.5;
        circle(GfVec3d( hx,0,0), Y, Z, r);   // end caps (in local YZ)
        circle(GfVec3d(-hx,0,0), Y, Z, r);
        pushLine(GfVec3d(-hx, r,0), GfVec3d(hx, r,0));   // connecting lines
        pushLine(GfVec3d(-hx,-r,0), GfVec3d(hx,-r,0));
        pushLine(GfVec3d(-hx,0, r), GfVec3d(hx,0, r));
        pushLine(GfVec3d(-hx,0,-r), GfVec3d(hx,0,-r));
    }
    else if (lightPrim.IsA<pxr::UsdLuxDistantLight>()) {
        // Sun: small disc facing -Z plus parallel rays along -Z.
        double r = marker;
        circle(O, X, Y, r);
        const GfVec3d off[5] = { O, GfVec3d(r,0,0), GfVec3d(-r,0,0), GfVec3d(0,r,0), GfVec3d(0,-r,0) };
        for (const auto& o : off) pushLine(o, o + GfVec3d(0,0,-marker*2.0));
    }
    else if (lightPrim.IsA<pxr::UsdLuxDomeLight>()) {
        // Environment dome: 3 large display-scaled circles around the origin.
        double r = std::max(scale, 1.0);
        circle(O, X, Y, r); circle(O, X, Z, r); circle(O, Y, Z, r);
    }
    else if (lightPrim.IsA<pxr::UsdLuxSphereLight>()) {
        float r = 0.5f; pxr::UsdLuxSphereLight sl(lightPrim);
        sl.GetRadiusAttr().Get(&r, m_currentTime);
        bool pointish = false; sl.GetTreatAsPointAttr().Get(&pointish, m_currentTime);
        if (pointish || r < 1e-4f) {
            double s = marker;             // point light: axis cross + tiny disc
            pushLine(GfVec3d(-s,0,0), GfVec3d(s,0,0));
            pushLine(GfVec3d(0,-s,0), GfVec3d(0,s,0));
            pushLine(GfVec3d(0,0,-s), GfVec3d(0,0,s));
            circle(O, X, Y, s * 0.4);
        } else {
            circle(O, X, Y, r); circle(O, X, Z, r); circle(O, Y, Z, r);
        }
    }
    else {
        double s = marker;                 // unknown light: generic point marker
        pushLine(GfVec3d(-s,0,0), GfVec3d(s,0,0));
        pushLine(GfVec3d(0,-s,0), GfVec3d(0,s,0));
        pushLine(GfVec3d(0,0,-s), GfVec3d(0,0,s));
    }
}

// ===========================================================================
// StageHasAuthoredLights
// ===========================================================================

bool UsdSceneRenderer::StageHasAuthoredLights() const
{
    if (!m_stage) return false;
    for (const pxr::UsdPrim& prim : m_stage->Traverse()) {
        if (prim.HasAPI<pxr::UsdLuxLightAPI>())
            return true;
    }
    return false;
}

// ===========================================================================
// DrawLightWireframes  (mirrors DrawCameraWireframes)
// ===========================================================================

void UsdSceneRenderer::DrawLightWireframes(
    pxr::UsdStageRefPtr        stage,
    const pxr::SdfPathVector&  selectedPaths,
    const pxr::GfMatrix4d&     viewProjMatrix,
    double                     viewportCameraDist)
{
    if (!stage) return;
    if (!m_bboxProgram || !m_lightWireVAO || !m_lightWireVBO) return;
    if (!m_drawTarget) return;

    // Collect every light prim (anything carrying UsdLuxLightAPI).
    pxr::SdfPathVector lightPaths;
    for (const pxr::UsdPrim& prim : stage->Traverse()) {
        if (prim.HasAPI<pxr::UsdLuxLightAPI>())
            lightPaths.push_back(prim.GetPath());
    }
    if (lightPaths.empty()) return;

    double rawScale = viewportCameraDist * 0.12;
    double scale    = std::min(rawScale, 50.0);
    scale = std::max(scale, 0.5);

    std::unordered_set<std::string> selectedSet;
    for (const auto& p : selectedPaths) selectedSet.insert(p.GetString());

    pxr::GfVec2i sz = m_drawTarget->GetSize();
    m_drawTarget->Bind();
    glViewport(0, 0, sz[0], sz[1]);

    GLboolean prevDepthTest = glIsEnabled(GL_DEPTH_TEST);
    GLboolean prevDepthMask; glGetBooleanv(GL_DEPTH_WRITEMASK, &prevDepthMask);
    GLboolean prevBlend = glIsEnabled(GL_BLEND);

    // Draw on top of everything (like camera wireframes) so lights stay visible.
    glDisable(GL_DEPTH_TEST);
    glDepthMask(GL_FALSE);
    glEnable(GL_BLEND);
    glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA);
    if (m_aaEnabled) glEnable(GL_LINE_SMOOTH);

    glUseProgram(m_bboxProgram);
    pxr::GfMatrix4f vp = pxr::GfMatrix4f(viewProjMatrix);

    pxr::UsdGeomXformCache xformCache(m_currentTime);

    for (const auto& lpath : lightPaths) {
        pxr::UsdPrim prim = stage->GetPrimAtPath(lpath);
        if (!prim) continue;

        pxr::GfMatrix4d l2w = xformCache.GetLocalToWorldTransform(prim);

        bool isSelected = (selectedSet.count(lpath.GetString()) > 0);
        pxr::GfVec4f col = isSelected
            ? pxr::GfVec4f(1.00f, 0.75f, 0.10f, 1.00f)   // accent orange
            : pxr::GfVec4f(1.00f, 0.90f, 0.35f, 0.90f);  // warm yellow
        glUniform4f(m_bboxUniformColor, col[0], col[1], col[2], col[3]);

        std::vector<float> verts;
        BuildLightWireframeLines(prim, l2w, scale, verts);
        if (verts.empty()) continue;

        int vertCount = static_cast<int>(verts.size() / 3);
        glBindBuffer(GL_ARRAY_BUFFER, m_lightWireVBO);
        glBufferData(GL_ARRAY_BUFFER,
                     static_cast<GLsizeiptr>(verts.size() * sizeof(float)),
                     verts.data(), GL_DYNAMIC_DRAW);
        glUniformMatrix4fv(m_bboxUniformMVP, 1, GL_TRUE, vp.GetArray());
        glBindVertexArray(m_lightWireVAO);
        glDrawArrays(GL_LINES, 0, vertCount);
        glBindVertexArray(0);
    }

    if (m_aaEnabled) glDisable(GL_LINE_SMOOTH);
    if (prevDepthTest)  glEnable(GL_DEPTH_TEST); else glDisable(GL_DEPTH_TEST);
    glDepthMask(prevDepthMask);
    if (!prevBlend) glDisable(GL_BLEND);

    glUseProgram(0);
    m_drawTarget->Unbind();
}

// ===========================================================================
// PickCameraAtPoint
// ===========================================================================

bool UsdSceneRenderer::PickCameraAtPoint(
    pxr::UsdStageRefPtr    stage,
    float mouseX, float mouseY,
    const pxr::GfMatrix4d& viewProjMatrix,
    float imagePosX, float imagePosY,
    int viewW, int viewH,
    double                 viewportCameraDist,
    pxr::SdfPath*          outCameraPath)
{
    if (!stage || !outCameraPath) return false;

    // Rebuild cache fresh (same logic as DrawCameraWireframes — always traverse
    // so newly created cameras are immediately pickable).
    m_cachedCameraPaths.clear();
    for (const pxr::UsdPrim& prim : stage->Traverse()) {
        if (prim.IsA<pxr::UsdGeomCamera>())
            m_cachedCameraPaths.push_back(prim.GetPath());
    }
    if (m_cachedCameraPaths.empty()) return false;

    static constexpr float kPickRadius = 10.0f;

    double rawScale = viewportCameraDist * 0.12;
    double scale    = std::min(rawScale, 50.0);
    scale = std::max(scale, 0.05);

    float  bestDist = kPickRadius;
    pxr::SdfPath bestPath;

    for (const auto& camPath : m_cachedCameraPaths) {
        pxr::UsdPrim prim = stage->GetPrimAtPath(camPath);
        if (!prim) continue;

        pxr::UsdGeomCamera usdCam(prim);
        pxr::GfCamera gfCam = usdCam.GetCamera(m_currentTime);

        std::vector<float> verts;
        BuildCameraWireframeLines(gfCam, scale, verts);

        // verts = interleaved XYZ pairs (2 verts per segment = 6 floats per segment)
        for (size_t i = 0; i + 5 < verts.size(); i += 6) {
            pxr::GfVec3d wa(verts[i],   verts[i+1], verts[i+2]);
            pxr::GfVec3d wb(verts[i+3], verts[i+4], verts[i+5]);

            float ax, ay, bx, by;
            bool okA = WorldToScreen(wa, viewProjMatrix, viewW, viewH, imagePosX, imagePosY, ax, ay);
            bool okB = WorldToScreen(wb, viewProjMatrix, viewW, viewH, imagePosX, imagePosY, bx, by);
            if (!okA && !okB) continue;
            if (!okA) { ax = bx; ay = by; }
            if (!okB) { bx = ax; by = ay; }

            float d = PointToSegmentDist(mouseX, mouseY, ax, ay, bx, by);
            if (d < bestDist) {
                bestDist = d;
                bestPath = camPath;
            }
        }
    }

    if (bestPath.IsEmpty()) return false;
    *outCameraPath = bestPath;
    return true;
}

} // namespace UsdLayerManager