//
// Created by atdunbg on 2026/9/22.
//
// RenderGraph：声明式帧图内核。
//
// 设计原则（面向引擎，非 demo）：
//   1. 声明与执行分离：图结构（pass + 虚拟资源）只声明一次并 Compile，
//      每帧 Execute 零图结构开销、零堆分配（barrier 列表全部编译期预算）。
//   2. 虚拟资源 + 生成号句柄：上层永远不持有裸 vulkan 句柄，
//      生命周期、layout、物理内存全部由图管理。
//   3. 物理资源池：真实 GpuImage 按 (format, extent, usage) 池化复用，
//      每个虚拟资源按帧槽位拥有 kMaxFramesInFlight 个物理实例，
//      从根上避免帧间读写竞争，同时杜绝每帧重建 GPU 资源。
//   4. 自动同步：layout/访问状态链在 Compile 期解算，产出批量
//      vk::ImageMemoryBarrier2（Synchronization2），上层不再手写 barrier。
//
// 已知边界（后续演进点）：
//   - 暂不支持异步 compute queue / 跨队列信号量
//   - 暂不支持内存别名（需 VMA）与帧历史资源（TAA 等）
//   - 同一 pass 内对同一资源"读+写"（input attachment / RWM）未建模
//

#ifndef HELLOVULKAN_RENDERGRAPH_H
#define HELLOVULKAN_RENDERGRAPH_H

#include <Engine/Core/Types.h>
#include <Engine/Core/FrameContext.h>
#include <Engine/Vulkan/MemoryAllocator.h>
#include <vulkan/vulkan.hpp>

#include <functional>
#include <string>
#include <unordered_map>
#include <vector>

namespace Engine::RG {

    // ---------- 句柄 ----------
    // 带生成号的索引句柄：图永远不会删除条目（只追加），
    // 生成号用于上层误持过期句柄时的校验，防止 use-after-free 式错误。

    inline constexpr u32 kInvalidHandle = 0xFFFFFFFF;

    template<typename Tag>
    struct Handle {
        u32 index = kInvalidHandle;
        u32 generation = 0;

        bool Valid() const { return index != kInvalidHandle; }
        explicit operator bool() const { return Valid(); }
    };

    struct TextureTag;
    using TextureHandle = Handle<TextureTag>;

    // ---------- 资源使用方式 ----------
    // 声明 pass 怎么用资源；layout/access/stage 映射和 image usage
    // 都由这里推导，上层不需要（也不允许）手写。

    enum class ResourceUse : u8 {
        ColorAttachment,
        DepthAttachment,
        ShaderRead,
        TransferSrc,
        TransferDst,
        Present,        // 仅用于 swapchain 导入资源：transition 到 ePresentSrcKHR
    };

    // ---------- 虚拟纹理描述 ----------

    struct TextureDesc {
        vk::Format format = vk::Format::eUndefined;
        u32 width = 0;   // 0 = 跟随目标（swapchain）尺寸
        u32 height = 0;
    };

    // ---------- 物理镜像 ----------

    struct PhysicalImage {
        vk::Image image{};
        vk::ImageView view{};
        vk::Extent2D extent{};
        vk::Format format{};
        vk::ImageUsageFlags usage{};
        MemoryAllocator::Allocation allocation{};
        bool valid = false;
    };

    // ---------- Pass 上下文（Execute 回调收到的） ----------

    struct RenderGraphContext {
        vk::CommandBuffer cmd;
        FrameContext frame;
        vk::Extent2D extent{};

        // 按句柄取当前帧槽位的物理资源（image/view）
        vk::ImageView View(TextureHandle h) const;
        vk::Image Image(TextureHandle h) const;

        class RenderGraph* graph = nullptr;
    };

    using ExecuteFn = std::function<void(RenderGraphContext&)>;

    // ---------- 图内部结构 ----------

    struct TextureEntry {
        std::string name;
        u32 index = 0;                  // 自身在 m_Textures 中的下标
        TextureDesc desc;
        bool imported = false;          // swapchain 导入资源无物理实例
        u32 generation = 1;

        // Compile 期推导
        vk::ImageUsageFlags usage{};
        vk::ImageAspectFlagBits aspect = vk::ImageAspectFlagBits::eColor;
        i32 firstPass = -1;
        i32 lastPass = -1;

        // 物理实例：每个帧槽位一个，created 资源才拥有
        std::vector<PhysicalImage*> physical;

        // 导入资源每帧绑定
        vk::Image importedImage{};
        vk::ImageView importedView{};

        // 帧槽位状态：物理镜像是否已被使用过（决定 entry barrier 变体）
        struct SlotState {
            bool used = false;
        };
        std::vector<SlotState> slots;

        // Compile 期状态链解算用（执行序上的"当前"状态）
        vk::ImageLayout chainLayout = vk::ImageLayout::eUndefined;
        vk::PipelineStageFlags2 chainStage{};
        vk::AccessFlags2 chainAccess{};
        bool chainWasWrite = false;

        // 稳态：上一轮执行结束时（lastPass 之后）的状态。
        // 同一帧槽位再次进入图时，entry barrier 从这里出发。
        vk::ImageLayout restingLayout = vk::ImageLayout::eUndefined;
        vk::PipelineStageFlags2 restingStage{};
        vk::AccessFlags2 restingAccess{};
    };

    struct PassNode {
        std::string name;
        struct Use {
            u32 resIndex;
            ResourceUse use;
            vk::ShaderStageFlags readStage;   // ShaderRead 时的消费 stage
        };
        std::vector<Use> uses;
        ExecuteFn exec;
        bool sideEffect = false;              // 无读写但必须执行（compute 写 buffer 等）
        bool culled = false;

        // Compile 期预算好的 barrier：每个资源两个变体
        //   steady = 该槽位已被使用过（从 resting 状态出发）
        //   first  = 该槽位首次使用（从 Undefined 出发）
        // imported 资源恒用 first（acquire 后 layout 语义上为 Undefined）
        struct BarrierPair {
            u32 resIndex;
            vk::ImageMemoryBarrier2 steady;
            vk::ImageMemoryBarrier2 first;
        };
        std::vector<BarrierPair> barriers;
    };

    // ---------- RenderGraph ----------

    class RenderGraph {
    public:
        RenderGraph(Device& device, MemoryAllocator& allocator, u32 framesInFlight);
        ~RenderGraph();

        RenderGraph(const RenderGraph&) = delete;
        RenderGraph& operator=(const RenderGraph&) = delete;

        // ---- 声明阶段（AddPass 之后不得再声明资源；全部声明完成后 Compile）----

        // 声明一张图拥有的瞬时纹理（此时不分配任何 GPU 内存）
        TextureHandle CreateTexture(std::string name, const TextureDesc& desc);

        // 导入 swapchain 颜色附件（每帧通过 SetImportedImage 绑定物理 image）
        TextureHandle ImportSwapchain(std::string name, vk::Format format);

        // 每帧 acquire 后调用，把当帧 swapchain image 绑到导入资源
        void SetImportedImage(TextureHandle h, vk::Image image, vk::ImageView view);

        // ---- Pass 声明 ----

        class PassBuilder {
        public:
            PassBuilder& Write(TextureHandle h, ResourceUse use);
            PassBuilder& Read(TextureHandle h, ResourceUse use,
                              vk::ShaderStageFlags stage = vk::ShaderStageFlagBits::eFragment);
            PassBuilder& SideEffect();

        private:
            friend class RenderGraph;
            PassBuilder(RenderGraph& graph, u32 passIndex)
                : m_Graph(graph), m_PassIndex(passIndex) {}

            RenderGraph& m_Graph;
            u32 m_PassIndex;
        };

        [[nodiscard]] PassBuilder AddPass(std::string name, ExecuteFn exec);

        // ---- 编译与执行 ----

        // 编译：死 pass 剔除 → 拓扑排序 → 生命周期 → barrier 预算 → 物理分配。
        // 目标尺寸变化时调用 SetTargetExtent 会自动重编译并轮换物理资源。
        void Compile();

        void SetTargetExtent(vk::Extent2D extent);
        vk::Extent2D GetTargetExtent() const { return m_TargetExtent; }

        // 每帧执行：按编译好的顺序录命令。要求 cmd 处于 recording 状态。
        void Execute(const FrameContext& frame);

        // 物理镜像访问器：给"按帧槽位预写描述符"的场景用（如阴影深度图采样）。
        // 仅 created 资源有效；imported 资源（swapchain）每帧才绑定，无物理视图。
        vk::ImageView GetPhysicalView(TextureHandle h, u32 frameSlot) const;

        // ---- 统计 ----

        u32 GetPassCount() const;
        u32 GetActivePassCount() const;
        u32 GetTextureCount() const;
        u32 GetBarrierCount() const;
        u32 GetPooledImageCount() const;

    private:
        friend struct RenderGraphContext;

        // 编译子步骤
        void ResolveUsage(TextureEntry& entry);
        void CullPasses();
        void Toposort();
        void ComputeLifetimes();
        void PrecomputeBarriers();
        void AllocatePhysical();
        void ReleasePhysical();
        void PurgePool(vk::Extent2D target);
        PhysicalImage* AcquirePhysical(vk::Format format, vk::Extent2D extent,
                                       vk::ImageUsageFlags usage,
                                       vk::ImageAspectFlagBits aspect);

        TextureEntry& GetEntry(TextureHandle h);
        const TextureEntry& GetEntry(TextureHandle h) const;
        void CheckHandle(TextureHandle h) const;
        void CheckMutable() const;   // 声明期校验：Compile 后禁止再改图

        static vk::ImageLayout UseLayout(ResourceUse use);
        static vk::PipelineStageFlags2 UseDstStage(const PassNode::Use& use);
        static vk::PipelineStageFlags2 UseSrcStage(ResourceUse use);
        static vk::AccessFlags2 UseDstAccess(const PassNode::Use& use);
        static vk::AccessFlags2 UseSrcAccess(ResourceUse use);

        Device& m_Device;
        MemoryAllocator& m_Allocator;
        u32 m_FramesInFlight;
        bool m_Compiled = false;

        std::vector<TextureEntry> m_Textures;
        std::vector<PassNode> m_Passes;
        u32 m_ActivePassCount = 0;

        vk::Extent2D m_TargetExtent{};

        // 拓扑排序后的执行序（存 pass 下标）
        std::vector<u32> m_ExecutionOrder;

        // 物理资源池：key = (format, extent, usage)，value = 空闲镜像列表
        struct PoolKey {
            u32 format;
            u32 width;
            u32 height;
            u32 usage;
            bool operator==(const PoolKey& o) const {
                return format == o.format && width == o.width &&
                       height == o.height && usage == o.usage;
            }
        };
        struct PoolKeyHash {
            size_t operator()(const PoolKey& k) const {
                size_t h = std::hash<u64>{}(
                    (static_cast<u64>(k.format) << 32) ^ k.usage);
                h ^= std::hash<u64>{}((static_cast<u64>(k.width) << 32) ^ k.height) << 1;
                return h;
            }
        };
        std::unordered_map<PoolKey, std::vector<Scope<PhysicalImage>>, PoolKeyHash> m_Pool;

        // Execute 用的 barrier 暂存（成员变量，避免每帧堆分配）
        std::vector<vk::ImageMemoryBarrier2> m_ScratchBarriers;
    };

} // namespace Engine::RG

#endif //HELLOVULKAN_RENDERGRAPH_H