NodeTransformManager

#include "NodeTransformManager.h"
#include "NetImmerseOverride.h"

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

#include "skse/GameStreams.h"

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

extern OverrideMapHolder            g_overrideMapHolder;
extern SKSETaskInterface            * g_task;

void NodeTransformKeys::Save(SKSESerializationInterface * intfc, UInt32 kVersion)
{
    UInt32 numNodes = this->size();
    intfc->WriteRecordData(&numNodes, sizeof(numNodes));

    for (NodeTransformKeys::iterator it = this->begin(); it != this->end(); ++it)
    {
        intfc->OpenRecord('NOTM', kVersion);

        // Key
        WriteKey<BSFixedString>(intfc, it->first);

        // Value
        it->second.Save(intfc, kVersion);
    }
}

bool NodeTransformKeys::Load(SKSESerializationInterface * intfc, UInt32 kVersion)
{
    UInt32 type, length, version;
    bool error = false;

    // Handle count
    UInt32 numRegs = 0;
    if (!intfc->ReadRecordData(&numRegs, sizeof(numRegs)))
    {
        _MESSAGE("Error loading override registration count");
        error = true;
        return error;
    }

    for (UInt32 i = 0; i < numRegs; i++)
    {
        if (intfc->GetNextRecordInfo(&type, &version, &length))
        {
            switch (type)
            {
                case 'NOTM':
                {
                    BSFixedString key;
                    if (ReadKey<BSFixedString>(intfc, key)) {
                        _MESSAGE("Error loading node entry key");
                        error = true;
                        return error;
                    }

                    // operator[] not working for some odd reason
                    bool loadError = false;
                    NodeTransformKeys::iterator iter = this->find(key); // Find existing first
                    if (iter != this->end()) {
                        error = iter->second.Load(intfc, version);
                    }
                    else { // No existing, create
                        OverrideRegistration<BSFixedString> set;
                        error = set.Load(intfc, version);
                        insert(std::pair<BSFixedString, OverrideRegistration<BSFixedString>>(key, set));
                    }
                    if (loadError)
                    {
                        _MESSAGE("Error loading node overrides");
                        error = true;
                        return error;
                    }
                    break;
                }
                default:
                {
                    _MESSAGE("Error loading unexpected chunk type %08X (%.4s)", type, &type);
                    error = true;
                    return error;
                }
            }
        }
    }

    return error;
}

void NodeTransformRegistrationMapHolder::Save(SKSESerializationInterface* intfc, UInt32 kVersion)
{
    for (NodeTransformRegistrationMapHolder::RegMap::iterator it = m_data.begin(); it != m_data.end(); ++it) {
        intfc->OpenRecord('ACTM', kVersion);

        // Key
        UInt64 handle = it->first;
        intfc->WriteRecordData(&handle, sizeof(handle));

#ifdef _DEBUG
        _MESSAGE("%s - Saving Handle %016llX", __FUNCTION__, handle);
#endif

        // Value
        it->second.Save(intfc, kVersion);
    }
}

bool NodeTransformRegistrationMapHolder::Load(SKSESerializationInterface* intfc, UInt32 kVersion, UInt64 * outHandle)
{
    bool error = false;

    UInt64 handle = 0;
    // Key
    if (!intfc->ReadRecordData(&handle, sizeof(handle)))
    {
        _MESSAGE("Error loading reg key");
        error = true;
        return error;
    }

    MultiRegistration<MultiRegistration<NodeTransformKeys, 2>,2> reg;
    if (reg.Load(intfc, kVersion))
    {
        _MESSAGE("Error loading transform gender registrations");
        error = true;
        return error;
    }

    UInt64 newHandle = 0;

    // Skip if handle is no longer valid.
    if (!intfc->ResolveHandle(handle, &newHandle)) {
        *outHandle = 0;
        return error;
    }

    if (reg.empty()) {
        *outHandle = 0;
        return error;
    }

    *outHandle = newHandle;

    Lock();
    m_data[newHandle] = reg;
    Release();

#ifdef _DEBUG
    _MESSAGE("%s - Loaded Handle %016llX", __FUNCTION__, newHandle);
#endif

    //SetHandleProperties(newHandle, false);
    return error;
}

void NodeTransformManager::Save(SKSESerializationInterface * intfc, UInt32 kVersion)
{
    transformData.Save(intfc, kVersion);
}
bool NodeTransformManager::Load(SKSESerializationInterface* intfc, UInt32 kVersion)
{
    UInt64 handle = 0;
    if (!transformData.Load(intfc, kVersion, &handle))
    {
        //
        SetHandleNodeTransforms(handle);
    }

    return false;
}

bool NodeTransformManager::AddNodeTransform(TESObjectREFR * refr, bool firstPerson, bool isFemale, BSFixedString node, BSFixedString name, OVRValue & value)
{
    UInt64 handle = g_overrideMapHolder.GetHandle(refr, refr->formType);
    transformData.Lock();
    transformData.m_data[handle][isFemale ? 1 : 0][firstPerson ? 1 : 0][node][name].erase(value);
    transformData.m_data[handle][isFemale ? 1 : 0][firstPerson ? 1 : 0][node][name].insert(value);
    transformData.Release();
    return true;
}


bool NodeTransformManager::RemoveNodeTransform(TESObjectREFR * refr, bool firstPerson, bool isFemale, BSFixedString node, BSFixedString name)
{
    UInt8 gender = isFemale ? 1 : 0;
    UInt8 fp = firstPerson ? 1 : 0;
    UInt64 handle = g_overrideMapHolder.GetHandle(refr, refr->formType);
    auto it = transformData.m_data.find(handle);
    if (it != transformData.m_data.end())
    {
        auto ait = it->second[gender][fp].find(node);
        if (ait != it->second[gender][fp].end())
        {
            auto oit = ait->second.find(name);
            if (oit != ait->second.end())
            {
                ait->second.erase(oit);
                return true;
            }
        }
    }

    return false;
}

bool NodeTransformManager::RemoveNodeTransformComponent(TESObjectREFR * refr, bool firstPerson, bool isFemale, BSFixedString node, BSFixedString name, UInt16 key, UInt16 index)
{
    UInt8 gender = isFemale ? 1 : 0;
    UInt8 fp = firstPerson ? 1 : 0;
    UInt64 handle = g_overrideMapHolder.GetHandle(refr, refr->formType);
    auto it = transformData.m_data.find(handle);
    if (it != transformData.m_data.end())
    {
        auto ait = it->second[gender][fp].find(node);
        if (ait != it->second[gender][fp].end())
        {
            auto oit = ait->second.find(name);
            if (oit != ait->second.end())
            {
                OVRValue ovr;
                ovr.key = key;
                ovr.index = index;
                auto ost = oit->second.find(ovr);
                if (ost != oit->second.end())
                {
                    oit->second.erase(ost);
                    return true;
                }
            }
        }
    }

    return false;
}

bool NodeTransformManager::VisitNodeTransforms(TESObjectREFR * refr, bool firstPerson, bool isFemale, BSFixedString node, std::function<bool(NiAVObject*, NiTransform*, OverrideRegistration<BSFixedString>*)> functor)
{
    bool ret = false;
    UInt8 gender = isFemale ? 1 : 0;
    UInt8 fp = firstPerson ? 1 : 0;
    UInt64 handle = g_overrideMapHolder.GetHandle(refr, refr->formType);
    auto it = transformData.m_data.find(handle); // Find ActorHandle
    if (it != transformData.m_data.end())
    {
        NiNode * root = refr->GetNiRootNode(fp);
        if (root) {
            root->IncRef();
            NiAVObject * foundNode = root->GetObjectByName(&node.data);
            if (foundNode) {
                auto nodeIt = it->second[gender][fp].find(node);
                if (nodeIt != it->second[gender][fp].end()) {
                    NiTransform * baseTransform = transformCache.GetBaseTransform(GetRootModelPath(refr, firstPerson, isFemale), node);
                    if (baseTransform)
                        ret = functor(foundNode, baseTransform, &nodeIt->second);
                }
            }
            root->DecRef();
        }
    }

    return ret;
}

void NodeTransformManager::UpdateNodeTransforms(TESObjectREFR * ref, bool firstPerson, bool isFemale, BSFixedString node)
{
    VisitNodeTransforms(ref, firstPerson, isFemale, node, [&](NiAVObject* foundNode, NiTransform* baseTransform, OverrideRegistration<BSFixedString>* keys)
    {
        NiTransform combinedTransform;
        for (auto dit = keys->begin(); dit != keys->end(); ++dit) {// Loop Keys
            NiTransform localTransform;
            GetOverrideTransform(&dit->second, OVRValue::kParam_NodeTransformPosition, &localTransform);
            GetOverrideTransform(&dit->second, OVRValue::kParam_NodeTransformScale, &localTransform);
            GetOverrideTransform(&dit->second, OVRValue::kParam_NodeTransformRotation, &localTransform);
            combinedTransform = combinedTransform * localTransform;
        }

        foundNode->m_localTransform = (*baseTransform) * combinedTransform;

        if (g_task)
            g_task->AddTask(new NIOVTaskUpdateWorldData(foundNode));

        return false;
    });
}

bool NodeTransformManager::GetOverrideNodeTransform(TESObjectREFR * refr, bool firstPerson, bool isFemale, BSFixedString node, BSFixedString name, UInt16 key, NiTransform * result)
{
    return VisitNodeTransforms(refr, firstPerson, isFemale, node, [&](NiAVObject* foundNode, NiTransform* baseTransform, OverrideRegistration<BSFixedString>* keys)
    {
        auto it = keys->find(name);
        if (it != keys->end()) {
            GetOverrideTransform(&it->second, key, result);
            return true;
        }

        return false;
    });
}

void NodeTransformManager::UpdateNodeAllTransforms(TESObjectREFR * refr)
{
    UInt64 handle = g_overrideMapHolder.GetHandle(refr, refr->formType);
    SetHandleNodeTransforms(handle);
}

void NodeTransformManager::SetHandleNodeTransforms(UInt64 handle)
{
    TESObjectREFR * refr = (TESObjectREFR *)g_overrideMapHolder.GetObject(handle, TESObjectREFR::kTypeID);
    if (!refr) {
        transformData.Lock();
        transformData.m_data.erase(handle);
        transformData.Release();
        return;
    }

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


    auto it = transformData.m_data.find(handle); // Find ActorHandle
    if (it != transformData.m_data.end())
    {
        NiNode * lastNode = NULL;
        for (UInt8 i = 0; i <= 1; i++)
        {
            NiNode * root = refr->GetNiRootNode(i);
            if (root == lastNode) // First and Third are the same, skip
                continue;

            if (root)
            {
                root->IncRef();
                for (auto ait = it->second[gender][i].begin(); ait != it->second[gender][i].end(); ++ait) // Loop Nodes
                {
                    NiTransform * baseTransform = transformCache.GetBaseTransform(GetRootModelPath(refr, i >= 1 ? true : false, gender >= 1 ? true : false), ait->first);
                    if (baseTransform)
                    {
                        NiTransform combinedTransform;
                        for (auto dit = ait->second.begin(); dit != ait->second.end(); ++dit) {// Loop Keys
                            NiTransform localTransform;
                            GetOverrideTransform(&dit->second, OVRValue::kParam_NodeTransformPosition, &localTransform);
                            GetOverrideTransform(&dit->second, OVRValue::kParam_NodeTransformScale, &localTransform);
                            GetOverrideTransform(&dit->second, OVRValue::kParam_NodeTransformRotation, &localTransform);
                            combinedTransform = combinedTransform * localTransform;
                        }

                        BSFixedString nodeName = ait->first;
                        NiAVObject * transformable = root->GetObjectByName(&nodeName.data);
                        if (transformable) {
                            transformable->m_localTransform = (*baseTransform) * combinedTransform;
                        }
                    }
                }
                root->DecRef();
            }

            lastNode = root;

            if (g_task)
                g_task->AddTask(new NIOVTaskUpdateWorldData(root));
        }
    }
}

void NodeTransformManager::GetOverrideTransform(OverrideSet * set, UInt16 key, NiTransform * result)
{
    OVRValue value;
    OverrideSet::iterator it;
    switch (key) {
        case OVRValue::kParam_NodeTransformPosition:
        {
            value.key = OVRValue::kParam_NodeTransformPosition;
            value.index = 0;
            it = set->find(value);
            if (it != set->end()) {
                result->pos.x = it->data.f;
            }
            value.index = 1;
            it = set->find(value);
            if (it != set->end()) {
                result->pos.y = it->data.f;
            }
            value.index = 2;
            it = set->find(value);
            if (it != set->end()) {
                result->pos.y = it->data.f;
            }
            break;
        }
        case OVRValue::kParam_NodeTransformScale:
        {
            value.key = OVRValue::kParam_NodeTransformScale;
            value.index = 0;
            it = set->find(value);
            if (it != set->end()) {
                result->scale = it->data.f;
            }
        }
        break;
        case OVRValue::kParam_NodeTransformRotation:
        {
            value.key = OVRValue::kParam_NodeTransformRotation;
            value.index = 0;
            it = set->find(value);
            if (it != set->end()) {
                result->rot.data[0][0] = it->data.f;
            }
            value.index = 1;
            it = set->find(value);
            if (it != set->end()) {
                result->rot.data[0][1] = it->data.f;
            }
            value.index = 2;
            it = set->find(value);
            if (it != set->end()) {
                result->rot.data[0][2] = it->data.f;
            }
            value.index = 3;
            it = set->find(value);
            if (it != set->end()) {
                result->rot.data[1][0] = it->data.f;
            }
            value.index = 4;
            it = set->find(value);
            if (it != set->end()) {
                result->rot.data[1][1] = it->data.f;
            }
            value.index = 5;
            it = set->find(value);
            if (it != set->end()) {
                result->rot.data[1][2] = it->data.f;
            }
            value.index = 6;
            it = set->find(value);
            if (it != set->end()) {
                result->rot.data[2][0] = it->data.f;
            }
            value.index = 7;
            it = set->find(value);
            if (it != set->end()) {
                result->rot.data[2][1] = it->data.f;
            }
            value.index = 8;
            it = set->find(value);
            if (it != set->end()) {
                result->rot.data[2][2] = it->data.f;
            }
        }
        break;
    }
}


NiTransform * NodeTransformCache::GetBaseTransform(BSFixedString rootModel, BSFixedString nodeName)
{
    SimpleLocker<std::unordered_map<BSFixedString, std::unordered_map<BSFixedString, NiTransform>>> lock(this);
    auto it = m_data.find(rootModel);
    if (it != m_data.end()) {
        auto nodeIt = it->second.find(nodeName);
        if (nodeIt != it->second.end()) {
            return &nodeIt->second;
        }
    } else {
        char pathBuffer[MAX_PATH];
        memset(pathBuffer, 0, MAX_PATH);
        sprintf_s(pathBuffer, MAX_PATH, "meshes\\%s", rootModel.data);

        BSFixedString newPath(pathBuffer);

        // No skeleton path found, why is this?
        BSResourceNiBinaryStream binaryStream(newPath.data);
        if (!binaryStream.IsValid()) {
            return NULL;
        }

        std::unordered_map<BSFixedString, NiTransform> transformMap;
        NiTransform * foundTransform = NULL;

        UInt8 niStreamMemory[0x5B4];
        memset(niStreamMemory, 0, 0x5B4);
        NiStream * niStream = (NiStream *)niStreamMemory;
        CALL_MEMBER_FN(niStream, ctor)();

        niStream->LoadStream(&binaryStream);
        if (niStream->m_rootObjects.m_data)
        {
            for (UInt32 i = 0; i < niStream->m_rootObjects.m_size; i++) {
                NiObject * object = niStream->m_rootObjects.m_data[i];
                if (object) {
                    NiAVObject * node = ni_cast(object, NiAVObject);
                    if (node) {
                        VisitObjects(node, [&](NiAVObject* child)
                        {
                            BSFixedString localName(node->m_name);
                            auto ret = transformMap.insert(std::make_pair(localName, node->m_localTransform));
                            if (ret.second && localName == nodeName) {
                                foundTransform = &ret.first->second;
                            }
                            return false;
                        });
                    }
                }
            }
        }

        CALL_MEMBER_FN(niStream, dtor)();
        m_data.insert(std::make_pair(rootModel, transformMap));
        return foundTransform;
    }

    return NULL;
}

BSFixedString NodeTransformManager::GetRootModelPath(TESObjectREFR * refr, bool firstPerson, bool isFemale)
{
    TESModel * model = NULL;
    Character * character = DYNAMIC_CAST(refr, TESObjectREFR, Character);
    if (character) {
        if (firstPerson) {
            Setting * setting = (*g_gameSettingCollection)->Get("sFirstPersonSkeleton");
            if (setting && setting->GetType() == Setting::kType_String)
                return BSFixedString(setting->data.s);
        }

        TESRace * race = character->race;
        if (!race) {
            TESNPC * actorBase = DYNAMIC_CAST(refr->baseForm, TESForm, TESNPC);
            if (actorBase)
                race = actorBase->race.race;
        }

        if (race)
            model = &race->models[isFemale ? 1 : 0];
    }
    else
        model = DYNAMIC_CAST(refr->baseForm, TESForm, TESModel);

    if (model)
        return BSFixedString(model->GetModelName());

    return BSFixedString("");
}