#include "ConnectShaderAttrsCommand.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 {

ConnectShaderAttrsCommand::ConnectShaderAttrsCommand(pxr::UsdStageRefPtr stage,
                                                     const pxr::SdfPath&  destNode,
                                                     const std::string&   destInput,
                                                     const pxr::SdfValueTypeName& destType,
                                                     const pxr::SdfPath&  sourceNode,
                                                     const std::string&   sourceOutput,
                                                     const pxr::SdfValueTypeName& sourceType,
                                                     pxr::SdfLayerHandle  editLayer)
    : m_stage(stage)
    , m_destNode(destNode)
    , m_destInput(destInput)
    , m_destType(destType)
    , m_sourceNode(sourceNode)
    , m_sourceOutput(sourceOutput)
    , m_sourceType(sourceType)
    , m_editLayer(editLayer)
    , m_description("Connect " + sourceNode.GetName() + "." + sourceOutput +
                    " -> " + destNode.GetName() + "." + destInput)
{
    if (!stage) return;

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

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

    for (const auto& source : pxr::UsdShadeConnectableAPI::GetConnectedSources(existingInput)) {
        if (!source.IsValid()) continue;
        m_hadPriorConnection = true;
        m_priorSourceNode    = source.source.GetPrim().GetPath();
        m_priorSourceOutput  = source.sourceName.GetString();
        break; // shader inputs carry a single upstream connection in this editor's model
    }
}

void ConnectShaderAttrsCommand::Execute() {
    if (!m_stage) return;
    try {
        pxr::UsdShadeShader destShader(m_stage->GetPrimAtPath(m_destNode));
        pxr::UsdShadeShader sourceShader(m_stage->GetPrimAtPath(m_sourceNode));
        if (!destShader || !sourceShader) {
            LOG_ERROR("ConnectShaderAttrsCommand: invalid shader prim(s)");
            return;
        }

        if (m_editLayer) {
            pxr::UsdEditContext ec(m_stage, m_editLayer);
            pxr::UsdShadeInput destInput = destShader.CreateInput(pxr::TfToken(m_destInput), m_destType);
            pxr::UsdShadeOutput sourceOutput = sourceShader.CreateOutput(pxr::TfToken(m_sourceOutput), m_sourceType);
            pxr::UsdShadeConnectableAPI::ConnectToSource(destInput, sourceOutput);
        } else {
            pxr::UsdShadeInput destInput = destShader.CreateInput(pxr::TfToken(m_destInput), m_destType);
            pxr::UsdShadeOutput sourceOutput = sourceShader.CreateOutput(pxr::TfToken(m_sourceOutput), m_sourceType);
            pxr::UsdShadeConnectableAPI::ConnectToSource(destInput, sourceOutput);
        }
    } catch (const std::exception& e) {
        LOG_ERROR(std::string("ConnectShaderAttrsCommand::Execute error: ") + e.what());
    }
}

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

        auto doUndo = [&]() {
            if (m_hadPriorConnection) {
                pxr::UsdShadeShader priorSourceShader(m_stage->GetPrimAtPath(m_priorSourceNode));
                pxr::UsdShadeOutput priorSourceOutput = priorSourceShader.GetOutput(pxr::TfToken(m_priorSourceOutput));
                if (priorSourceShader && priorSourceOutput)
                    pxr::UsdShadeConnectableAPI::ConnectToSource(destInput, priorSourceOutput);
            } else {
                pxr::UsdShadeConnectableAPI::DisconnectSource(destInput);
            }
        };

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

} // namespace UsdLayerManager
