Extra2D/include/renderer/material.h

#pragma once

#include <renderer/rhi/rhi.h>
#include <renderer/shader.h>
#include <renderer/texture.h>
#include <types/ptr/ref_counted.h>
#include <types/ptr/intrusive_ptr.h>
#include <types/math/vec2.h>
#include <types/math/color.h>
#include <string>
#include <unordered_map>
#include <vector>

namespace extra2d {

// 前向声明
class Material;

/**
 * @brief 材质参数类型
 */
enum class MaterialParamType : uint8_t {
    Float,      // 浮点数
    Vec2,       // 二维向量
    Vec3,       // 三维向量
    Vec4,       // 四维向量
    Color,      // 颜色
    Mat4,       // 4x4 矩阵
    Texture     // 纹理
};

/**
 * @brief 获取材质参数类型的大小
 * @param type 参数类型
 * @return 参数大小（字节）
 */
inline uint32_t getMaterialParamSize(MaterialParamType type) {
    switch (type) {
        case MaterialParamType::Float: return sizeof(float);
        case MaterialParamType::Vec2:  return sizeof(float) * 2;
        case MaterialParamType::Vec3:  return sizeof(float) * 3;
        case MaterialParamType::Vec4:  return sizeof(float) * 4;
        case MaterialParamType::Color: return sizeof(float) * 4;
        case MaterialParamType::Mat4:  return sizeof(float) * 16;
        case MaterialParamType::Texture: return sizeof(uint32_t);
        default: return 0;
    }
}

/**
 * @brief 材质参数信息
 */
struct MaterialParamInfo {
    MaterialParamType type;
    uint32_t offset;
    uint32_t size;
};

/**
 * @brief 纹理槽位信息
 */
struct TextureSlot {
    Ptr<Texture> texture;
    uint32_t slot;
    std::string uniformName;
};

/**
 * @brief 材质布局类
 * 
 * 定义材质参数的布局，可被多个材质共享
 */
class MaterialLayout : public RefCounted {
public:
    /**
     * @brief 默认构造函数
     */
    MaterialLayout();
    
    /**
     * @brief 添加参数
     * @param name 参数名称
     * @param type 参数类型
     */
    void addParam(const std::string& name, MaterialParamType type);
    
    /**
     * @brief 完成布局定义，计算偏移和总大小
     */
    void finalize();
    
    /**
     * @brief 获取参数信息
     * @param name 参数名称
     * @return 参数信息指针，不存在返回 nullptr
     */
    const MaterialParamInfo* getParam(const std::string& name) const;
    
    /**
     * @brief 获取缓冲区大小
     * @return 缓冲区大小
     */
    uint32_t getBufferSize() const { return bufferSize_; }
    
    /**
     * @brief 检查是否已最终化
     * @return 是否已最终化
     */
    bool isFinalized() const { return finalized_; }

private:
    // 使用vector保持参数添加顺序，确保与着色器中的声明顺序一致
    std::vector<std::pair<std::string, MaterialParamInfo>> params_;
    std::unordered_map<std::string, size_t> paramIndexMap_;  // 用于快速查找
    uint32_t bufferSize_ = 0;
    bool finalized_ = false;
};

/**
 * @brief 材质类
 * 
 * 基于 RHI 的材质包装类，管理：
 * - 着色器程序
 * - 材质参数（通过 UBO 上传）
 * - 纹理绑定
 * 
 * 使用示例：
 * @code
 * // 创建自定义着色器
 * auto shader = makePtr<Shader>();
 * shader->loadFromFile("custom.vert", "custom.frag");
 * 
 * // 创建材质布局
 * auto layout = makePtr<MaterialLayout>();
 * layout->addParam("uTintColor", MaterialParamType::Color);
 * layout->addParam("uOpacity", MaterialParamType::Float);
 * layout->finalize();
 * 
 * // 创建材质
 * auto material = makePtr<Material>();
 * material->setShader(shader);
 * material->setLayout(layout);
 * material->setColor("uTintColor", Color::White);
 * material->setFloat("uOpacity", 1.0f);
 * 
 * // 设置纹理
 * material->setTexture("uTexture", texture, 0);
 * @endcode
 */
class Material : public RefCounted {
public:
    /**
     * @brief 默认构造函数
     */
    Material();
    
    // ========================================
    // 着色器和布局设置
    // ========================================
    
    /**
     * @brief 设置材质布局
     * @param layout 材质布局
     */
    void setLayout(Ptr<MaterialLayout> layout);
    
    /**
     * @brief 设置着色器
     * @param shader 着色器
     */
    void setShader(Ptr<Shader> shader);
    
    /**
     * @brief 获取着色器
     * @return 着色器
     */
    Ptr<Shader> getShader() const { return shader_; }
    
    /**
     * @brief 获取 RHI 管线句柄
     * @return RHI 管线句柄
     */
    PipelineHandle getPipeline() const;
    
    // ========================================
    // 参数设置
    // ========================================
    
    /**
     * @brief 设置 float 参数
     * @param name 参数名称
     * @param value 值
     */
    void setFloat(const std::string& name, float value);
    
    /**
     * @brief 设置 vec2 参数
     * @param name 参数名称
     * @param value 值
     */
    void setVec2(const std::string& name, const Vec2& value);
    
    /**
     * @brief 设置 vec4 参数
     * @param name 参数名称
     * @param x X 分量
     * @param y Y 分量
     * @param z Z 分量
     * @param w W 分量
     */
    void setVec4(const std::string& name, float x, float y, float z, float w);
    
    /**
     * @brief 设置颜色参数
     * @param name 参数名称
     * @param value 颜色值
     */
    void setColor(const std::string& name, const Color& value);
    
    /**
     * @brief 设置 mat4 参数
     * @param name 参数名称
     * @param value 矩阵数据指针
     */
    void setMat4(const std::string& name, const float* value);
    
    // ========================================
    // 纹理管理
    // ========================================
    
    /**
     * @brief 设置纹理
     * @param uniformName 着色器中的采样器 uniform 名称
     * @param texture 纹理
     * @param slot 纹理槽位（0-15）
     */
    void setTexture(const std::string& uniformName, Ptr<Texture> texture, uint32_t slot);
    
    /**
     * @brief 获取所有纹理槽位
     * @return 纹理槽位列表
     */
    const std::vector<TextureSlot>& getTextures() const { return textures_; }
    
    /**
     * @brief 清除所有纹理
     */
    void clearTextures();
    
    // ========================================
    // 数据访问
    // ========================================
    
    /**
     * @brief 获取材质数据指针
     * @return 数据指针
     */
    const void* getData() const { return data_.data(); }
    
    /**
     * @brief 获取材质数据大小
     * @return 数据大小
     */
    uint32_t getDataSize() const { return static_cast<uint32_t>(data_.size()); }
    
    /**
     * @brief 获取材质布局
     * @return 材质布局
     */
    Ptr<MaterialLayout> getLayout() const { return layout_; }

private:
    Ptr<MaterialLayout> layout_;    // 材质布局
    Ptr<Shader> shader_;            // 着色器
    std::vector<uint8_t> data_;     // 材质数据（UBO 缓冲）
    std::vector<TextureSlot> textures_;  // 纹理槽位列表
};

} // namespace extra2d