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

#pragma once

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

#include "gpu_resource_manager.h"
#include "nodes.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 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(StagingBuffer *stagingBuffer, tinygltf::Image *image, bool isSrgb) 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);
};