#ifndef RENDER_IBL_BAKER_H
#define RENDER_IBL_BAKER_H

#include "texture.h"

#include <DirectXMath.h>
#include <d3d11.h>
#include <filesystem>
#include <memory>
#include <wrl/client.h>

// [新增] 把导入的环境 TextureCube 在 GPU 上一次性烘焙为：
// 1. 漫反射 irradiance cube；
// 2. 按 roughness 分布在 mip 链上的 GGX prefiltered cube；
// 3. Split-Sum BRDF 二维 LUT。
class IBLBaker final {
public:
    struct Settings {
        UINT irradianceSize = 32;
        UINT prefilteredSize = 128;
        UINT brdfLutSize = 256;

        UINT irradianceSampleCount = 128;
        UINT prefilterSampleCount = 256;
        UINT brdfSampleCount = 256;
    };

    struct Result {
        std::unique_ptr<TextureCube> irradiance;
        std::unique_ptr<TextureCube> prefiltered;
        std::unique_ptr<Texture2D> brdfLut;

        [[nodiscard]] bool IsValid() const noexcept {
            return irradiance && prefiltered && brdfLut;
        }
    };

public:
    [[nodiscard]] bool Initialize(
        ID3D11Device *device,
        const std::filesystem::path &bakeShaderPath);

    [[nodiscard]] bool Bake(
        ID3D11DeviceContext *context,
        const TextureCube &sourceEnvironment,
        Result &outResult);

    [[nodiscard]] bool Bake(
        ID3D11DeviceContext *context,
        const TextureCube &sourceEnvironment,
        Result &outResult,
        const Settings &settings);

private:
    struct alignas(16) BakeConstantBufferData {
        // x=roughness, y=源 Cube 尺寸, z=源 Cube 最大 mip, w=保留。
        DirectX::XMFLOAT4 floatParams{};
        // x=目标 face, y=采样数, z/w=保留。
        DirectX::XMUINT4 uintParams{};
    };
    static_assert(sizeof(BakeConstantBufferData) % 16 == 0);

private:
    [[nodiscard]] bool UpdateConstantBuffer(
        ID3D11DeviceContext *context,
        float roughness,
        float sourceSize,
        float sourceMaxMip,
        UINT faceIndex,
        UINT sampleCount) const;

    [[nodiscard]] bool RenderCube(
        ID3D11DeviceContext *context,
        const TextureCube &sourceEnvironment,
        TextureCube &destination,
        ID3D11PixelShader *pixelShader,
        UINT sampleCount,
        bool roughnessFromMip) const;

    [[nodiscard]] bool RenderBrdfLut(
        ID3D11DeviceContext *context,
        Texture2D &destination,
        UINT size,
        UINT sampleCount) const;

    [[nodiscard]] static UINT CalculateMipCount(UINT size) noexcept;

private:
    Microsoft::WRL::ComPtr<ID3D11Device> m_device;

    Microsoft::WRL::ComPtr<ID3D11VertexShader> m_fullscreenVS;
    Microsoft::WRL::ComPtr<ID3D11PixelShader> m_irradiancePS;
    Microsoft::WRL::ComPtr<ID3D11PixelShader> m_prefilterPS;
    Microsoft::WRL::ComPtr<ID3D11PixelShader> m_brdfLutPS;

    Microsoft::WRL::ComPtr<ID3D11Buffer> m_constantBuffer;
    Microsoft::WRL::ComPtr<ID3D11SamplerState> m_sourceSampler;
    Microsoft::WRL::ComPtr<ID3D11RasterizerState> m_rasterizerState;
};

#endif // RENDER_IBL_BAKER_H
