// =============================================
//  Aster: model_loader.h
//  Copyright (c) 2020-2024 Anish Bhobe
// =============================================

#pragma once

#include "buffer.h"
#include "global.h"

#include "render_resource_manager.h"

#include <tiny_gltf.h>

struct TextureHandle;
struct Texture;

constexpr auto GLTF_ASCII_FILE_EXTENSION = ".gltf";
constexpr auto GLTF_BINARY_FILE_EXTENSION = ".glb";

struct MeshPrimitive
{
    u32 m_VertexOffset;
    i32 m_NormalOffset; // <0 for invalid
    i32 m_TexCoord0Offset; // <0 for invalid
    u32 m_FirstIndex;
    u32 m_IndexCount;
    i32 m_MaterialIdx; // <0 for invalid
    i32 m_TransformIdx;
};

struct Nodes
{
    eastl::vector<mat4> m_Transforms;
    eastl::vector<mat4> m_GlobalTransforms;
    /// Parents are also used for bookkeeping
    eastl::vector<u32> m_Parents_;
    bool m_Dirty = true;

    constexpr static u32 ROOT_BIT = 1u << 31;
    constexpr static u32 DIRTY_BIT = 1u << 30;
    constexpr static u32 PARENT_MASK = ~(ROOT_BIT | DIRTY_BIT);

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

        return index;
    }

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

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

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

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

    [[nodiscard]] u32
    Count() const
    {
        return Cast<u32>(m_Transforms.size());
    }

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

    [[nodiscard]] const mat4 *
    GetGlobalTransformPtr() const
    {
        return m_GlobalTransforms.data();
    }

    bool
    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] * *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;
    }
};

struct Material
{
    vec4 m_AlbedoFactor;            // 16   16
    vec3 m_EmissionFactor;          // 12   28
    f32 m_MetalFactor;              // 04   32
    f32 m_RoughFactor;              // 04   36
    TextureHandle m_AlbedoTex;      // 04   40
    TextureHandle m_NormalTex;      // 04   44
    TextureHandle m_MetalRoughTex;  // 04   48
    TextureHandle m_OcclusionTex;   // 04   52
    TextureHandle m_EmissionTex;    // 04   56
};

struct Model
{
    GpuResourceManager *m_ResourceManager;

    eastl::vector<TextureHandle> m_TextureHandles;
    Nodes m_Nodes;

    BufferHandle m_MaterialsHandle;
    BufferHandle m_VertexPositionHandle;
    BufferHandle m_NormalHandle;
    BufferHandle m_TexCoord0Handle;
    BufferHandle m_VertexColorHandle;
    BufferHandle m_NodeHandle;

    IndexBuffer m_IndexBuffer;
    eastl::vector<MeshPrimitive> m_MeshPrimitives;

    [[nodiscard]] const mat4 &GetModelTransform() const;
    void SetModelTransform(const mat4 &transform);
    void Update();

    Model(GpuResourceManager *resourceManager, eastl::vector<TextureHandle> &&textureHandles, Nodes&& nodes,
          BufferHandle materialsHandle, BufferHandle vertexPosHandle, BufferHandle normalHandle, BufferHandle uv0Handle, BufferHandle vertexColor,
          const IndexBuffer &indexBuffer, const eastl::vector<MeshPrimitive> &meshPrimitives, BufferHandle nodeHandle);

    Model(Model &&other) noexcept;
    Model &operator=(Model &&other) noexcept;
    ~Model();

    DISALLOW_COPY_AND_ASSIGN(Model);
};

struct ModelLoader
{
    GpuResourceManager *m_ResourceManager;
    vk::CommandPool m_CommandPool;
    vk::CommandBuffer m_CommandBuffer;
    vk::Queue m_TransferQueue;
    u32 m_TransferQueueIndex;
    u32 m_GraphicsQueueIndex;

    ModelLoader(GpuResourceManager *resourceManager, vk::Queue transferQueue, u32 transferQueueIndex,
                u32 graphicsQueueIndex);
    ~ModelLoader();

    TextureHandle LoadImage(vk::CommandBuffer commandBuffer, StagingBuffer *stagingBuffer,
                            tinygltf::Image *image) const;
    Model LoadModel(cstr path, cstr name = nullptr, bool batched = false);

    constexpr static auto ANormal = "NORMAL";
    constexpr static auto APosition = "POSITION";
    constexpr static auto ATangent = "TANGENT";
    constexpr static auto ATexCoord0 = "TEXCOORD_0";
    constexpr static auto ATexCoord1 = "TEXCOORD_1";
    constexpr static auto AColor0 = "COLOR_0";
    constexpr static auto AJoints0 = "JOINTS_0";
    constexpr static auto AWeights0 = "WEIGHTS_0";
};