/* Copyright (c) 2024-2026 The Khronos Group Inc.
 * Copyright (c) 2024-2026 Valve Corporation
 * Copyright (c) 2024-2026 LunarG, Inc.
 * Modifications Copyright (C) 2025-2026 Advanced Micro Devices, Inc. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

#pragma once
#include <vulkan/vulkan.h>
#include <memory>
#include "containers/custom_containers.h"

namespace vku {
struct safe_VkPipelineShaderStageCreateInfo;
struct safe_VkShaderCreateInfoEXT;
struct safe_VkSpecializationInfo;
}  // namespace vku

namespace vvl {
struct ShaderModule;
struct DescriptorSetLayoutList;
}  // namespace vvl

namespace spirv {
struct Module;
struct EntryPoint;
class Instruction;
}  // namespace spirv

// This is a wrapper around the Shader as it may come from a Pipeline or Shader Object.
//
// TODO - This class has some very nasty flaws...
// A good example is Descriptor Heap flags, we have a nice |descriptor_heap_mode| member we want to query, because it requires the
// pipelineFlags2 pNext struct. Because we don't have the Pipeline/ShaderObject state object we have to basically reimplement
// various state.
// The core issue with this is we use it everything for Pipeline, but for ShaderObject, we only use at vkCreateShadersEXT()
// Also for Pipeline, we create the state object prior to validating the vkCreatePipeline call
struct ShaderStageState {
    // We use this over a spirv::Module because there are times we need to create empty objects
    std::shared_ptr<const vvl::ShaderModule> module_state;
    std::shared_ptr<const spirv::Module> spirv_state;
    const vku::safe_VkPipelineShaderStageCreateInfo *pipeline_create_info;
    const vku::safe_VkShaderCreateInfoEXT *shader_object_create_info;
    // Can be null because GPL
    const vvl::DescriptorSetLayoutList *descriptor_set_layouts;
    // Will be NULL if coming from ShaderObject or Descriptor Heaps
    VkPipelineLayout pipeline_layout;
    // If null, means it is an empty object, no SPIR-V backing it
    std::shared_ptr<const spirv::EntryPoint> entrypoint;
    // VK_EXT_descriptor_heap
    const struct Heap {
        // Because this can come from a different struct depending on the Pipeline type, have it passed on creation
        bool descriptor_heap_mode = false;
        // If the heap is used by any variable (mapping or untyped)
        bool uses_resource_heap = false;
        bool uses_sampler_heap = false;

        bool uses_push_data = false;
    } heap;

    ShaderStageState(const vku::safe_VkPipelineShaderStageCreateInfo* pipeline_create_info,
                     const vku::safe_VkShaderCreateInfoEXT* shader_object_create_info,
                     const vvl::DescriptorSetLayoutList* descriptor_set_layouts,
                     std::shared_ptr<const vvl::ShaderModule> module_state, std::shared_ptr<const spirv::Module> spirv_state,
                     const VkPipelineLayout pipeline_layout, bool descriptor_heap_mode);

    bool HasPipeline() const { return pipeline_create_info != nullptr; }
    const char *GetPName() const;
    VkShaderStageFlagBits GetStage() const;
    vku::safe_VkSpecializationInfo *GetSpecializationInfo() const;
    const void* GetPNext() const;

    // Handles things like VK_EXT_shader_module_identifier and VK_SHADER_CODE_TYPE_SPIRV_EXT
    bool HasSpirv() const { return spirv_state.get() != nullptr && entrypoint.get() != nullptr; }

  private:
    Heap GetHeapInfo(bool descriptor_heap_mode);
};

namespace spirv {
struct ResourceInterfaceVariable;
}  // namespace spirv

struct DescriptorRequirement {
    uint64_t revalidate_hash;
    const spirv::ResourceInterfaceVariable *variable;
    DescriptorRequirement() : revalidate_hash(0), variable(nullptr) {}
};

inline bool operator==(const DescriptorRequirement &a, const DescriptorRequirement &b) noexcept {
    return a.revalidate_hash == b.revalidate_hash;
}

inline bool operator<(const DescriptorRequirement &a, const DescriptorRequirement &b) noexcept {
    return a.revalidate_hash < b.revalidate_hash;
}

// < binding index (of descriptor set) : meta data >
using BindingVariableMap = std::unordered_multimap<uint32_t, DescriptorRequirement>;

// Capture which slots (set#->bindings) are actually used by the shaders of this pipeline/shaderObject.
// This is same as "statically used" in vkspec.html#shaders-staticuse
using ActiveSlotMap = vvl::unordered_map<uint32_t, BindingVariableMap>;

void GetActiveSlots(ActiveSlotMap &active_slots, const std::shared_ptr<const spirv::EntryPoint> &entrypoint);
ActiveSlotMap GetActiveSlots(const std::vector<ShaderStageState> &stage_states);
ActiveSlotMap GetActiveSlots(const std::shared_ptr<const spirv::EntryPoint> &entrypoint);

uint32_t GetMaxActiveSlot(const ActiveSlotMap &active_slots);
