// =============================================
//  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;
    u32 m_FirstIndex;
    u32 m_IndexCount;
    i32 m_MaterialIdx; // <0 for invalid
    i32 m_TransformIdx;
};

struct Nodes
{
    struct Transform
    {
        mat4 m_GlobalTransforms;
        mat4 m_NormalTransforms;

        explicit Transform(const mat4& transform)
            : m_GlobalTransforms(transform)
            , m_NormalTransforms(transpose(inverse(mat3{transform})))
        {
        }

        Transform &
        operator=(const mat4& transform)
        {
            m_GlobalTransforms = transform;
            m_NormalTransforms = transpose(inverse(mat3{transform}));
            return *this;
        }
    };

    eastl::vector<mat4> m_Transforms;
    eastl::vector<Transform> 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);
    [[nodiscard]] const mat4 &Get(const u32 index) const;
    void Set(const u32 index, const mat4 &transform);
    [[nodiscard]] u32 Count() const;

    [[nodiscard]] const mat4 &operator[](const u32 index) const;
    [[nodiscard]] mat4 &operator[](const u32 index);

    [[nodiscard]] usize GetGlobalTransformByteSize() const;
    [[nodiscard]] const Transform *GetGlobalTransformPtr() const;

    bool Update();
};

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 VertexData
{
    vec4 m_Normal;
    vec2 m_TexCoord0 = vec2{0.0f, 0.0f};
    vec2 m_TexCoord1 = vec2{0.0f, 0.0f};
    vec4 m_Color0 = vec4{1.0f, 0.0f, 1.0f, 1.0f};
};

struct Model
{
    GpuResourceManager *m_ResourceManager;

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

    struct ModelHandles
    {
        BufferHandle m_VertexPositionHandle;
        BufferHandle m_VertexDataHandle;
        BufferHandle m_MaterialsHandle;
        BufferHandle m_NodeHandle;
    } m_Handles;

    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,
          const ModelHandles& handles, const IndexBuffer &indexBuffer, const eastl::vector<MeshPrimitive> &meshPrimitives);
    ~Model();

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

    Model(const Model &) = delete;
    const Model &operator=(const Model &) = delete;
};

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

    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";

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

    ModelLoader(ModelLoader &&other) noexcept;
    ModelLoader &operator=(ModelLoader &&other) noexcept;

    DISALLOW_COPY_AND_ASSIGN(ModelLoader);
};