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

#pragma once

#include "buffer.h"
#include "global.h"
#include "image.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
};

struct Nodes
{
    eastl::vector<mat4> m_Transforms;
    eastl::vector<i32> m_Parents;

    void
    Add(const mat4& transform, i32 parent = -1)
    {
        m_Transforms.push_back(transform);
        m_Parents.push_back(parent);
    }

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

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

    [[nodiscard]] const mat4 *
    GetTransformPtr() const
    {
        return m_Transforms.data();
    }
};

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;

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

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

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

    ~Model();

    DISALLOW_COPY_AND_ASSIGN(Model);
};

struct ModelLoader
{
    GpuResourceManager *const 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";
};