#include "TransformCommand.h"
#include "../../utils/Logger.h"
#include <pxr/usd/usd/editContext.h>
#include <pxr/usd/usdGeom/xformOp.h>

namespace UsdLayerManager {

// True when attr has an authored time sample exactly at time t.
static bool HasSampleAtTime(const pxr::UsdAttribute& attr, double t)
{
    if (!attr) return false;
    std::vector<double> times;
    if (!attr.GetTimeSamples(&times)) return false;
    for (double s : times)
        if (s == t) return true;
    return false;
}

// Attribute name of the rotate op XformCommonAPI authors for a rotation order
// (e.g. "xformOp:rotateXYZ").
static pxr::TfToken RotateOpName(pxr::UsdGeomXformCommonAPI::RotationOrder ro)
{
    return pxr::UsdGeomXformOp::GetOpName(
        pxr::UsdGeomXformCommonAPI::ConvertRotationOrderToOpType(ro));
}

TransformCommand::SamplePresence TransformCommand::QuerySamplePresence(
    pxr::UsdStageRefPtr                       stage,
    const pxr::SdfPath&                       primPath,
    pxr::UsdGeomXformCommonAPI::RotationOrder rotOrder,
    pxr::UsdTimeCode                          timeCode)
{
    SamplePresence presence;
    if (timeCode.IsDefault() || !stage) return presence;
    pxr::UsdPrim prim = stage->GetPrimAtPath(primPath);
    if (!prim) return presence;

    double t = timeCode.GetValue();
    presence.translate = HasSampleAtTime(
        prim.GetAttribute(pxr::UsdGeomXformOp::GetOpName(
            pxr::UsdGeomXformOp::TypeTranslate)), t);
    presence.rotate = HasSampleAtTime(
        prim.GetAttribute(RotateOpName(rotOrder)), t);
    presence.scale = HasSampleAtTime(
        prim.GetAttribute(pxr::UsdGeomXformOp::GetOpName(
            pxr::UsdGeomXformOp::TypeScale)), t);
    return presence;
}

TransformCommand::TransformCommand(pxr::UsdStageRefPtr       stage,
                                   const pxr::SdfPath&        primPath,
                                   pxr::SdfLayerHandle        editLayer,
                                   const pxr::GfVec3d&        oldTranslate,
                                   const pxr::GfVec3f&        oldRotate,
                                   const pxr::GfVec3f&        oldScale,
                                   const pxr::GfVec3d&        newTranslate,
                                   const pxr::GfVec3f&        newRotate,
                                   const pxr::GfVec3f&        newScale,
                                   pxr::UsdGeomXformCommonAPI::RotationOrder rotOrder,
                                   pxr::UsdTimeCode           timeCode,
                                   SamplePresence             preEditSamples,
                                   std::string                description)
    : m_stage(stage)
    , m_primPath(primPath)
    , m_editLayer(editLayer)
    , m_oldTranslate(oldTranslate)
    , m_oldRotate(oldRotate)
    , m_oldScale(oldScale)
    , m_newTranslate(newTranslate)
    , m_newRotate(newRotate)
    , m_newScale(newScale)
    , m_rotOrder(rotOrder)
    , m_timeCode(timeCode)
    , m_description(std::move(description))
    , m_preEditSamples(preEditSamples)
{
}

void TransformCommand::Execute() {
    Apply(m_newTranslate, m_newRotate, m_newScale);
}

void TransformCommand::Undo() {
    Apply(m_oldTranslate, m_oldRotate, m_oldScale);
    ClearCreatedSamples();
}

void TransformCommand::Apply(const pxr::GfVec3d& t,
                              const pxr::GfVec3f& r,
                              const pxr::GfVec3f& s)
{
    if (!m_stage || m_primPath.IsEmpty()) return;
    pxr::UsdPrim prim = m_stage->GetPrimAtPath(m_primPath);
    if (!prim) return;

    pxr::UsdGeomXformCommonAPI api(prim);
    if (!api) return;

    try {
        if (m_editLayer) {
            pxr::UsdEditContext ec(m_stage, m_editLayer);
            api.SetTranslate(t, m_timeCode);
            api.SetRotate(r, m_rotOrder, m_timeCode);
            api.SetScale(s, m_timeCode);
        } else {
            api.SetTranslate(t, m_timeCode);
            api.SetRotate(r, m_rotOrder, m_timeCode);
            api.SetScale(s, m_timeCode);
        }
    } catch (const std::exception& e) {
        LOG_ERROR(std::string("TransformCommand::Apply error: ") + e.what());
    }
}

void TransformCommand::ClearCreatedSamples()
{
    if (m_timeCode.IsDefault()) return;
    if (!m_stage || m_primPath.IsEmpty()) return;
    pxr::UsdPrim prim = m_stage->GetPrimAtPath(m_primPath);
    if (!prim) return;

    auto clearOp = [&](const pxr::TfToken& opName, bool hadSample) {
        if (hadSample) return;
        pxr::UsdAttribute attr = prim.GetAttribute(opName);
        if (attr) attr.ClearAtTime(m_timeCode);
    };

    auto clearAll = [&]() {
        clearOp(pxr::UsdGeomXformOp::GetOpName(
                    pxr::UsdGeomXformOp::TypeTranslate), m_preEditSamples.translate);
        clearOp(RotateOpName(m_rotOrder), m_preEditSamples.rotate);
        clearOp(pxr::UsdGeomXformOp::GetOpName(
                    pxr::UsdGeomXformOp::TypeScale), m_preEditSamples.scale);
    };

    try {
        if (m_editLayer) {
            pxr::UsdEditContext ec(m_stage, m_editLayer);
            clearAll();
        } else {
            clearAll();
        }
    } catch (const std::exception& e) {
        LOG_ERROR(std::string("TransformCommand::ClearCreatedSamples error: ") + e.what());
    }
}

} // namespace UsdLayerManager