/* Copyright (c) 2015-2026 The Khronos Group Inc.
 * Copyright (c) 2015-2026 Valve Corporation
 * Copyright (c) 2015-2026 LunarG, Inc.
 * Copyright (C) 2015-2025 Google Inc.
 * Modifications Copyright (C) 2020 Advanced Micro Devices, Inc. All rights reserved.
 * Modifications Copyright (C) 2022 RasterGrid Kft.
 *
 * 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 "state_tracker/state_object.h"
#include "state_tracker/submission_reference.h"
#include "containers/span.h"
#include <vulkan/utility/vk_safe_struct.hpp>
#include <optional>
#include <deque>

namespace vvl {
class DeviceState;
class Fence;
class Semaphore;
class Surface;
class Swapchain;
class SwapchainSubState;
class Image;
}  // namespace vvl

struct GpuQueue {
    VkPhysicalDevice gpu;
    uint32_t queue_family_index;
};

inline bool operator==(GpuQueue const &lhs, GpuQueue const &rhs) {
    return (lhs.gpu == rhs.gpu && lhs.queue_family_index == rhs.queue_family_index);
}

namespace std {
template <>
struct hash<GpuQueue> {
    size_t operator()(GpuQueue gq) const throw() {
        return hash<uint64_t>()((uint64_t)(gq.gpu)) ^ hash<uint32_t>()(gq.queue_family_index);
    }
};
}  // namespace std

namespace vvl {
class Swapchain;

// Tracks the status of the fence or semaphore specified by acquire operation.
// The reason we can't simply check the fence or semaphore's own status is that the sync
// primitive can be reused, and it might report "not ready" due to other usage even if we
// have already waited on it after acquire.
enum class AcquireSyncStatus { NotSpecified, Signaled, WasWaitedOn };

struct SwapchainImage {
    vvl::Image* image_state = nullptr;

    // Acquire state
    bool acquired = false;
    std::shared_ptr<vvl::Semaphore> acquire_semaphore;
    std::shared_ptr<vvl::Fence> acquire_fence;
    AcquireSyncStatus acquire_semaphore_status = AcquireSyncStatus::NotSpecified;
    AcquireSyncStatus acquire_fence_status = AcquireSyncStatus::NotSpecified;
    void ResetAcquireState();

    // Queue location (seq) for present operation that presented this image.
    // When this image is reacquired, the acquire fence can synchronize with this location.
    std::optional<SubmissionReference> present_submission_ref;

    // Wait semaphores from the presentation request
    small_vector<std::shared_ptr<vvl::Semaphore>, 1> present_wait_semaphores;
    void ResetPresentWaitSemaphores();
};

// State for VkSwapchainKHR objects.
// Parent -> child relationships in the object usage tree:
//    vvl::Swapchain [N] -> [1] vvl::Surface
//    However, only 1 swapchain for each surface can be !retired.
class Swapchain : public RefcountedStateObject, public SubStateManager<SwapchainSubState> {
  public:
    const vku::safe_VkSwapchainCreateInfoKHR safe_create_info;
    const VkSwapchainCreateInfoKHR& create_info;
    // VkImageUsageFlags2CreateInfoKHR can be used instead of the VkSwapchainCreateInfoKHR::imageUsage
    const VkImageUsageFlags2KHR image_usage;

    std::vector<VkPresentModeKHR> present_modes;
    std::vector<SwapchainImage> images;
    bool retired = false;
    bool exclusive_full_screen_access;
    const bool shared_presentable;
    uint64_t max_present_id = 0;
    const vku::safe_VkImageCreateInfo image_create_info;

    std::shared_ptr<vvl::Surface> surface;
    DeviceState &dev_data;
    uint32_t acquired_images = 0;

    // Image acquire history ring buffer (most recently acquired image indices)
    static constexpr uint32_t acquire_history_ring_buffer_size = 16;
    std::array<uint32_t, acquire_history_ring_buffer_size> acquire_history_ring_buffer;
    uint32_t acquire_request_count = 0;
    uint32_t GetAcquireHistoryLength() const;
    uint32_t GetAcquiredImageIndexFromHistory(uint32_t acquire_history_index) const;

    // Present ids and associated information ring buffer.
    // It is not reasonable to keep all present id, instead, track the fixed
    // number of the most recent present ids. This works well in practice
    // since synchronization often involves only the last few frames.
    // The storage for the ring buffer is allocated on the first usage.
    static constexpr uint32_t present_id_ring_buffer_size = 24;
    struct PresentIdInfo {
        uint64_t present_id;
        SubmissionReference present_submission_ref;
    };
    std::vector<PresentIdInfo> present_id_ring_buffer;
    uint32_t present_id_request_count = 0;
    uint32_t GetPresentIdInfoCount() const;
    PresentIdInfo GetPresentIdInfo(uint32_t present_id_info_index);

    // Old swapchain state:
    // The new swapchain is a swapchain for which *this* swapchain is the oldSwapchain.
    // We still can present (pre-acquired) images from the old swapchain.
    // The new swapchain is responsible for tracking in-use status of present wait semaphores from
    // the last few presentations (if any) of the old swapchain.
    std::shared_ptr<Swapchain> new_swapchain;

    // New swapchain state:
    // Present wait semaphores from the the old swapchain presentations.
    std::vector<std::shared_ptr<vvl::Semaphore>> old_swapchain_present_wait_semaphores;
    std::weak_ptr<Swapchain> old_swapchain;

    // Number of bits set in VkPresentTimingInfoEXT::presentStageQueries for each present that hasn't been completed
    struct PresentTimingInfo {
        uint32_t num_present_stage_queries;
        SubmissionReference present_submission_ref;
    };
    std::deque<std::pair<uint64_t, PresentTimingInfo>> present_timing_stage_queries;
    // Present timing queue size for the swapchain set by vkSetSwapchainPresentTimingQueueSizeEXT
    uint32_t present_timing_queue_size = 0;
    // Present timing id to domain map
    std::unordered_map<uint64_t, VkTimeDomainKHR> time_domains;
    bool present_at_absolute_time_supported;
    bool present_at_relative_time_supported;

    Swapchain(DeviceState &dev_data, const VkSwapchainCreateInfoKHR *pCreateInfo, VkSwapchainKHR handle);

    ~Swapchain() {
        if (!Destroyed()) {
            Destroy();
        }
    }

    const VulkanTypedHandle* InUse() const override;
    VkSwapchainKHR VkHandle() const { return handle_.Cast<VkSwapchainKHR>(); }

    void PresentImage(uint32_t image_index, uint64_t present_id, const SubmissionReference &present_submission_ref,
                      vvl::span<std::shared_ptr<vvl::Semaphore>> present_wait_semaphores);

    void ReleaseImage(uint32_t image_index);

    void AcquireImage(uint32_t image_index, const std::shared_ptr<vvl::Semaphore> &semaphore_state,
                      const std::shared_ptr<vvl::Fence> &fence_state);

    void Destroy() override;

    SwapchainImage GetSwapChainImage(uint32_t index) const;

    std::shared_ptr<const vvl::Image> GetSwapChainImageShared(uint32_t index) const;

    std::shared_ptr<const Swapchain> shared_from_this() const { return SharedFromThisImpl(this); }
    std::shared_ptr<Swapchain> shared_from_this() { return SharedFromThisImpl(this); }

  protected:
    void NotifyInvalidate(const StateObject::NodeList &invalid_nodes, bool unlink) override;
};

class SwapchainSubState {
  public:
    explicit SwapchainSubState(Swapchain &sc) : base(sc) {}
    SwapchainSubState(const SwapchainSubState &) = delete;
    SwapchainSubState &operator=(const SwapchainSubState &) = delete;
    virtual ~SwapchainSubState() {}
    virtual void Destroy() {}

    Swapchain &base;
};

enum class SurfaceType {
    Unknown,

    Android,
    DisplayPlane,
    Headless,
    ImagePipe_FUSCHIA,
    IOS_MVK,
    MacOS_MVK,
    Metal,
    Screen_QNX,
    Wayland,
    Win32,
    Xcb,
    Xlib,
};


// Parent -> child relationships in the object usage tree:
//    vvl::Surface -> nothing
class Surface : public StateObject {
  public:
    Surface(VkSurfaceKHR handle, SurfaceType type) : StateObject(handle, kVulkanObjectTypeSurfaceKHR) {
        type_ = type;
    }

    ~Surface() {
        if (!Destroyed()) {
            Destroy();
        }
    }

    VkSurfaceKHR VkHandle() const { return handle_.Cast<VkSurfaceKHR>(); }

    void Destroy() override;

    void RemoveParent(StateObject *parent_node) override;

    void SetQueueSupport(VkPhysicalDevice phys_dev, uint32_t qfi, bool supported);
    bool GetQueueSupport(VkPhysicalDevice phys_dev, uint32_t qfi) const;

    void SetPresentModes(VkPhysicalDevice phys_dev, vvl::span<const VkPresentModeKHR> modes);
    std::vector<VkPresentModeKHR> GetPresentModes(VkPhysicalDevice phys_dev) const;

    void SetFormats(VkPhysicalDevice phys_dev, std::vector<vku::safe_VkSurfaceFormat2KHR> &&fmts);
    vvl::span<const vku::safe_VkSurfaceFormat2KHR> GetFormats(bool get_surface_capabilities2, VkPhysicalDevice phys_dev,
                                                              const void *surface_info2_pnext) const;

    // Cache capabilities that do not depend on the present mode
    void UpdateCapabilitiesCache(VkPhysicalDevice phys_dev, const VkSurfaceCapabilitiesKHR &surface_caps);
    // Cache per present mode capabilities
    void UpdateCapabilitiesCache(VkPhysicalDevice phys_dev, const VkSurfaceCapabilities2KHR &surface_caps,
                                 VkPresentModeKHR present_mode);

    bool IsLastCapabilityQueryUsedPresentMode(VkPhysicalDevice phys_dev) const;
    VkSurfaceCapabilitiesKHR GetSurfaceCapabilities(VkPhysicalDevice phys_dev, const void *surface_info_pnext) const;
    VkSurfaceCapabilitiesKHR GetPresentModeSurfaceCapabilities(VkPhysicalDevice phys_dev, VkPresentModeKHR present_mode) const;
    VkSurfacePresentScalingCapabilitiesKHR GetPresentModeScalingCapabilities(VkPhysicalDevice phys_dev,
                                                                             VkPresentModeKHR present_mode) const;
    std::vector<VkPresentModeKHR> GetCompatibleModes(VkPhysicalDevice phys_dev, VkPresentModeKHR present_mode) const;

    SurfaceType GetSurfaceType() { return type_; }

    vvl::Swapchain *swapchain{nullptr};

  private:
    // Contains per present mode capabilities
    struct PresentModeInfo {
        VkPresentModeKHR present_mode;
        VkSurfaceCapabilitiesKHR surface_capabilities;
        std::optional<VkSurfacePresentScalingCapabilitiesKHR> scaling_capabilities;
        std::optional<std::vector<VkPresentModeKHR>> compatible_present_modes;
    };
    // Cached information per physical device. Optional indicates if element is in the cache.
    //
    // NOTE: One of the reasons to cache surface caps is to prevent a false-positive
    // when the surface change happens (e.g. resize) after the surface caps are queried
    // and before the swapchain is created. The assumption is that with the current API,
    // the app can't do better than this (no atomicity between query and swapchain creation).
    // The caching ensures that validation sees the same surface state as the application.
    //
    // The priority is to avoid false-positives for correctly written application.
    // When the application behaves incorrectly (e.g. forgets to query surface caps after
    // it processed the resize event), then the caching can hide a problem, since validation
    // will think that application respects surface caps values.
    struct PhysDevCache {
        std::optional<std::vector<VkPresentModeKHR>> present_modes;
        std::optional<VkSurfaceCapabilitiesKHR> capabilities;
        std::vector<PresentModeInfo> present_mode_infos;
        bool last_capability_query_used_present_mode = false;

        const PresentModeInfo *GetPresentModeInfo(VkPresentModeKHR present_mode) const;
    };
    const PhysDevCache *GetPhysDevCache(VkPhysicalDevice phys_dev) const;

  private:
    std::unique_lock<std::mutex> Lock() const { return std::unique_lock<std::mutex>(lock_); }
    // TODO: make mutex shared, so multiple Validate can read simultaneously. Remove remaining mutations in Validate first
    mutable std::mutex lock_;
    mutable vvl::unordered_map<GpuQueue, bool> gpu_queue_support_;
    mutable vvl::unordered_map<VkPhysicalDevice, std::vector<vku::safe_VkSurfaceFormat2KHR>> formats_;

    vvl::unordered_map<VkPhysicalDevice, PhysDevCache> cache_;
    SurfaceType type_;
};

}  // namespace vvl
