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

#include <pxr/imaging/hio/image.h>
#include <pxr/imaging/hio/types.h>
#include <pxr/base/gf/half.h>

#include <algorithm>
#include <cmath>
#include <cstdint>

namespace UsdLayerManager {

namespace {
// Mix a graph revision with a node path so switching which node the shared
// preview renderer draws always reads as "changed" to MaterialPreviewRenderer's
// revision gate (it treats the value as an opaque change token), while a real
// graph edit (revision bump) still forces a re-render.
size_t RenderToken(size_t revision, const std::string& nodeKey) {
    size_t h = revision * 1099511628211ull;
    h ^= std::hash<std::string>{}(nodeKey) + 0x9e3779b97f4a7c15ull + (h << 6) + (h >> 2);
    return h;
}
} // namespace

NodeThumbnailCache::NodeThumbnailCache() {
    // Prime Hio's plugin registry on the main thread so the worker never
    // first-touches plugin discovery off-thread.
    pxr::HioImage::IsSupportedImageFile("prime.png");

    m_running = true;
    m_worker = std::thread(&NodeThumbnailCache::WorkerLoop, this);
}

NodeThumbnailCache::~NodeThumbnailCache() {
    m_running = false;
    m_jobCv.notify_all();
    if (m_worker.joinable())
        m_worker.join();
    Clear();
    if (m_readFbo) glDeleteFramebuffers(1, &m_readFbo);
    if (m_drawFbo) glDeleteFramebuffers(1, &m_drawFbo);
}

// ---------------------------------------------------------------------------
// Texture thumbnails
// ---------------------------------------------------------------------------

ImTextureID NodeThumbnailCache::GetTextureThumbnail(const pxr::SdfPath& node,
                                                    const std::string& resolvedFile) {
    const std::string key = node.GetString();
    TexEntry& e = m_texEntries[key];
    e.touched = true;
    if (resolvedFile.empty())
        return e.texId;
    // (Re)request when the file differs from both the uploaded and in-flight one.
    if (resolvedFile != e.fileKey && resolvedFile != e.reqFile) {
        e.reqFile = resolvedFile;
        {
            std::lock_guard<std::mutex> lk(m_jobMutex);
            m_jobs.push_back(TexJob{key, resolvedFile});
        }
        m_jobCv.notify_one();
    }
    return e.texId;
}

void NodeThumbnailCache::WorkerLoop() {
    for (;;) {
        TexJob job;
        {
            std::unique_lock<std::mutex> lk(m_jobMutex);
            m_jobCv.wait(lk, [this] { return !m_running || !m_jobs.empty(); });
            if (!m_running) return;
            job = std::move(m_jobs.front());
            m_jobs.pop_front();
        }
        TexResult result;
        result.nodeKey = job.nodeKey;
        result.file    = job.file;
        result.ok      = DecodeThumbnail(job.file, result);
        {
            std::lock_guard<std::mutex> lk(m_resultMutex);
            m_results.push_back(std::move(result));
        }
    }
}

bool NodeThumbnailCache::DecodeThumbnail(const std::string& file, TexResult& out) {
    pxr::HioImageSharedPtr img = pxr::HioImage::OpenForReading(file, /*subimage=*/0, /*mip=*/0,
                                                              pxr::HioImage::SourceColorSpace::Auto,
                                                              /*suppressErrors=*/true);
    if (!img) return false;
    const int sw = img->GetWidth();
    const int sh = img->GetHeight();
    if (sw <= 0 || sh <= 0) return false;

    // Hio's stb reader does no format conversion: it requires the requested
    // StorageSpec.format to *exactly* equal the file's own format (channel
    // count + component type + sRGB-ness) or it raises "Image format mismatch".
    // So read into the image's native format and normalize to RGBA8 ourselves.
    const pxr::HioFormat fmt   = img->GetFormat();
    const int            nComp = pxr::HioGetComponentCount(fmt);
    const pxr::HioType   type  = pxr::HioGetHioType(fmt);
    if (nComp < 1 || nComp > 4) return false;

    size_t bpc = 0; // bytes per component
    switch (type) {
        case pxr::HioTypeUnsignedByte:
        case pxr::HioTypeUnsignedByteSRGB:
        case pxr::HioTypeSignedByte:    bpc = 1; break;
        case pxr::HioTypeUnsignedShort:
        case pxr::HioTypeSignedShort:
        case pxr::HioTypeHalfFloat:     bpc = 2; break;
        case pxr::HioTypeUnsignedInt:
        case pxr::HioTypeInt:
        case pxr::HioTypeFloat:         bpc = 4; break;
        case pxr::HioTypeDouble:        bpc = 8; break;
        default: return false;
    }

    const size_t pixels = static_cast<size_t>(sw) * sh;
    std::vector<unsigned char> raw(pixels * nComp * bpc);
    pxr::HioImage::StorageSpec spec;
    spec.width  = sw;
    spec.height = sh;
    spec.depth  = 1;
    spec.format = fmt;
    spec.flipped = false;               // keep top row first for ImGui display
    spec.data   = raw.data();
    if (!img->Read(spec)) return false;

    // Collapse whatever component type the file uses down to an 8-bit value.
    // sRGB byte data is passed through as-is (ImGui does no color management,
    // so the encoded bytes display correctly); float/half is clamped to [0,1].
    const void* rp = raw.data();
    auto toU8 = [&](size_t i) -> unsigned char {
        switch (type) {
            case pxr::HioTypeUnsignedByte:
            case pxr::HioTypeUnsignedByteSRGB:
                return reinterpret_cast<const uint8_t*>(rp)[i];
            case pxr::HioTypeSignedByte:
                return static_cast<unsigned char>(std::max(0, static_cast<int>(reinterpret_cast<const int8_t*>(rp)[i]) * 2));
            case pxr::HioTypeUnsignedShort:
                return static_cast<unsigned char>(reinterpret_cast<const uint16_t*>(rp)[i] >> 8);
            case pxr::HioTypeSignedShort:
                return static_cast<unsigned char>(std::clamp(static_cast<int>(reinterpret_cast<const int16_t*>(rp)[i]) >> 7, 0, 255));
            case pxr::HioTypeHalfFloat:
                return static_cast<unsigned char>(std::clamp(static_cast<float>(reinterpret_cast<const pxr::GfHalf*>(rp)[i]), 0.f, 1.f) * 255.f + 0.5f);
            case pxr::HioTypeFloat:
                return static_cast<unsigned char>(std::clamp(reinterpret_cast<const float*>(rp)[i], 0.f, 1.f) * 255.f + 0.5f);
            case pxr::HioTypeUnsignedInt:
                return static_cast<unsigned char>(reinterpret_cast<const uint32_t*>(rp)[i] >> 24);
            case pxr::HioTypeInt:
                return static_cast<unsigned char>(std::clamp(reinterpret_cast<const int32_t*>(rp)[i] >> 23, 0, 255));
            case pxr::HioTypeDouble:
                return static_cast<unsigned char>(std::clamp(reinterpret_cast<const double*>(rp)[i], 0.0, 1.0) * 255.0 + 0.5);
            default: return 0;
        }
    };

    // Expand to RGBA8: 1ch → grey, 2ch → grey+alpha, 3ch → opaque, 4ch → as-is.
    std::vector<unsigned char> native(pixels * 4);
    for (size_t px = 0; px < pixels; ++px) {
        const size_t s = px * nComp;
        const unsigned char c0 = toU8(s);
        unsigned char* d = &native[px * 4];
        d[0] = c0;
        d[1] = nComp >= 3 ? toU8(s + 1) : c0;
        d[2] = nComp >= 3 ? toU8(s + 2) : c0;
        d[3] = nComp == 4 ? toU8(s + 3) : (nComp == 2 ? toU8(s + 1) : 255);
    }

    const int longSide = std::max(sw, sh);
    const float scale  = longSide > kThumbPx ? static_cast<float>(kThumbPx) / longSide : 1.0f;
    const int dw = std::max(1, static_cast<int>(std::lround(sw * scale)));
    const int dh = std::max(1, static_cast<int>(std::lround(sh * scale)));

    if (dw == sw && dh == sh) {
        out.rgba = std::move(native);
    } else {
        out.rgba.assign(static_cast<size_t>(dw) * dh * 4, 0);
        for (int y = 0; y < dh; ++y) {
            const int sy0 = y * sh / dh, sy1 = std::max(sy0 + 1, (y + 1) * sh / dh);
            for (int x = 0; x < dw; ++x) {
                const int sx0 = x * sw / dw, sx1 = std::max(sx0 + 1, (x + 1) * sw / dw);
                unsigned int acc[4] = {0, 0, 0, 0};
                unsigned int n = 0;
                for (int sy = sy0; sy < sy1; ++sy)
                    for (int sx = sx0; sx < sx1; ++sx) {
                        const unsigned char* p = &native[(static_cast<size_t>(sy) * sw + sx) * 4];
                        acc[0] += p[0]; acc[1] += p[1]; acc[2] += p[2]; acc[3] += p[3];
                        ++n;
                    }
                unsigned char* d = &out.rgba[(static_cast<size_t>(y) * dw + x) * 4];
                d[0] = static_cast<unsigned char>(acc[0] / n);
                d[1] = static_cast<unsigned char>(acc[1] / n);
                d[2] = static_cast<unsigned char>(acc[2] / n);
                d[3] = static_cast<unsigned char>(acc[3] / n);
            }
        }
    }
    out.w = dw;
    out.h = dh;
    return true;
}

void NodeThumbnailCache::UploadTexture(TexEntry& e, const TexResult& r) {
    DeleteTex(e.texId);
    GLuint tex = 0;
    glGenTextures(1, &tex);
    glBindTexture(GL_TEXTURE_2D, tex);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
    glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA, r.w, r.h, 0, GL_RGBA, GL_UNSIGNED_BYTE, r.rgba.data());
    glBindTexture(GL_TEXTURE_2D, 0);
    e.texId = static_cast<ImTextureID>(static_cast<ImU64>(tex));
}

// ---------------------------------------------------------------------------
// Material thumbnails
// ---------------------------------------------------------------------------

ImTextureID NodeThumbnailCache::GetMaterialThumbnail(const pxr::UsdStageRefPtr& stage,
                                                     const pxr::SdfPath& materialPath,
                                                     const pxr::SdfPath& node,
                                                     const std::string& output, bool terminal,
                                                     size_t revision) {
    MatEntry& e = m_matEntries[node.GetString()];
    e.touched      = true;
    e.stage        = stage;
    e.materialPath = materialPath;
    e.nodePath     = node;
    e.output       = output;
    e.terminal     = terminal;
    if (revision != e.revision || e.texId == 0) {
        e.revision = revision;
        e.stale    = true;
    }
    return e.texId;
}

void NodeThumbnailCache::EnsureBlitFbos() {
    if (!m_readFbo) glGenFramebuffers(1, &m_readFbo);
    if (!m_drawFbo) glGenFramebuffers(1, &m_drawFbo);
}

void NodeThumbnailCache::BlitToEntry(MatEntry& e, uint32_t srcTex) {
    int sw = 0, sh = 0;
    glBindTexture(GL_TEXTURE_2D, srcTex);
    glGetTexLevelParameteriv(GL_TEXTURE_2D, 0, GL_TEXTURE_WIDTH, &sw);
    glGetTexLevelParameteriv(GL_TEXTURE_2D, 0, GL_TEXTURE_HEIGHT, &sh);
    glBindTexture(GL_TEXTURE_2D, 0);
    if (sw <= 0 || sh <= 0) return;

    if (e.texId == 0) {
        GLuint tex = 0;
        glGenTextures(1, &tex);
        glBindTexture(GL_TEXTURE_2D, tex);
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
        glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA, kThumbPx, kThumbPx, 0, GL_RGBA, GL_UNSIGNED_BYTE, nullptr);
        glBindTexture(GL_TEXTURE_2D, 0);
        e.texId = static_cast<ImTextureID>(static_cast<ImU64>(tex));
    }
    const GLuint dst = static_cast<GLuint>(static_cast<ImU64>(e.texId));

    EnsureBlitFbos();
    GLint prevRead = 0, prevDraw = 0;
    glGetIntegerv(GL_READ_FRAMEBUFFER_BINDING, &prevRead);
    glGetIntegerv(GL_DRAW_FRAMEBUFFER_BINDING, &prevDraw);
    glBindFramebuffer(GL_READ_FRAMEBUFFER, m_readFbo);
    glFramebufferTexture2D(GL_READ_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, GL_TEXTURE_2D, srcTex, 0);
    glBindFramebuffer(GL_DRAW_FRAMEBUFFER, m_drawFbo);
    glFramebufferTexture2D(GL_DRAW_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, GL_TEXTURE_2D, dst, 0);
    // Flip Y (dst rows reversed): the draw target is a GL bottom-up texture, so
    // this leaves the cache texture top-down for uniform ImGui (0,0)-(1,1) UVs.
    glBlitFramebuffer(0, 0, sw, sh, 0, kThumbPx, kThumbPx, 0, GL_COLOR_BUFFER_BIT, GL_LINEAR);
    glBindFramebuffer(GL_READ_FRAMEBUFFER, static_cast<GLuint>(prevRead));
    glBindFramebuffer(GL_DRAW_FRAMEBUFFER, static_cast<GLuint>(prevDraw));
}

// ---------------------------------------------------------------------------
// Per-frame pump / lifecycle
// ---------------------------------------------------------------------------

void NodeThumbnailCache::PumpMainThread() {
    // 1. Upload finished texture decodes.
    std::vector<TexResult> done;
    {
        std::lock_guard<std::mutex> lk(m_resultMutex);
        done.swap(m_results);
    }
    for (const TexResult& r : done) {
        auto it = m_texEntries.find(r.nodeKey);
        if (it == m_texEntries.end()) continue;
        TexEntry& e = it->second;
        if (r.file != e.reqFile) continue; // superseded by a newer request
        if (r.ok) {
            UploadTexture(e, r);
            e.fileKey = r.file;
        }
        e.reqFile.clear();
    }

    // 2. Render at most one stale material thumbnail (amortized).
    for (auto& kv : m_matEntries) {
        MatEntry& e = kv.second;
        if (!e.stale || !e.touched || !e.stage) continue;
        m_matRenderer.SetMaterial(e.stage, e.materialPath, RenderToken(e.revision, kv.first),
                                  e.nodePath, e.output, e.terminal);
        const uint32_t src = m_matRenderer.Render(kThumbPx, kThumbPx);
        if (src != 0) {
            BlitToEntry(e, src);
            e.stale = false;
        }
        break; // one per frame
    }

    // 3. Prune entries not requested since the previous pump (node deleted or
    //    no longer a thumbnailed category), then reset touch flags for this frame.
    for (auto it = m_texEntries.begin(); it != m_texEntries.end();) {
        if (!it->second.touched) { DeleteTex(it->second.texId); it = m_texEntries.erase(it); }
        else { it->second.touched = false; ++it; }
    }
    for (auto it = m_matEntries.begin(); it != m_matEntries.end();) {
        if (!it->second.touched) { DeleteTex(it->second.texId); it = m_matEntries.erase(it); }
        else { it->second.touched = false; ++it; }
    }
}

void NodeThumbnailCache::Clear() {
    for (auto& kv : m_texEntries) DeleteTex(kv.second.texId);
    for (auto& kv : m_matEntries) DeleteTex(kv.second.texId);
    m_texEntries.clear();
    m_matEntries.clear();
}

void NodeThumbnailCache::DeleteTex(ImTextureID& id) {
    if (id) {
        GLuint t = static_cast<GLuint>(static_cast<ImU64>(id));
        glDeleteTextures(1, &t);
        id = 0;
    }
}

} // namespace UsdLayerManager