// =============================================
//  Aster: nodes.cpp
//  Copyright (c) 2020-2024 Anish Bhobe
// =============================================

#include "nodes.h"

u32
Nodes::Add(const mat4 &transform, const i32 parent)
{
    m_Dirty = true;
    const u32 index = Count();
    m_Transforms.push_back(transform);
    m_GlobalTransforms.emplace_back(transform);
    const u32 parentVal = (parent < 0 ? ROOT_BIT : parent & PARENT_MASK) | DIRTY_BIT;
    m_Parents_.push_back(parentVal);

    return index;
}

const mat4 &
Nodes::Get(const u32 index) const
{
    return m_Transforms[index];
}

void
Nodes::Set(const u32 index, const mat4 &transform)
{
    m_Dirty = true;
    m_Transforms[index] = transform;
    m_Parents_[index] |= DIRTY_BIT;
}

const mat4 &
Nodes::operator[](const u32 index) const
{
    return m_Transforms[index];
}

mat4 &
Nodes::operator[](const u32 index)
{
    m_Dirty = true;
    m_Parents_[index] |= DIRTY_BIT;
    return m_Transforms[index];
}

u32
Nodes::Count() const
{
    return Cast<u32>(m_Transforms.size());
}

usize
Nodes::GetGlobalTransformByteSize() const
{
    return m_GlobalTransforms.size() * sizeof m_GlobalTransforms[0];
}

const Nodes::Transform *
Nodes::GetGlobalTransformPtr() const
{
    return m_GlobalTransforms.data();
}

bool
Nodes::Update()
{
    if (!m_Dirty)
        return false;

    auto transformIter = m_Transforms.begin();
    auto globalTransformIter = m_GlobalTransforms.begin();
    auto parentIter = m_Parents_.begin();
    const auto parentEnd = m_Parents_.end();

    while (parentIter != parentEnd)
    {
        const bool isRoot = *parentIter & ROOT_BIT;
        const bool isDirty = *parentIter & DIRTY_BIT;

        if (isRoot)
        {
            if (isDirty)
            {
                // Copy-update if the root is dirty.
                *globalTransformIter = *transformIter;
            }
        }
        else
        {
            const u32 parentIdx = *parentIter & PARENT_MASK;
            const bool isParentDirty = m_Parents_[parentIdx] & DIRTY_BIT;
            if (isDirty || isParentDirty)
            {
                // Update w.r.t parent if either local or parent transforms updated.
                *globalTransformIter = m_GlobalTransforms[parentIdx].m_GlobalTransforms * *transformIter;
                m_Parents_[parentIdx] |= DIRTY_BIT; // Set dirty to propagate the update.
            }
        }

        ++parentIter;
        ++globalTransformIter;
        ++transformIter;
    }

    for (u32 &parentValue : m_Parents_)
    {
        parentValue &= ~DIRTY_BIT; // Unset dirty.
    }

    m_Dirty = false;
    return true;
}