#include "DisconnectShaderAttrCommand.h"
#include "../../utils/Logger.h"
#include <pxr/usd/usd/editContext.h>
#include <pxr/usd/usdShade/shader.h>
#include <pxr/usd/usdShade/connectableAPI.h>
#include <pxr/base/tf/token.h>

namespace UsdLayerManager {

DisconnectShaderAttrCommand::DisconnectShaderAttrCommand(pxr::UsdStageRefPtr stage,
                                                         const pxr::SdfPath&  destNode,
                                                         const std::string&   destInput,
                                                         pxr::SdfLayerHandle  editLayer)
    : m_stage(stage)
    , m_destNode(destNode)
    , m_destInput(destInput)
    , m_editLayer(editLayer)
    , m_description("Disconnect " + destNode.GetName() + "." + destInput)
{
    if (!stage) return;

    pxr::UsdShadeShader destShader(stage->GetPrimAtPath(destNode));
    if (!destShader) return;

    pxr::UsdShadeInput input = destShader.GetInput(pxr::TfToken(destInput));
    if (!input) return;

    for (const auto& source : pxr::UsdShadeConnectableAPI::GetConnectedSources(input)) {
        if (!source.IsValid()) continue;
        m_hadConnection  = true;
        m_sourceNode     = source.source.GetPrim().GetPath();
        m_sourceOutput   = source.sourceName.GetString();
        break;
    }
}

void DisconnectShaderAttrCommand::Execute() {
    if (!m_stage) return;
    try {
        pxr::UsdShadeShader destShader(m_stage->GetPrimAtPath(m_destNode));
        if (!destShader) return;
        pxr::UsdShadeInput input = destShader.GetInput(pxr::TfToken(m_destInput));
        if (!input) return;

        if (m_editLayer) {
            pxr::UsdEditContext ec(m_stage, m_editLayer);
            pxr::UsdShadeConnectableAPI::DisconnectSource(input);
        } else {
            pxr::UsdShadeConnectableAPI::DisconnectSource(input);
        }
    } catch (const std::exception& e) {
        LOG_ERROR(std::string("DisconnectShaderAttrCommand::Execute error: ") + e.what());
    }
}

void DisconnectShaderAttrCommand::Undo() {
    if (!m_stage || !m_hadConnection) return;
    try {
        pxr::UsdShadeShader destShader(m_stage->GetPrimAtPath(m_destNode));
        pxr::UsdShadeShader sourceShader(m_stage->GetPrimAtPath(m_sourceNode));
        if (!destShader || !sourceShader) return;

        pxr::UsdShadeInput input = destShader.GetInput(pxr::TfToken(m_destInput));
        pxr::UsdShadeOutput output = sourceShader.GetOutput(pxr::TfToken(m_sourceOutput));
        if (!input || !output) return;

        if (m_editLayer) {
            pxr::UsdEditContext ec(m_stage, m_editLayer);
            pxr::UsdShadeConnectableAPI::ConnectToSource(input, output);
        } else {
            pxr::UsdShadeConnectableAPI::ConnectToSource(input, output);
        }
    } catch (const std::exception& e) {
        LOG_ERROR(std::string("DisconnectShaderAttrCommand::Undo error: ") + e.what());
    }
}

} // namespace UsdLayerManager