#include "GroupPrimsCommand.h"
#include "../../utils/Logger.h"
#include <pxr/usd/usd/namespaceEditor.h>
#include <pxr/usd/usd/prim.h>
#include <pxr/base/tf/token.h>

PXR_NAMESPACE_USING_DIRECTIVE

namespace UsdLayerManager {

GroupPrimsCommand::GroupPrimsCommand(UsdStageRefPtr              stage,
                                     const std::vector<SdfPath>& srcPaths,
                                     const SdfPath&              groupPath)
    : m_stage(stage)
    , m_groupPath(groupPath)
    , m_srcPaths(srcPaths)
    , m_description("Group " + std::to_string(srcPaths.size()) +
                     " prim(s) under " + groupPath.GetString())
{
}

/// Compute a unique child name under @p parent (collision check against live stage).
static std::string UniqueNameUnder(UsdStageRefPtr stage,
                                   const SdfPath& parent,
                                   const std::string& baseName) {
    std::string name = baseName;
    int n = 1;
    while (stage->GetPrimAtPath(parent.AppendChild(TfToken(name))).IsValid())
        name = baseName + "_" + std::to_string(n++);
    return name;
}

static UsdPrim GetParentPrim(UsdStageRefPtr stage, const SdfPath& parentPath) {
    if (parentPath == SdfPath("/") || parentPath == SdfPath::AbsoluteRootPath())
        return stage->GetPseudoRoot();
    return stage->GetPrimAtPath(parentPath);
}

void GroupPrimsCommand::Execute() {
    // Create the Xform group.
    UsdPrim groupPrim = m_stage->DefinePrim(m_groupPath, TfToken("Xform"));
    if (!groupPrim.IsValid()) {
        LOG_ERROR("GroupPrimsCommand::Execute: failed to create group: " + m_groupPath.GetString());
        return;
    }

    m_movedNames.clear();

    for (const SdfPath& srcPath : m_srcPaths) {
        UsdPrim src = m_stage->GetPrimAtPath(srcPath);
        if (!src.IsValid()) {
            LOG_WARNING("GroupPrimsCommand::Execute: prim not found (skipping): " + srcPath.GetString());
            m_movedNames.emplace_back(srcPath.GetName());
            continue;
        }

        // Compute a collision-free name inside the group.
        std::string name = UniqueNameUnder(m_stage, m_groupPath, srcPath.GetName());
        m_movedNames.push_back(name);

        UsdPrim gPrim = m_stage->GetPrimAtPath(m_groupPath);
        UsdNamespaceEditor editor(m_stage);
        editor.ReparentPrim(src, gPrim, TfToken(name));
        std::string whyNot;
        if (!editor.CanApplyEdits(&whyNot)) {
            LOG_ERROR("GroupPrimsCommand::Execute: cannot reparent " + srcPath.GetString() + ": " + whyNot);
            continue;
        }
        if (!editor.ApplyEdits()) {
            LOG_ERROR("GroupPrimsCommand::Execute: ApplyEdits failed for " + srcPath.GetString());
        }
    }
}

void GroupPrimsCommand::Undo() {
    if (m_movedNames.size() != m_srcPaths.size()) {
        LOG_ERROR("GroupPrimsCommand::Undo: state mismatch, cannot undo safely");
        return;
    }

    // Reparent children back in reverse order.
    for (int i = (int)m_srcPaths.size() - 1; i >= 0; --i) {
        SdfPath movedPath = m_groupPath.AppendChild(TfToken(m_movedNames[i]));
        UsdPrim moved = m_stage->GetPrimAtPath(movedPath);
        if (!moved.IsValid()) {
            LOG_WARNING("GroupPrimsCommand::Undo: prim not in group (skipping): " + movedPath.GetString());
            continue;
        }

        SdfPath origParentPath = m_srcPaths[i].GetParentPath();
        UsdPrim origParent = GetParentPrim(m_stage, origParentPath);
        if (!origParent.IsValid()) {
            LOG_ERROR("GroupPrimsCommand::Undo: original parent not found: " + origParentPath.GetString());
            continue;
        }

        UsdNamespaceEditor editor(m_stage);
        editor.ReparentPrim(moved, origParent, TfToken(m_srcPaths[i].GetName()));
        std::string whyNot;
        if (!editor.CanApplyEdits(&whyNot)) {
            LOG_ERROR("GroupPrimsCommand::Undo: cannot reparent back " + movedPath.GetString() + ": " + whyNot);
            continue;
        }
        if (!editor.ApplyEdits()) {
            LOG_ERROR("GroupPrimsCommand::Undo: ApplyEdits failed for " + movedPath.GetString());
        }
    }

    // Remove the (now empty) group prim.
    m_stage->RemovePrim(m_groupPath);
}

} // namespace UsdLayerManager