#pragma once

#include "MaterialPreviewRenderer.h"
#include <pxr/usd/usd/stage.h>
#include <pxr/usd/sdf/path.h>
#include <pxr/base/tf/token.h>
#include <imgui.h>
#include <atomic>
#include <condition_variable>
#include <cstdint>
#include <deque>
#include <mutex>
#include <string>
#include <thread>
#include <unordered_map>
#include <vector>

namespace UsdLayerManager {

/// Owns the in-canvas node thumbnails for the material editor.
///
/// Two kinds, both handed back as an ImTextureID (0 while not ready):
///   - Texture nodes: the source image, decoded on a background worker thread
///     (USD HioImage / OpenImageIO) and uploaded to GL on the main thread.
///   - Material nodes: a shader-ball render through a dedicated
///     MaterialPreviewRenderer. Hydra is GL-/main-thread-bound, so these are
///     rendered one-per-frame (amortized) and cached by graph revision.
///
/// Lifecycle is driven from the panel: Get*Thumbnail() per node during the
/// canvas draw, then PumpMainThread() once per frame (GL context current).
class NodeThumbnailCache {
public:
    NodeThumbnailCache();
    ~NodeThumbnailCache();

    NodeThumbnailCache(const NodeThumbnailCache&) = delete;
    NodeThumbnailCache& operator=(const NodeThumbnailCache&) = delete;

    /// Displayed edge length; also the material render size.
    static constexpr int kThumbPx = 96;

    /// Returns the node's texture thumbnail (0 = not ready). Enqueues a
    /// background decode when resolvedFile is new or has changed. An empty
    /// resolvedFile (unresolvable asset) yields whatever is cached (usually 0).
    ImTextureID GetTextureThumbnail(const pxr::SdfPath& node, const std::string& resolvedFile);

    /// Returns the node's material shader-ball thumbnail (0 = not ready yet).
    /// Marks it stale (to be re-rendered on a later pump) when revision changes.
    /// output/terminal select the node output to preview (a token-typed
    /// terminal is wired as the surface; otherwise routed into diffuseColor).
    ImTextureID GetMaterialThumbnail(const pxr::UsdStageRefPtr& stage,
                                     const pxr::SdfPath& materialPath,
                                     const pxr::SdfPath& node,
                                     const std::string& output, bool terminal,
                                     size_t revision);

    /// Keeps the material shader-ball thumbnails' render delegate and color
    /// correction in sync with the main preview swatch's current settings
    /// (they'd otherwise render under this cache's own independent Hydra
    /// instance's defaults, mismatching the preview and the app's global
    /// color-correction). Cheap to call every frame: no-ops unless something
    /// actually changed, and only then invalidates cached thumbnails so
    /// PumpMainThread re-renders them under the new settings. An empty
    /// rendererPlugin means "not resolved yet" and is left alone.
    void SyncRenderSettings(const pxr::TfToken& rendererPlugin, int ccMode,
                            const std::string& ocioDisplay, const std::string& ocioView,
                            const std::string& ocioColorSpace, const std::string& ocioLook);

    /// Main thread, once per frame, GL context current: drains finished
    /// decodes and uploads them, renders at most one stale material thumbnail,
    /// and prunes entries not requested since the previous pump.
    void PumpMainThread();

    /// Frees every thumbnail (call on material/stage change). Main thread.
    void Clear();

private:
    // ---- texture thumbnails (background decode) ----
    struct TexJob    { std::string nodeKey; std::string file; };
    struct TexResult { std::string nodeKey; std::string file;
                       std::vector<unsigned char> rgba; int w = 0, h = 0; bool ok = false; };
    struct TexEntry  { ImTextureID texId = 0; std::string fileKey; std::string reqFile;
                       bool touched = false; };

    void WorkerLoop();
    static bool DecodeThumbnail(const std::string& file, TexResult& out);
    void UploadTexture(TexEntry& e, const TexResult& r);

    std::thread             m_worker;
    std::atomic<bool>       m_running{false};
    std::mutex              m_jobMutex;
    std::condition_variable m_jobCv;
    std::deque<TexJob>      m_jobs;
    std::mutex              m_resultMutex;
    std::vector<TexResult>  m_results;
    std::unordered_map<std::string, TexEntry> m_texEntries;

    // ---- material thumbnails (main-thread, amortized) ----
    struct MatEntry {
        ImTextureID         texId = 0;
        size_t              revision = 0;
        bool                stale = true;
        bool                touched = false;
        pxr::UsdStageRefPtr stage;
        pxr::SdfPath        materialPath;
        pxr::SdfPath        nodePath;
        std::string         output;
        bool                terminal = false;
    };
    std::unordered_map<std::string, MatEntry> m_matEntries;
    MaterialPreviewRenderer m_matRenderer;
    unsigned int m_readFbo = 0;   // scratch GL FBOs for the render->cache blit
    unsigned int m_drawFbo = 0;

    // Last render settings applied via SyncRenderSettings, so repeated
    // per-frame calls with unchanged values are cheap no-ops.
    bool         m_renderSettingsInitialized = false;
    pxr::TfToken m_lastRendererPlugin;
    int          m_lastCcMode = -1;
    std::string  m_lastOcioDisplay, m_lastOcioView, m_lastOcioColorSpace, m_lastOcioLook;

    void EnsureBlitFbos();
    void BlitToEntry(MatEntry& e, uint32_t srcTex);
    static void DeleteTex(ImTextureID& id);
};

} // namespace UsdLayerManager
