#include "MaterialManager.h"
#include <pxr/usd/sdr/registry.h>
#include <pxr/usd/sdr/shaderProperty.h>
#include <pxr/base/tf/stringUtils.h>
#include <pxr/usd/usd/primRange.h>
#include <pxr/usd/usdShade/material.h>
#include <pxr/usd/usdShade/shader.h>
#include <pxr/usd/usdShade/nodeGraph.h>
#include <pxr/usd/usdShade/connectableAPI.h>
#include <pxr/usd/usdShade/utils.h>
#include <pxr/base/tf/token.h>
#include <pxr/base/vt/value.h>
#include <algorithm>
#include <cctype>
#include <set>

namespace UsdLayerManager {

namespace {

// Coarse Hypershade-style bucket for the create-node list. Sdr metadata is
// parser-dependent (glslfx vs MaterialX vs Arnold vs Cycles fill different
// fields), so classify from the strongest signal available: context first,
// then terminal output types, then role, then name keywords.
std::string DeriveCategory(pxr::SdrShaderNodeConstPtr node) {
    const std::string context = pxr::TfStringify(node->GetContext());
    if (context == "surface" || context == "volume" || context == "displacement")
        return "Material";
    if (context == "light" || context == "lightFilter")
        return "Light";

    // MaterialX/Arnold terminals often lack a context but type their outputs
    // as surfaceshader/volumeshader/displacementshader/closure.
    for (const pxr::TfToken& outName : node->GetShaderOutputNames()) {
        if (const auto* prop = node->GetShaderOutput(outName)) {
            const std::string outType = pxr::TfStringify(prop->GetType());
            if (outType.find("shader") != std::string::npos ||
                outType == "terminal" || outType == "closure")
                return "Material";
        }
    }

    const std::string role = pxr::TfStringify(node->GetRole());
    if (role == "texture") return "Texture";
    if (role == "primvar") return "Geometry";
    if (role == "math")    return "Utility";

    const std::string name = pxr::TfStringToLower(
        node->GetIdentifier().GetString() + " " + node->GetName());
    for (const char* kw : {"texture", "image", "checker", "noise", "ramp",
                           "fractal", "worley", "voronoi", "triplanar"})
        if (name.find(kw) != std::string::npos) return "Texture";
    for (const char* kw : {"primvar", "texcoord", "geomprop", "position",
                           "normal", "tangent", "bitangent", "uv_"})
        if (name.find(kw) != std::string::npos) return "Geometry";

    return "Utility";
}

// Fixed display order for the known buckets; anything unexpected sorts after.
int CategoryRank(const std::string& category) {
    if (category == "Material") return 0;
    if (category == "Texture")  return 1;
    if (category == "Geometry") return 2;
    if (category == "Utility")  return 3;
    if (category == "Light")    return 4;
    return 5;
}

// Top-level group from the Sdr source type. glslfx and OSL collapse into one
// "USD" group — the core usd* nodes register under both parsers and would
// otherwise show up twice under different headers.
std::string DeriveSource(pxr::SdrShaderNodeConstPtr node) {
    std::string st = pxr::TfStringToLower(pxr::TfStringify(node->GetSourceType()));
    if (st == "glslfx" || st == "osl") return "USD";
    if (st == "mtlx")                  return "MaterialX";
    if (st == "arnold")                return "Arnold";
    if (st == "cycles")                return "Cycles";
    if (st.empty())                    return "Other";
    st[0] = static_cast<char>(std::toupper(static_cast<unsigned char>(st[0])));
    return st;
}

int SourceRank(const std::string& source) {
    if (source == "USD")       return 0;
    if (source == "MaterialX") return 1;
    return 2; // renderers and anything else, alphabetical among themselves
}

} // namespace

const std::vector<ShaderNodeTypeInfo>& MaterialManager::GetAvailableShaderNodes() {
    if (m_shaderNodeCacheBuilt)
        return m_shaderNodeCache;

    for (pxr::SdrShaderNodeConstPtr node : pxr::SdrRegistry::GetInstance().GetShaderNodesByFamily()) {
        if (!node) continue;

        ShaderNodeTypeInfo info;
        info.identifier = node->GetIdentifier().GetString();
        info.label      = !node->GetLabel().IsEmpty() ? node->GetLabel().GetString() : node->GetName();
        info.family     = node->GetFamily().GetString();
        info.category   = DeriveCategory(node);
        info.source     = DeriveSource(node);
        m_shaderNodeCache.push_back(std::move(info));
    }

    std::sort(m_shaderNodeCache.begin(), m_shaderNodeCache.end(),
              [](const ShaderNodeTypeInfo& a, const ShaderNodeTypeInfo& b) {
                  const int sa = SourceRank(a.source), sb = SourceRank(b.source);
                  if (sa != sb) return sa < sb;
                  if (a.source != b.source) return a.source < b.source;
                  const int ra = CategoryRank(a.category), rb = CategoryRank(b.category);
                  if (ra != rb) return ra < rb;
                  if (a.category != b.category) return a.category < b.category;
                  return a.label < b.label;
              });

    m_shaderNodeCacheBuilt = true;
    return m_shaderNodeCache;
}

std::vector<pxr::SdfPath> MaterialManager::GetAllMaterials() const {
    std::vector<pxr::SdfPath> result;
    if (!m_stage) return result;
    for (const pxr::UsdPrim& prim : m_stage->Traverse()) {
        if (prim.IsA<pxr::UsdShadeMaterial>())
            result.push_back(prim.GetPath());
    }
    return result;
}

ShaderGraphSnapshot MaterialManager::GetShaderGraph(const pxr::SdfPath& materialPath) const {
    ShaderGraphSnapshot snapshot;
    if (!m_stage) return snapshot;

    pxr::UsdPrim materialPrim = m_stage->GetPrimAtPath(materialPath);
    if (!materialPrim.IsValid()) return snapshot;

    static const pxr::TfToken kUiPositionKey("uiPosition");
    static const pxr::TfToken kUiPinDisplayModeKey("uiPinDisplayMode");

    // Seed with every shader under the material, recursing into nested node
    // graphs (usdMtlx nests a material's nodes inside UsdShadeNodeGraph
    // scopes rather than authoring them as direct children).
    std::vector<pxr::UsdPrim> pending;
    std::vector<pxr::UsdPrim> scopes{materialPrim};
    while (!scopes.empty()) {
        pxr::UsdPrim scope = scopes.back();
        scopes.pop_back();
        for (const pxr::UsdPrim& child : scope.GetChildren()) {
            if (child.IsA<pxr::UsdShadeShader>())
                pending.push_back(child);
            else if (child.IsA<pxr::UsdShadeNodeGraph>())
                scopes.push_back(child);
        }
    }

    // Worklist walk: resolving each input's true upstream shader output pulls
    // in nodes living in node graphs outside the material prim too.
    std::set<pxr::SdfPath> visited;
    while (!pending.empty()) {
        pxr::UsdPrim child = pending.back();
        pending.pop_back();
        if (!visited.insert(child.GetPath()).second)
            continue;
        pxr::UsdShadeShader shader(child);
        if (!shader) continue;

        ShaderGraphNode node;
        node.path = child.GetPath();

        pxr::TfToken shaderId;
        shader.GetIdAttr().Get(&shaderId);
        node.shaderId = shaderId.GetString();

        pxr::VtValue posValue = child.GetCustomDataByKey(kUiPositionKey);
        if (posValue.IsHolding<pxr::GfVec2f>()) {
            node.uiPosition = posValue.UncheckedGet<pxr::GfVec2f>();
            node.hasAuthoredPosition = true;
        }

        pxr::VtValue pinModeValue = child.GetCustomDataByKey(kUiPinDisplayModeKey);
        if (pinModeValue.IsHolding<int>()) {
            int raw = pinModeValue.UncheckedGet<int>();
            if (raw >= static_cast<int>(PinDisplayMode::All) && raw <= static_cast<int>(PinDisplayMode::Collapsed))
                node.pinDisplayMode = static_cast<PinDisplayMode>(raw);
        }

        // Prefer the full Sdr-defined pin set (so unauthored pins can still be
        // dragged to create a connection); fall back to authored attributes
        // only for shader types the registry doesn't know about.
        pxr::SdrShaderNodeConstPtr sdrNode =
            pxr::SdrRegistry::GetInstance().GetShaderNodeByIdentifier(shaderId);
        if (sdrNode) {
            node.category = DeriveCategory(sdrNode);
            for (const pxr::TfToken& inputName : sdrNode->GetShaderInputNames()) {
                if (auto* prop = sdrNode->GetShaderInput(inputName)) {
                    ShaderPinInfo pin{inputName.GetString(), prop->GetTypeAsSdfType().GetSdfType()};
                    for (const auto& option : prop->GetOptions())
                        pin.allowedValues.push_back(option.first.GetString());
                    node.inputs.push_back(std::move(pin));
                }
            }
            for (const pxr::TfToken& outputName : sdrNode->GetShaderOutputNames()) {
                if (auto* prop = sdrNode->GetShaderOutput(outputName))
                    node.outputs.push_back({outputName.GetString(), prop->GetTypeAsSdfType().GetSdfType()});
            }
        } else {
            for (const auto& input : shader.GetInputs())
                node.inputs.push_back({input.GetBaseName().GetString(), input.GetTypeName()});
            for (const auto& output : shader.GetOutputs())
                node.outputs.push_back({output.GetBaseName().GetString(), output.GetTypeName()});
        }

        // Only authored inputs can carry a connection. Resolve each one
        // through node-graph boundaries (shaderOutputsOnly) so a MaterialX
        // connection routed via NodeGraph interface attrs links straight to
        // the shader output that actually produces the value.
        for (const auto& input : shader.GetInputs()) {
            for (const pxr::UsdAttribute& sourceAttr :
                 pxr::UsdShadeUtils::GetValueProducingAttributes(input, /*shaderOutputsOnly=*/true)) {
                pxr::UsdPrim sourcePrim = sourceAttr.GetPrim();
                if (!sourcePrim.IsA<pxr::UsdShadeShader>()) continue;
                ShaderGraphLink link;
                link.destNode      = child.GetPath();
                link.destInput     = input.GetBaseName().GetString();
                link.sourceNode    = sourcePrim.GetPath();
                link.sourceOutput  = pxr::UsdShadeOutput(sourceAttr).GetBaseName().GetString();
                snapshot.links.push_back(std::move(link));
                pending.push_back(sourcePrim); // may live outside the material
            }
        }

        snapshot.nodes.push_back(std::move(node));
    }

    return snapshot;
}

} // namespace UsdLayerManager
