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

#pragma once

#include "aster/aster.h"

#include "aster/core/buffer.h"

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

namespace systems
{
class TransferContext;
}

namespace systems
{
class RenderingDevice;
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
    vec4 m_EmissionFactor;                       // 16   32
    systems::ResId<TextureView> m_AlbedoTex;     // 08   40
    systems::ResId<TextureView> m_NormalTex;     // 08   48
    systems::ResId<TextureView> m_MetalRoughTex; // 08   56
    systems::ResId<TextureView> m_OcclusionTex;  // 08   64
    systems::ResId<TextureView> m_EmissionTex;   // 08   72
    f32 m_MetalFactor;                           // 04   76
    f32 m_RoughFactor;                           // 04   80
};

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
{
    eastl::vector<systems::ResId<TextureView>> m_TextureHandles;
    Nodes m_Nodes;

    struct ModelHandlesData
    {
        uptr m_VertexPositionHandle;
        uptr m_VertexDataHandle;
        uptr m_MaterialsHandle;
        uptr m_NodeHandle;
    };

    struct ModelHandles
    {
        Ref<Buffer> m_VertexPositionHandle;
        Ref<Buffer> m_VertexDataHandle;
        Ref<Buffer> m_MaterialsHandle;
        Ref<Buffer> m_NodeHandle;

        operator ModelHandlesData() const
        {
            return {
                .m_VertexPositionHandle = m_VertexPositionHandle->GetDeviceAddress(),
                .m_VertexDataHandle = m_VertexDataHandle->GetDeviceAddress(),
                .m_MaterialsHandle = m_MaterialsHandle->GetDeviceAddress(),
                .m_NodeHandle = m_NodeHandle->GetDeviceAddress(),
            };
        }
    } m_Handles;

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

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

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

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

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

struct AssetLoader
{
    systems::RenderingDevice *m_Device;

    Ref<TextureView> LoadHdrImage(cstr path, cstr name = nullptr) 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";

    explicit AssetLoader(systems::RenderingDevice &device);

  private:
    systems::ResId<TextureView>
    LoadImageToGpu(systems::TransferContext &context, tinygltf::Image *image, bool isSrgb, cstr name = nullptr) const;
};

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

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

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