// =============================================
//  Aster: asset_loader.h
//  Copyright (c) 2020-2025 Anish Bhobe
// =============================================

#pragma once

#include "aster/aster.h"

#include "aster/core/buffer.h"

#include "aster/systems/image_manager.h"
#include "aster/systems/resource.h"
#include "aster/systems/view_manager.h"
#include "nodes.h"
#include "tiny_gltf.h"

namespace systems
{
class ResourceManager;
class SamplerManager;
class BufferManager;
class ImageManager;
class CommitManager;
} // namespace systems

namespace tinygltf
{
struct Image;
}

struct Image;
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
    systems::ResId<TextureView> m_AlbedoTex;     // 04   40
    systems::ResId<TextureView> m_NormalTex;     // 04   44
    systems::ResId<TextureView> m_MetalRoughTex; // 04   48
    systems::ResId<TextureView> m_OcclusionTex;  // 04   52
    systems::ResId<TextureView> 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, 1.0f, 1.0f, 1.0f};
};

struct Model
{
    systems::CommitManager *m_ResourceManager;

    eastl::vector<systems::ResId<TextureView>> m_TextureHandles;
    Nodes m_Nodes;

    struct ModelHandles
    {
        systems::ResId<Buffer> m_VertexPositionHandle;
        systems::ResId<Buffer> m_VertexDataHandle;
        systems::ResId<Buffer> m_MaterialsHandle;
        systems::ResId<Buffer> m_NodeHandle;
    } m_Handles;

    Ref<Buffer> m_NodeBuffer;
    Ref<Buffer> m_IndexBuffer;
    eastl::vector<MeshPrimitive> m_MeshPrimitives;

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

    Model(systems::CommitManager *resourceManager, eastl::vector<systems::ResId<TextureView>> &textureHandles,
          Nodes &&nodes, Ref<Buffer> nodeBuffer, ModelHandles &handles, Ref<Buffer> indexBuffer,
          const eastl::vector<MeshPrimitive> &meshPrimitives);
    ~Model() = default;

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

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

struct AssetLoader
{
    systems::ResourceManager *m_ResourceManager;
    systems::CommitManager *m_CommitManager;

    vk::CommandPool m_CommandPool;
    vk::CommandBuffer m_CommandBuffer;
    vk::Queue m_TransferQueue;
    u32 m_TransferQueueIndex;
    u32 m_GraphicsQueueIndex;

    Ref<TextureView> LoadHdrImage(cstr path, cstr name = nullptr) const;
    std::tuple<systems::ResId<TextureView>, Ref<Buffer>> LoadImageToGpu(tinygltf::Image *image, bool isSrgb) const;
    Model LoadModelToGpu(cstr path, cstr name = nullptr);

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

    AssetLoader(systems::ResourceManager *resourceManager, systems::CommitManager *commitManager,
                vk::Queue transferQueue, u32 transferQueueIndex, u32 graphicsQueueIndex);
    ~AssetLoader();

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

    DISALLOW_COPY_AND_ASSIGN(AssetLoader);
};

void
GenerateMipMaps(vk::CommandBuffer commandBuffer, const Ref<Texture> &textureView, vk::ImageLayout initialLayout,
                vk::ImageLayout finalLayout, vk::PipelineStageFlags2 prevStage, vk::PipelineStageFlags2 finalStage);

void
GenerateMipMaps(vk::CommandBuffer commandBuffer, concepts::SampledImageRef auto &texture, vk::ImageLayout initialLayout,
                vk::ImageLayout finalLayout,
                vk::PipelineStageFlags2 prevStage = vk::PipelineStageFlagBits2::eAllCommands,
                vk::PipelineStageFlags2 finalStage = vk::PipelineStageFlagBits2::eAllCommands)
{
    GenerateMipMaps(commandBuffer, systems::CastImage<Texture>(texture), initialLayout, finalLayout, prevStage,
                    finalStage);
}

void
GenerateMipMaps(vk::CommandBuffer commandBuffer, concepts::SampledImageViewRef auto &texture,
                vk::ImageLayout initialLayout, vk::ImageLayout finalLayout,
                vk::PipelineStageFlags2 prevStage = vk::PipelineStageFlagBits2::eAllCommands,
                vk::PipelineStageFlags2 finalStage = vk::PipelineStageFlagBits2::eAllCommands)
{
    GenerateMipMaps(commandBuffer, systems::CastImage<Texture>(texture->m_Image), initialLayout, finalLayout, prevStage,
                    finalStage);
}

static_assert(concepts::SampledImageRef<Ref<Texture>>);