#include "MaterialPreviewRenderer.h"
#include "../utils/Logger.h"
#include "../utils/PathUtils.h"
#include <pxr/base/tf/stringUtils.h>
#include <pxr/usd/usdGeom/sphere.h>
#include <pxr/usd/usdGeom/cube.h>
#include <pxr/usd/usdGeom/cylinder.h>
#include <pxr/usd/usdGeom/mesh.h>
#include <pxr/usd/usdGeom/primvarsAPI.h>
#include <pxr/usd/usdGeom/scope.h>
#include <pxr/usd/usdGeom/metrics.h>
#include <pxr/usd/usdGeom/tokens.h>
#include <pxr/usd/usdVol/volume.h>
#include <pxr/usd/usdVol/openVDBAsset.h>
#include <pxr/base/vt/array.h>
#include <pxr/base/gf/vec2f.h>
#include <pxr/base/gf/vec3f.h>
#include <pxr/usd/usdShade/materialBindingAPI.h>
#include <pxr/usd/usdShade/material.h>
#include <pxr/usd/usdShade/shader.h>
#include <pxr/usd/usdShade/connectableAPI.h>
#include <pxr/usd/usdLux/distantLight.h>
#include <pxr/usd/usdLux/domeLight.h>
#include <pxr/usd/usdLux/lightAPI.h>
#include <pxr/usd/usdLux/tokens.h>
#include <pxr/usd/usdGeom/xformCommonAPI.h>
#include <pxr/usd/usd/references.h>
#include <pxr/usd/usd/editContext.h>
#include <pxr/usd/sdf/assetPath.h>
#include <pxr/usd/sdf/primSpec.h>
#include <imgui.h>
#include <filesystem>

namespace UsdLayerManager {

namespace {
const pxr::SdfPath kShapesRootPath("/Preview/Shapes");
const pxr::SdfPath kDistantLightPath("/Preview/Light");
const pxr::SdfPath kDomeLightPath("/Preview/DomeLight");

/// Selectable preview geometry; prim name is a child of kShapesRootPath.
/// Shapes backed by an asset file are only created when the file exists.
/// frameFit: camera framing margin — the camera frames the shape's bounding
/// box, and a cube fills its box completely while a sphere only inscribes it,
/// so boxy shapes need a larger margin to read at a similar visual size.
struct ShapePreset {
    const char* label;
    const char* primName;
    double      frameFit;
};
const ShapePreset kShapePresets[] = {
    {"Sphere",       "Sphere",   1.6},
    {"Cube",         "Cube",     2.1},
    {"Cylinder",     "Cylinder", 1.8},
    {"Teapot",       "Teapot",   1.6},
    {"Hair",         "Hair",     1.6},
    {"Cloud Volume", "Cloud",    1.6},
};
constexpr int kShapePresetCount = static_cast<int>(sizeof(kShapePresets) / sizeof(kShapePresets[0]));

/// Scratch UsdPreviewSurface authored inside the (referenced) material scope
/// when previewing a non-terminal node output — child of the material so the
/// connection stays encapsulated.
const pxr::TfToken kNodePreviewShaderName("__uslm_node_preview");

struct LightPreset {
    const char* label;
    const char* file; ///< under resources/hdri/ (CC0, Poly Haven; fetched by CMake)
};
const LightPreset kLightPresets[] = {
    {"External", "kloofendal_48d_partly_cloudy_puresky_1k.exr"},
    {"Room",     "lebombo_1k.exr"},
    {"Interior", "artist_workshop_1k.exr"},
    {"Sunset",   "venice_sunset_1k.exr"},
};
constexpr int kLightPresetCount = static_cast<int>(sizeof(kLightPresets) / sizeof(kLightPresets[0]));

/// Authors a polygonal UsdGeomMesh from prebuilt arrays: vertex-interpolated
/// normals + a `st` texCoord2f primvar, no subdivision, single-sided. Shared by
/// every preview shape so they all light and texture identically.
void DefineUvMesh(const pxr::UsdStageRefPtr& stage, const pxr::SdfPath& path,
                  const pxr::VtVec3fArray& points, const pxr::VtVec3fArray& normals,
                  const pxr::VtVec2fArray& st, const pxr::VtIntArray& fvCounts,
                  const pxr::VtIntArray& fvIndices,
                  const pxr::GfVec3f& extentMin, const pxr::GfVec3f& extentMax) {
    pxr::UsdGeomMesh mesh = pxr::UsdGeomMesh::Define(stage, path);
    mesh.CreatePointsAttr(pxr::VtValue(points));
    mesh.CreateNormalsAttr(pxr::VtValue(normals));
    mesh.SetNormalsInterpolation(pxr::UsdGeomTokens->vertex);
    mesh.CreateFaceVertexCountsAttr(pxr::VtValue(fvCounts));
    mesh.CreateFaceVertexIndicesAttr(pxr::VtValue(fvIndices));
    // Polygonal (not subdivided) so the authored smooth normals and UVs are used
    // verbatim. Single-sided (the default): a closed shape never shows a
    // backface, and double-sided shading runs a view-dependent normal flip that
    // corrupts the mirror-reflection vector near the silhouette.
    mesh.CreateSubdivisionSchemeAttr(pxr::VtValue(pxr::UsdGeomTokens->none));
    mesh.CreateExtentAttr(pxr::VtValue(pxr::VtVec3fArray{extentMin, extentMax}));
    pxr::UsdGeomPrimvarsAPI(mesh).CreatePrimvar(
        pxr::TfToken("st"), pxr::SdfValueTypeNames->TexCoord2fArray, pxr::UsdGeomTokens->vertex)
        .Set(st);
}

/// Authors a UV-mapped mesh sphere (Y-up) at `path` with vertex normals and a
/// `st` texCoord2f primvar. The implicit UsdGeomSphere carries no texture
/// coordinates (GeomUtilSphereMeshGenerator emits only points/normals/topology),
/// so a UsdUVTexture fed by a UsdPrimvarReader_float2("st") would only ever read
/// the reader's fallback and paint a flat colour. This mesh gives the shader
/// ball real UVs so textures map across every render delegate.
void CreateUvSphere(const pxr::UsdStageRefPtr& stage, const pxr::SdfPath& path, double radius) {
    constexpr int    rings   = 48;   // latitude bands  (theta 0..pi)
    constexpr int    sectors = 96;   // longitude bands (phi 0..2pi)
    constexpr double kPi     = 3.14159265358979323846;
    const float      r       = static_cast<float>(radius);

    pxr::VtVec3fArray points, normals;
    pxr::VtVec2fArray st;
    points.reserve((rings + 1) * (sectors + 1));
    normals.reserve((rings + 1) * (sectors + 1));
    st.reserve((rings + 1) * (sectors + 1));
    // A duplicated seam column (sectors+1 wide) gives distinct u=0 and u=1
    // vertices so the texture wraps without a smear across the seam.
    for (int i = 0; i <= rings; ++i) {
        const float v    = static_cast<float>(i) / rings;
        const float th   = v * static_cast<float>(kPi);
        const float sinT = std::sin(th), cosT = std::cos(th);
        for (int j = 0; j <= sectors; ++j) {
            const float u    = static_cast<float>(j) / sectors;
            const float phi  = u * 2.0f * static_cast<float>(kPi);
            const pxr::GfVec3f n(sinT * std::cos(phi), cosT, sinT * std::sin(phi));
            points.push_back(n * r);
            normals.push_back(n);
            st.push_back(pxr::GfVec2f(u, 1.0f - v)); // flip V so image top → sphere top
        }
    }

    pxr::VtIntArray fvCounts, fvIndices;
    const int stride = sectors + 1;
    for (int i = 0; i < rings; ++i)
        for (int j = 0; j < sectors; ++j) {
            const int a = i * stride + j;
            fvCounts.push_back(4);
            // Wind CCW as seen from outside so the (single-sided) front face and
            // the outward vertex normals agree; a reversed winding would either
            // cull the visible surface or flip its shading normal.
            fvIndices.push_back(a);
            fvIndices.push_back(a + 1);
            fvIndices.push_back(a + stride + 1);
            fvIndices.push_back(a + stride);
        }

    DefineUvMesh(stage, path, points, normals, st, fvCounts, fvIndices,
                 pxr::GfVec3f(-r, -r, -r), pxr::GfVec3f(r, r, r));
}

/// UV-mapped box: 24 vertices (4 per face) so every face gets its own outward
/// normal and a full 0..1 UV square. halfSize is the half-edge length.
void CreateUvCube(const pxr::UsdStageRefPtr& stage, const pxr::SdfPath& path, double halfSize) {
    const float h = static_cast<float>(halfSize);
    pxr::VtVec3fArray points, normals;
    pxr::VtVec2fArray st;
    pxr::VtIntArray fvCounts, fvIndices;

    // p0..p3 are wound CCW as seen from outside, so the front face agrees with n.
    auto addFace = [&](const pxr::GfVec3f& p0, const pxr::GfVec3f& p1,
                       const pxr::GfVec3f& p2, const pxr::GfVec3f& p3,
                       const pxr::GfVec3f& n) {
        const int base = static_cast<int>(points.size());
        points.push_back(p0); points.push_back(p1); points.push_back(p2); points.push_back(p3);
        for (int k = 0; k < 4; ++k) normals.push_back(n);
        st.push_back(pxr::GfVec2f(0, 0)); st.push_back(pxr::GfVec2f(1, 0));
        st.push_back(pxr::GfVec2f(1, 1)); st.push_back(pxr::GfVec2f(0, 1));
        fvCounts.push_back(4);
        fvIndices.push_back(base + 0); fvIndices.push_back(base + 1);
        fvIndices.push_back(base + 2); fvIndices.push_back(base + 3);
    };
    addFace({ h,-h, h}, { h,-h,-h}, { h, h,-h}, { h, h, h}, { 1, 0, 0}); // +X
    addFace({-h,-h,-h}, {-h,-h, h}, {-h, h, h}, {-h, h,-h}, {-1, 0, 0}); // -X
    addFace({-h, h,-h}, {-h, h, h}, { h, h, h}, { h, h,-h}, { 0, 1, 0}); // +Y
    addFace({-h,-h,-h}, { h,-h,-h}, { h,-h, h}, {-h,-h, h}, { 0,-1, 0}); // -Y
    addFace({-h,-h, h}, { h,-h, h}, { h, h, h}, {-h, h, h}, { 0, 0, 1}); // +Z
    addFace({ h,-h,-h}, {-h,-h,-h}, {-h, h,-h}, { h, h,-h}, { 0, 0,-1}); // -Z

    DefineUvMesh(stage, path, points, normals, st, fvCounts, fvIndices,
                 pxr::GfVec3f(-h, -h, -h), pxr::GfVec3f(h, h, h));
}

/// UV-mapped cylinder (Y axis): radial side wall + two end caps. Side normals
/// are radial, caps ±Y; UVs wrap the wall (u around, v along height) and map
/// each cap as a disk.
void CreateUvCylinder(const pxr::UsdStageRefPtr& stage, const pxr::SdfPath& path,
                      double radius, double halfHeight) {
    constexpr int    sectors = 96;
    constexpr double kPi     = 3.14159265358979323846;
    const float      R       = static_cast<float>(radius);
    const float      hy      = static_cast<float>(halfHeight);

    pxr::VtVec3fArray points, normals;
    pxr::VtVec2fArray st;
    pxr::VtIntArray fvCounts, fvIndices;

    // Side wall: sectors+1 columns (duplicated seam for a clean u wrap), each a
    // bottom+top pair interleaved so column j owns indices 2j (bottom), 2j+1 (top).
    for (int j = 0; j <= sectors; ++j) {
        const float u = static_cast<float>(j) / sectors;
        const float phi = u * 2.0f * static_cast<float>(kPi);
        const float c = std::cos(phi), s = std::sin(phi);
        points.push_back(pxr::GfVec3f(R * c, -hy, R * s));
        normals.push_back(pxr::GfVec3f(c, 0, s));
        st.push_back(pxr::GfVec2f(u, 0));
        points.push_back(pxr::GfVec3f(R * c,  hy, R * s));
        normals.push_back(pxr::GfVec3f(c, 0, s));
        st.push_back(pxr::GfVec2f(u, 1));
    }
    for (int j = 0; j < sectors; ++j) {
        const int b0 = 2 * j, t0 = 2 * j + 1, b1 = 2 * (j + 1), t1 = 2 * (j + 1) + 1;
        fvCounts.push_back(4);           // b_j, t_j, t_{j+1}, b_{j+1} → outward radial
        fvIndices.push_back(b0); fvIndices.push_back(t0);
        fvIndices.push_back(t1); fvIndices.push_back(b1);
    }

    // End caps as triangle fans around a center vertex; ring vertices carry the
    // cap normal (not the radial one) and disk UVs.
    auto addCap = [&](float y, float ny) {
        const int center = static_cast<int>(points.size());
        points.push_back(pxr::GfVec3f(0, y, 0));
        normals.push_back(pxr::GfVec3f(0, ny, 0));
        st.push_back(pxr::GfVec2f(0.5f, 0.5f));
        const int ringBase = static_cast<int>(points.size());
        for (int j = 0; j < sectors; ++j) {
            const float phi = static_cast<float>(j) / sectors * 2.0f * static_cast<float>(kPi);
            const float c = std::cos(phi), s = std::sin(phi);
            points.push_back(pxr::GfVec3f(R * c, y, R * s));
            normals.push_back(pxr::GfVec3f(0, ny, 0));
            st.push_back(pxr::GfVec2f(0.5f + 0.5f * c, 0.5f + 0.5f * s));
        }
        for (int j = 0; j < sectors; ++j) {
            const int a = ringBase + j, b = ringBase + (j + 1) % sectors;
            fvCounts.push_back(3);
            fvIndices.push_back(center);
            // Wind so the tri normal matches the cap direction (+Y vs -Y).
            if (ny > 0) { fvIndices.push_back(b); fvIndices.push_back(a); }
            else        { fvIndices.push_back(a); fvIndices.push_back(b); }
        }
    };
    addCap( hy,  1.0f);
    addCap(-hy, -1.0f);

    DefineUvMesh(stage, path, points, normals, st, fvCounts, fvIndices,
                 pxr::GfVec3f(-R, -hy, -R), pxr::GfVec3f(R, hy, R));
}
} // namespace

void MaterialPreviewRenderer::EnsureInitialized() {
    if (m_initialized) return;

    m_previewStage = pxr::UsdStage::CreateInMemory();
    // Everything authored here (sphere, lights, reference) must go to the
    // root layer: the session layer is reserved for node-preview overrides
    // and is wiped wholesale on every SetMaterial.
    m_previewStage->SetEditTarget(pxr::UsdEditTarget(m_previewStage->GetRootLayer()));
    // Pin the up axis rather than inheriting the site fallback: the light
    // rig, dome-pole compensation, and camera framing all assume Y-up.
    pxr::UsdGeomSetStageUpAxis(m_previewStage, pxr::UsdGeomTokens->y);

    // All selectable preview shapes live under one scope; the material is
    // bound on the scope so it inherits to whichever shape is active, and
    // ApplyPreviewShape keeps exactly one child active.
    pxr::UsdGeomScope::Define(m_previewStage, kShapesRootPath);

    // UV-mapped meshes (not implicit UsdGeom prims) so textured materials —
    // e.g. a UsdUVTexture wired into diffuseColor via a UsdPrimvarReader "st" —
    // map onto the preview geometry instead of sampling a single texel. The
    // implicit primitives generate no `st` primvar, so they can't show textures.
    CreateUvSphere(m_previewStage, kShapesRootPath.AppendChild(pxr::TfToken("Sphere")), 1.0);
    CreateUvCube(m_previewStage, kShapesRootPath.AppendChild(pxr::TfToken("Cube")), 0.6);
    CreateUvCylinder(m_previewStage, kShapesRootPath.AppendChild(pxr::TfToken("Cylinder")), 0.7, 0.8);

    // Asset-backed shapes, created only when their file is bundled.
    for (const char* name : {"Teapot", "Hair"}) {
        const std::string file = ResourcePath(
            std::string("resources/preview/") + pxr::TfStringToLower(name) + ".usdc");
        if (!std::filesystem::exists(file)) continue;
        pxr::UsdPrim prim = m_previewStage->DefinePrim(
            kShapesRootPath.AppendChild(pxr::TfToken(name)));
        prim.GetReferences().AddReference(file);
    }

    const std::string vdbFile = ResourcePath("resources/vdb/bunny_cloud.vdb");
    if (std::filesystem::exists(vdbFile)) {
        pxr::UsdVolVolume cloud = pxr::UsdVolVolume::Define(
            m_previewStage, kShapesRootPath.AppendChild(pxr::TfToken("Cloud")));
        pxr::UsdVolOpenVDBAsset density = pxr::UsdVolOpenVDBAsset::Define(
            m_previewStage, cloud.GetPath().AppendChild(pxr::TfToken("density")));
        density.CreateFilePathAttr().Set(pxr::SdfAssetPath(vdbFile));
        density.CreateFieldNameAttr().Set(pxr::TfToken("density"));
        cloud.CreateFieldRelationship(pxr::TfToken("density"), density.GetPath());
    }

    // Storm renders fine with zero authored lights (it injects a GL headlight
    // as a fallback), but proper Hydra delegates like Embree/Arnold shade
    // through scene lights only — without one they render solid black.
    // This distant light is the fallback when the HDR dome preset's texture
    // isn't bundled; ApplyLightPreset toggles between it and the dome.
    pxr::UsdLuxDistantLight light = pxr::UsdLuxDistantLight::Define(m_previewStage, kDistantLightPath);
    // pxr::UsdLuxLightAPI(light.GetPrim()).CreateIntensityAttr(pxr::VtValue(15000.0f));
    pxr::UsdLuxLightAPI(light.GetPrim()).CreateIntensityAttr(pxr::VtValue(1500.0f));
    // Default distant-light orientation shines along local -Z, which lights
    // the far side of the sphere from the camera's default view. Angle it
    // like a classic front-upper-left key light instead. Adding this scene
    // light suppresses UsdSceneRenderer's GL headlight fallback (it only
    // kicks in when the stage has no authored lights), so this needs to
    // carry the whole preview on its own.
    pxr::UsdGeomXformCommonAPI(light.GetPrim()).SetRotate(pxr::GfVec3f(-35.0f, -35.0f, 0.0f));

    m_camera.SetStage(m_previewStage);
    m_renderer.SetStage(m_previewStage);
    m_renderer.SetShowGrid(false);
    // Raised refinement so the implicit sphere tessellates smoothly instead
    // of showing faceted silhouettes at swatch scale.
    m_renderer.SetComplexity(1.3f);

    ApplyLightPreset(m_lightPreset);
    ApplyPreviewShape(m_previewShape); // activates one shape + frames the camera

    // Maya-style default 3/4 view: orbit the framed camera ~45° around and
    // ~25° down instead of the dead-on front view (subsequent shape switches
    // keep whatever orientation the user has orbited to).
    m_camera.Tumble(-45.0, 25.0);
    m_renderer.SetCameraStateFromGfCamera(
        m_camera.ComputeGfCamera(m_previewBBox, /*autoClip=*/true));

    m_initialized = true;
    m_dirty = true;
}

void MaterialPreviewRenderer::SetMaterial(const pxr::UsdStageRefPtr& sourceStage,
                                          const pxr::SdfPath& materialPath,
                                          size_t graphRevision,
                                          const pxr::SdfPath& previewNodePath,
                                          const std::string& previewNodeOutput,
                                          bool previewOutputIsTerminal) {
    EnsureInitialized();

    if (!sourceStage || materialPath.IsEmpty()) return;
    if (materialPath == m_materialPath && graphRevision == m_lastGraphRevision &&
        previewNodePath == m_previewNodePath && previewNodeOutput == m_previewNodeOutput)
        return; // nothing the shader ball cares about has changed

    if (!sourceStage->GetPrimAtPath(materialPath)) {
        LOG_WARNING("MaterialPreviewRenderer: no prim at " + materialPath.GetString());
        return;
    }

    // Compose the material into the preview stage via a reference to the
    // source stage's root layer, rather than SdfCopySpec from one local
    // layer: a material pulled in by a reference (e.g. a .mtlx asset) has no
    // spec in any local layer — it only exists composed — and even local
    // materials can be spread across sublayers. Referencing at the identical
    // path keeps internal connection paths valid without remapping, and
    // relative asset paths keep resolving against their original layers.
    // Authored only when the material changes: the shared in-memory layers
    // recompose live, so revision bumps need no re-referencing.
    if (materialPath != m_materialPath) {
        pxr::UsdPrim mirrorPrim = m_previewStage->OverridePrim(materialPath);
        if (!mirrorPrim) {
            LOG_ERROR("MaterialPreviewRenderer: failed to create override at " + materialPath.GetString());
            return;
        }
        mirrorPrim.GetReferences().ClearReferences();
        mirrorPrim.GetReferences().AddReference(
            sourceStage->GetRootLayer()->GetIdentifier(), materialPath);
    }

    pxr::UsdPrim shapesRoot = m_previewStage->GetPrimAtPath(kShapesRootPath);
    pxr::UsdShadeMaterial material(m_previewStage->GetPrimAtPath(materialPath));
    if (!material) {
        LOG_WARNING("MaterialPreviewRenderer: no material composed at " + materialPath.GetString());
        return;
    }
    if (shapesRoot) // bound on the scope: inherits to whichever shape is active
        pxr::UsdShadeMaterialBindingAPI::Apply(shapesRoot).Bind(material);

    // ── Selected-node preview (Hypershade-style) ─────────────────────────
    // All node-preview overrides (the scratch wrapper shader and the local
    // opinions on the material's surface outputs) live exclusively in the
    // preview stage's session layer, which is stronger than the root layer's
    // reference. Wiping them is then just clearing that layer — no per-spec
    // surgery (RemoveNameChild/RemoveProperty) that could interact badly with
    // Sdf's inert-spec cleanup, and no stale overrides on other materials.
    pxr::SdfLayerHandle sessionLayer = m_previewStage->GetSessionLayer();
    sessionLayer->Clear();

    // Route the material's surface through the selected node. Nodes outside
    // the material's subtree (possible in hand-authored cross-scope networks)
    // aren't composed into the scratch stage, so those keep the whole-material
    // preview.
    if (!previewNodePath.IsEmpty() && !previewNodeOutput.empty() &&
        previewNodePath.HasPrefix(materialPath)) {
        pxr::UsdShadeShader nodeShader(m_previewStage->GetPrimAtPath(previewNodePath));
        if (nodeShader) {
            pxr::UsdEditContext sessionEdit(m_previewStage,
                                            pxr::UsdEditTarget(sessionLayer));
            pxr::UsdShadeConnectableAPI sourceApi = nodeShader.ConnectableAPI();
            pxr::TfToken sourceName(previewNodeOutput);
            if (!previewOutputIsTerminal) {
                // Pattern/texture output: show it as the diffuse color of a
                // scratch UsdPreviewSurface, like Maya's swatch for textures.
                pxr::UsdShadeShader wrap = pxr::UsdShadeShader::Define(
                    m_previewStage, materialPath.AppendChild(kNodePreviewShaderName));
                wrap.CreateIdAttr(pxr::VtValue(pxr::TfToken("UsdPreviewSurface")));
                wrap.CreateInput(pxr::TfToken("roughness"), pxr::SdfValueTypeNames->Float)
                    .Set(0.4f);
                wrap.CreateInput(pxr::TfToken("diffuseColor"), pxr::SdfValueTypeNames->Color3f)
                    .ConnectToSource(sourceApi, sourceName);
                sourceApi  = wrap.ConnectableAPI();
                sourceName = pxr::TfToken("surface");
            }
            // Override both the universal and the mtlx render-context outputs
            // so the preview target wins regardless of which context the
            // material authored (a .mtlx import only authors outputs:mtlx:*,
            // which Storm would otherwise prefer).
            material.CreateSurfaceOutput().ConnectToSource(sourceApi, sourceName);
            material.CreateSurfaceOutput(pxr::TfToken("mtlx")).ConnectToSource(sourceApi, sourceName);
        }
    }

    m_materialPath      = materialPath;
    m_lastGraphRevision = graphRevision;
    m_previewNodePath   = previewNodePath;
    m_previewNodeOutput = previewNodeOutput;
    m_dirty             = true;
}

void MaterialPreviewRenderer::ApplyLightPreset(int index) {
    if (index < 0 || index >= kLightPresetCount) index = 0;
    m_lightPreset = index;

    const std::string file = ResourcePath(std::string("resources/hdri/") + kLightPresets[index].file);
    const bool haveHdr = std::filesystem::exists(file);

    pxr::UsdPrim distant = m_previewStage->GetPrimAtPath(kDistantLightPath);
    if (haveHdr) {
        pxr::UsdLuxDomeLight dome = pxr::UsdLuxDomeLight::Define(m_previewStage, kDomeLightPath);
        dome.CreateTextureFileAttr().Set(pxr::SdfAssetPath(file));
        // Declare the projection explicitly: left at "automatic", hdArnold
        // falls through to Arnold's skydome default ("angular" fisheye),
        // which squeezes an equirect panorama into a blob at the center.
        // Storm always samples latlong regardless, so this only affects
        // spec-following delegates.
        dome.CreateTextureFormatAttr().Set(pxr::UsdLuxTokens->latlong);
        dome.GetPrim().SetActive(true);
        if (distant) distant.SetActive(false); // dome carries the lighting alone
        UpdateDomeOrientation();
    } else {
        LOG_WARNING("MaterialPreviewRenderer: HDRI not found (" + file +
                    ") — falling back to distant light. Re-run CMake configure to download presets.");
        if (pxr::UsdPrim dome = m_previewStage->GetPrimAtPath(kDomeLightPath))
            dome.SetActive(false);
        if (distant) distant.SetActive(true);
    }
    m_dirty = true;
}

void MaterialPreviewRenderer::ApplyPreviewShape(int index) {
    if (index < 0 || index >= kShapePresetCount) index = 0;
    // Fall back to the sphere if the requested shape's asset wasn't bundled.
    if (!m_previewStage->GetPrimAtPath(
            kShapesRootPath.AppendChild(pxr::TfToken(kShapePresets[index].primName))))
        index = 0;
    m_previewShape = index;

    for (int i = 0; i < kShapePresetCount; ++i) {
        pxr::UsdPrim prim = m_previewStage->GetPrimAtPath(
            kShapesRootPath.AppendChild(pxr::TfToken(kShapePresets[i].primName)));
        if (prim)
            prim.SetActive(i == index);
    }

    // Shapes differ wildly in size (the cloud VDB spans hundreds of units) —
    // reframe on the active shape's bounds; inactive prims don't contribute.
    m_previewBBox = pxr::GfBBox3d(m_renderer.ComputeStageBounds());
    m_camera.FrameSelection(m_previewBBox, kShapePresets[index].frameFit);
    m_renderer.SetCameraStateFromGfCamera(
        m_camera.ComputeGfCamera(m_previewBBox, /*autoClip=*/true));
    m_dirty = true;
}

void MaterialPreviewRenderer::RenderShapeDropdown(float width) {
    if (!m_initialized) return;
    ImGui::SetNextItemWidth(width);
    if (ImGui::BeginCombo("##PreviewShape", kShapePresets[m_previewShape].label)) {
        for (int i = 0; i < kShapePresetCount; ++i) {
            if (!m_previewStage->GetPrimAtPath(
                    kShapesRootPath.AppendChild(pxr::TfToken(kShapePresets[i].primName))))
                continue; // asset not bundled
            if (ImGui::Selectable(kShapePresets[i].label, i == m_previewShape) && i != m_previewShape)
                ApplyPreviewShape(i);
        }
        ImGui::EndCombo();
    }
    if (ImGui::IsItemHovered())
        ImGui::SetTooltip("Shader-ball preview geometry");
}

void MaterialPreviewRenderer::UpdateDomeOrientation() {
    if (!m_previewStage) return;
    pxr::UsdPrim dome = m_previewStage->GetPrimAtPath(kDomeLightPath);
    if (!dome) return;

    // Storm samples the dome texture with the pole along the light's local
    // +Y (domeLight.glslfx: v = acos(dir.y)/pi), while spec-following
    // delegates map it with the pole along local +Z (hdCycles feeds the
    // dome-local direction to Cycles' Z-up EnvironmentTextureNode; hdArnold's
    // skydome follows the same UsdLux convention). Rotating the dome -90°
    // about X for those points their +Z pole at world +Y, which reproduces
    // Storm's orientation exactly — including the horizontal (longitude)
    // alignment, since both parameterize u from the same axes.
    const pxr::TfToken rendererId = m_renderer.GetCurrentRendererId();
    const bool yUpPole = rendererId.IsEmpty() ||
                         rendererId == pxr::TfToken("HdStormRendererPlugin");
    if (rendererId == pxr::TfToken("HdArnoldRendererPlugin")) {
        // With texture:format explicitly latlong, hdArnold aligns its skydome
        // with Storm on its own — earlier rotation offsets here were chasing
        // what turned out to be the angular-projection bug (format left
        // "automatic" fell through to Arnold's fisheye default). If a real
        // residual yaw shows up against Storm, put it in the Y component.
        pxr::UsdGeomXformCommonAPI(dome).SetRotate(pxr::GfVec3f(0.0f, 0.0f, 0.0f));
    } else {
        if(rendererId == pxr::TfToken("HdCyclesPlugin")) {
            // hdCycles' latlong mapping is flipped 180° from Storm's, so add
            // a yaw to match Storm's orientation. The pole is still along +Z.
            pxr::UsdGeomXformCommonAPI(dome).SetRotate(pxr::GfVec3f(yUpPole ? 0.0f : -90.0f, 270.0f, 0.0f));
        } else {
            pxr::UsdGeomXformCommonAPI(dome).SetRotate(pxr::GfVec3f(yUpPole ? 0.0f : -90.0f, 0.0f, 0.0f));
        }
        
    }

    m_dirty = true;
}

void MaterialPreviewRenderer::RenderLightingDropdown(float width) {
    if (!m_initialized) return;
    ImGui::SetNextItemWidth(width);
    if (ImGui::BeginCombo("##PreviewLighting", kLightPresets[m_lightPreset].label)) {
        for (int i = 0; i < kLightPresetCount; ++i) {
            if (ImGui::Selectable(kLightPresets[i].label, i == m_lightPreset) && i != m_lightPreset)
                ApplyLightPreset(i);
        }
        ImGui::EndCombo();
    }
    if (ImGui::IsItemHovered())
        ImGui::SetTooltip("Shader-ball HDR environment");
}

uint32_t MaterialPreviewRenderer::Render(int width, int height) {
    EnsureInitialized();
    if (width <= 0 || height <= 0) return 0;

    // Re-render while a progressive delegate (Arnold/Cycles/Embree) is still
    // accumulating — rendering only once would freeze its first noisy sample
    // pass into the cached texture. Storm converges in one pass, so the
    // dirty-flag caching still applies there.
    if (m_dirty || width != m_lastWidth || height != m_lastHeight || !m_renderer.IsConverged()) {
        // 2x supersample: the texture is displayed at (width, height) by
        // ImGui, so rendering larger gives the swatch cheap anti-aliasing.
        m_renderer.Render(width * 2, height * 2);
        m_dirty      = false;
        m_lastWidth  = width;
        m_lastHeight = height;
    }
    return m_renderer.GetColorTextureID();
}

void MaterialPreviewRenderer::OrbitDrag(float deltaX, float deltaY) {
    if (!m_initialized) return;
    m_camera.Tumble(deltaX * 0.5, deltaY * 0.5); // matches the viewport's tumble sensitivity
    m_renderer.SetCameraStateFromGfCamera(
        m_camera.ComputeGfCamera(m_previewBBox, /*autoClip=*/true));
    m_dirty = true;
}

void MaterialPreviewRenderer::SetColorCorrection(int ccMode, const std::string& ocioDisplay,
                                                 const std::string& ocioView,
                                                 const std::string& ocioColorSpace,
                                                 const std::string& ocioLook) {
    m_renderer.SetColorCorrectionMode(static_cast<ColorCorrectionMode>(ccMode));
    m_renderer.SetOcioDisplay(ocioDisplay);
    m_renderer.SetOcioView(ocioView);
    m_renderer.SetOcioColorSpace(ocioColorSpace);
    m_renderer.SetOcioLook(ocioLook);
    m_dirty = true;
}

void MaterialPreviewRenderer::SetRendererPlugin(const pxr::TfToken& pluginId) {
    if (pluginId == m_renderer.GetCurrentRendererId()) return;
    m_renderer.SetRendererPlugin(pluginId);
    UpdateDomeOrientation(); // pole convention differs per delegate
    m_dirty = true;
}

pxr::TfToken MaterialPreviewRenderer::GetRendererPlugin() const {
    return m_renderer.GetCurrentRendererId();
}

void MaterialPreviewRenderer::RenderRendererDropdown(float width) {
    pxr::TfToken currentId = m_renderer.GetCurrentRendererId();
    std::string displayName = currentId.IsEmpty()
        ? "Renderer"
        : UsdSceneRenderer::GetRendererDisplayName(currentId);

    ImGui::Button(displayName.c_str(), ImVec2(width, 0.0f));
    if (ImGui::IsItemHovered())
        ImGui::SetTooltip("Shader-ball render delegate");

    if (ImGui::BeginPopupContextItem("MaterialPreviewRendererPopup", ImGuiPopupFlags_MouseButtonLeft)) {
        for (const auto& pluginId : UsdSceneRenderer::GetRendererPlugins()) {
            std::string name = UsdSceneRenderer::GetRendererDisplayName(pluginId);
            if (name.empty()) name = pluginId.GetString();
            bool selected = (pluginId == currentId);
            if (ImGui::MenuItem(name.c_str(), nullptr, selected) && !selected)
                SetRendererPlugin(pluginId);
        }
        ImGui::EndPopup();
    }
}

} // namespace UsdLayerManager
