SkeletonExtender

#include <set>
#include <algorithm>
#include <iterator>

#include "SkeletonExtender.h"
#include "ShaderUtilities.h"
#include "json\json.h"
#include "interfaces/NiTransformInterface.h"

#include "skse/NiNodes.h"
#include "skse/NiObjects.h"
#include "skse/NiExtraData.h"
#include "skse/NiSerialization.h"

#include "skse/GameReferences.h"
#include "skse/GameStreams.h"
#include "skse/GameRTTI.h"

extern NiTransformInterface g_transformInterface;

void SkeletonExtender::AddTransforms(TESObjectREFR * refr, bool isFirstPerson, NiNode * skeleton, NiAVObject * objectRoot)
{
    std::set<BSFixedString> previous_nodes;
    std::set<BSFixedString> current_nodes;
    std::set<BSFixedString> diffs;

    UInt8 gender = 0;
    TESNPC * actorBase = DYNAMIC_CAST(refr->baseForm, TESForm, TESNPC);
    if (actorBase)
        gender = CALL_MEMBER_FN(actorBase, GetSex)();

    NiStringsExtraData * globalData = ni_cast(FindExtraData(skeleton, "BNDT"), NiStringsExtraData);
    if (globalData)
    {
        for (int i = 0; i < globalData->m_size; i++)
        {
            BSFixedString node(globalData->m_data[i]);
            previous_nodes.insert(node);
        }
    }

    VisitObjects(skeleton, [&](NiAVObject*object)
    {
        NiStringExtraData * stringData = ni_cast(object->GetExtraData(BSFixedString("SDTA").data), NiStringExtraData);
        if (stringData)
        {
            try
            {
                Json::Features features;
                features.all();

                Json::Value root;
                Json::Reader reader(features);

                bool parseSuccess = reader.parse(stringData->m_pString, root);
                if (parseSuccess)
                {
                    for (auto & objects : root)
                    {
                        BSFixedString node = objects["name"].asCString();
                        current_nodes.insert(node);
                    }
                }
            }
            catch (...)
            {
                _ERROR("%s - Error - Failed to parse skeleton transform data", __FUNCTION__);
            }
        }

        return false;
    });

    // Differences here means we lost nodes
    std::set_symmetric_difference(current_nodes.begin(), current_nodes.end(),
                        previous_nodes.begin(), previous_nodes.end(),
                        std::inserter(diffs, diffs.begin()));

    for (auto & node : diffs)
    {
        g_transformInterface.RemoveNodeTransform(refr, isFirstPerson, gender == 1, node, "internal");
    }

    diffs.clear();

    NiStringExtraData * stringData = ni_cast(FindExtraData(objectRoot, "SDTA"), NiStringExtraData);
    ReadTransforms(refr, stringData, isFirstPerson, gender == 1, current_nodes);

    std::set_symmetric_difference(current_nodes.begin(), current_nodes.end(),
                        previous_nodes.begin(), previous_nodes.end(),
                        std::inserter(diffs, diffs.begin()));

    for (auto & node : diffs)
    {
        g_transformInterface.UpdateNodeTransforms(refr, isFirstPerson, gender == 1, node);
    }

    std::vector<BSFixedString> newNodes;
    for (auto & node : current_nodes)
    {
        newNodes.push_back(node);
    }

    // Was already there, set the new current nodes
    if (globalData)
    {
        if (newNodes.size() == 0)
        {
            globalData->SetData(nullptr, 0);
        }
        else
        {
            globalData->SetData(&newNodes.at(0), newNodes.size());
        }

    }

    // No previous nodes, and we have new nodes
    if (!globalData && current_nodes.size() > 0)
    {
        NiStringsExtraData * strData = NiStringsExtraData::Create("BNDT", &newNodes.at(0), newNodes.size());
        skeleton->AddExtraData(strData);
    }
}

void SkeletonExtender::ReadTransforms(TESObjectREFR * refr, NiStringExtraData * stringData, bool isFirstPerson, bool isFemale, std::set<BSFixedString> & nodes)
{
    if (stringData)
    {
        try
        {
            Json::Features features;
            features.all();

            Json::Value root;
            Json::Reader reader(features);

            bool parseSuccess = reader.parse(stringData->m_pString, root);
            if (parseSuccess)
            {
                for (auto & objects : root)
                {
                    BSFixedString node = objects["name"].asCString();
                    nodes.insert(node);

                    Json::Value pos = objects["pos"];
                    if (pos.isArray() && pos.size() == 3)
                    {
                        float position[3];
                        position[0] = pos[0].asFloat();
                        position[1] = pos[1].asFloat();
                        position[2] = pos[2].asFloat();

                        OverrideVariant posV[3];
                        for (UInt32 i = 0; i < 3; i++) {
                            PackValue<float>(&posV[i], OverrideVariant::kParam_NodeTransformPosition, i, &position[i]);
                            g_transformInterface.AddNodeTransform(refr, isFirstPerson, isFemale, node, "internal", posV[i]);
                        }
                    }

                    NiMatrix33 rotation;
                    Json::Value rot = objects["rot"];
                    if (pos.isArray() && pos.size() == 3)
                    {
                        float rotationEuler[3];
                        rotationEuler[0] = rot[0].asFloat() * MATH_PI / 180;
                        rotationEuler[1] = rot[1].asFloat() * MATH_PI / 180;
                        rotationEuler[2] = rot[2].asFloat() * MATH_PI / 180;
                        rotation.SetEulerAngles(rotationEuler[0], rotationEuler[1], rotationEuler[2]);

                        OverrideVariant rotV[9];
                        for (UInt32 i = 0; i < 9; i++) {
                            PackValue<float>(&rotV[i], OverrideVariant::kParam_NodeTransformRotation, i, &rotation.arr[i]);
                            g_transformInterface.AddNodeTransform(refr, isFirstPerson, isFemale, node, "internal", rotV[i]);
                        }
                    }

                    Json::Value scale = objects["scale"];
                    if (scale.isDouble() && scale.asFloat() > 0)
                    {
                        float fScale = scale.asFloat();
                        OverrideVariant scaleV;
                        PackValue<float>(&scaleV, OverrideVariant::kParam_NodeTransformScale, 0, &fScale);
                        g_transformInterface.AddNodeTransform(refr, isFirstPerson, isFemale, node, "internal", scaleV);
                    }

                }
            }
        }
        catch (...)
        {
            _ERROR("%s - Error - Failed to parse skeleton transform data", __FUNCTION__);
        }
    }
}