// Copyright (c) 2019-2021, NVIDIA CORPORATION. All rights reserved.
//
// NVIDIA CORPORATION and its licensors retain all intellectual property
// and proprietary rights in and to this software, related documentation
// and any modifications thereto.  Any use, reproduction, disclosure or
// distribution of this software and related documentation without an express
// license agreement from NVIDIA CORPORATION is strictly prohibited.
//
#pragma once

#include "IGpuFoundation.h"

#include <carb/IObject.h>
#include <carb/Interface.h>
#include <carb/InterfaceUtils.h>
#include <carb/tasking/TaskingTypes.h>

#include <omni/kit/IAppWindow.h>

namespace carb
{
namespace events
{
class IEventStream;
}
namespace graphics
{
enum class DescriptorType : int32_t;
struct Device;
class DeviceGroup;
struct CommandList;
struct CommandQueue;
struct DescriptorPool;
struct ResourceBindingSignature;
struct TextureDesc;
// TextureGpuData
struct Texture;
struct Descriptor;
struct DescriptorSet;
}
namespace graphicsmux
{
class CommandList;
}
namespace renderer
{
struct TextureHandle;
}
}

namespace gpu
{
enum class GfResult : int32_t;
}

namespace rtx
{
namespace resourcemanager
{
class RpResource;
typedef uint32_t SyncScopeId;
}
}

namespace omni
{
namespace kit
{
namespace renderer
{

typedef int64_t CommandListHandle;

/**
 * Defines a simple interface for texture GPU resources
 */
class TextureGpuData : public carb::IObject
{
public:
    virtual carb::graphics::Texture* getGraphicsTexture() const = 0;
    virtual carb::graphics::Descriptor* getGraphicsDescriptor() const = 0;

    virtual carb::graphics::DescriptorSet* getGraphicsDescriptorSet() const = 0;

    virtual rtx::resourcemanager::RpResource* getManagedResource() const = 0;
    virtual uint32_t getBindlessDescriptorIndex() const = 0;
};

struct TextureGpuReference
{
    union {
        uint32_t gpuIndex;
        void* ptr;
    };
};

/**
 * @brief The data for all the viewports in GUI. We use it when present thread
 * is enabled to render viewport with more FPS that GUI.
 */
struct IRendererHydraViewportData
{
    using Clock = ::std::chrono::high_resolution_clock;
    using Duration = Clock::duration;
    using TimePoint = Clock::time_point;

    uint32_t viewportId;
    size_t subFrameCount;
    uint32_t subFrames[4];
    size_t currentSubFrame;
    TimePoint publishTime;
    Duration frameAverageTime;
};

/**
 * @brief Options for texture creation.
 */
struct IRendererTextureOptions
{
    uint32_t gpuDeviceMask = 0;        ///< Device mask to create texture on (0 for UI device only)
    uint32_t textureUsageFlags = 0;    ///< carb::graphics::TextureUsageFlags
    uint32_t resourceUsageFlags = 0;   ///< rtx::ResourceUsageFlags
    uint32_t unusedExtPadding = 0;     ///< Unused padding and expansion
};

/**
 * Defines an interface for the small version of editor that provides simple GUI.
 */
class IRenderer
{
public:
    CARB_PLUGIN_INTERFACE("omni::kit::renderer::IRenderer", 1, 9);

    virtual void startup() = 0;
    virtual void shutdown() = 0;

    /**
     * Render events stream. Render events and pushed and popped every frame.
     */
    virtual carb::events::IEventStream* getPreBeginFrameEventStream(IAppWindow* appWindow) = 0;
    virtual carb::events::IEventStream* getPreBeginRenderPassEventStream(IAppWindow* appWindow) = 0;
    virtual carb::events::IEventStream* getRenderFrameEventStream(IAppWindow* appWindow) = 0;
    virtual carb::events::IEventStream* getPostEndRenderPassEventStream(IAppWindow* appWindow) = 0;
    virtual carb::events::IEventStream* getPostEndRenderFrameEventStream(IAppWindow* appWindow) = 0;

    virtual carb::graphics::Device* getGraphicsDevice(IAppWindow* appWindow) = 0;
    virtual carb::graphics::DeviceGroup* getGraphicsDeviceGroup(IAppWindow* appWindow) = 0;
    virtual uint32_t getGraphicsDeviceIndex(IAppWindow* appWindow) = 0;
    virtual carb::Format getGraphicsSwapchainFormat(IAppWindow* appWindow) = 0;
    virtual carb::graphics::CommandQueue* getGraphicsCommandQueue(IAppWindow* appWindow) = 0;

    typedef void (*UploadCompleteCbDeprecated)(void* arg,
                                               void* textureHandle,
                                               carb::graphics::Descriptor* textureDescriptor,
                                               size_t width,
                                               size_t height,
                                               void* userData);
    struct UploadTextureDescDeprecated
    {
        const char* assetUri;
        UploadCompleteCbDeprecated callback;
        void* arg;
    };

    // TODO: remove these functions as they are merely wrappers around GPU data creation with integrated assets
    // management
    virtual void uploadTextureDeprecated(const UploadTextureDescDeprecated& desc, void* userData = nullptr) = 0;
    virtual void destroyTextureDeprecated(void* texture) = 0;
    virtual TextureGpuData* getTextureDataFromHandle(void* textureHandle) = 0;

    static constexpr uint32_t kMaxDescriptorSetCount = 8192;
    virtual carb::graphics::DescriptorPool* getDescriptorPool() = 0;

    /**
     * These functions are needed because carb::graphics doesn't have descriptor set deallocations, we have very
     * finite set of allocations - we have to track them to make sure we can emit a warning if we're reaching the limit.
     */
    virtual void notifyDescriptorPoolAllocation(carb::graphics::DescriptorType descriptorType) = 0;
    virtual int getFreeDescriptorPoolSlotCount(carb::graphics::DescriptorType descriptorType) = 0;

    /**
     * Functions that allocate id to use in shaders, given the created/imported GPU data.
     */
    virtual TextureGpuReference allocateReferenceForTexture(
        carb::graphics::ResourceBindingSignature* resourceBindingSignature) = 0;
    virtual TextureGpuReference allocateReferenceForDrawable(
        carb::graphics::ResourceBindingSignature* resourceBindingSignature) = 0;

    virtual TextureGpuReference updateReferenceFromTextureData(TextureGpuReference textureGpuReference,
                                                               TextureGpuData* textureGpuData) = 0;
    virtual TextureGpuReference updateReferenceFromDrawableData(TextureGpuReference* textureGpuReference,
                                                                TextureGpuData* drawableGpuData) = 0;

    /**
     * Functions to manipulate GPU data that was created outside of the renderer.
     */
    virtual TextureGpuData* createExternalTextureGpuData() = 0;
    virtual void updateExternalTextureGpuData(
        TextureGpuData* externalTextureGpuData,
        carb::graphics::Descriptor* externalDescriptor,
        uint32_t bindlessDescriptorIndex = CARB_UINT32_MAX) = 0;

    /**
     * Functions to create managed GPU data.
     */
    virtual TextureGpuData* createGpuResourcesForTexture(carb::Format format,
                                                         size_t mipMapCount,
                                                         size_t* mipWidths,
                                                         size_t* mipHeights,
                                                         const uint8_t** mipDatas,
                                                         size_t* mipStrides) = 0;
    virtual void updateGpuResourcesForTexture(TextureGpuData* textureGpuData,
                                              carb::Format format,
                                              size_t mipMapCount,
                                              size_t* mipWidths,
                                              size_t* mipHeights,
                                              const uint8_t** mipDatas,
                                              size_t* mipStrides) = 0;

    virtual void resetGpuResourcesForTexture(TextureGpuData* textureGpuData) = 0;
    virtual void destroyGpuResourcesForTexture(TextureGpuData* textureGpuData) = 0;

    /**
     * Functions to extract command lists from handles.
     */
    virtual carb::graphics::CommandList* getGraphicsCommandList(CommandListHandle commandListHandle) = 0;
    virtual carb::graphicsmux::CommandList* getGraphicsMuxCommandList(CommandListHandle commandListHandle) = 0;

    virtual uint32_t getCurrentBackBufferIndex(IAppWindow* appWindow) = 0;
    virtual uint32_t getBackBufferCount(IAppWindow* appWindow) = 0;
    virtual size_t getFrameCount(IAppWindow* appWindow) = 0;

    /**
     * Framebuffer information.
     */
    virtual carb::graphics::Texture* getFramebufferTexture(IAppWindow* appWindow) = 0;
    virtual uint32_t getFramebufferWidth(IAppWindow* appWindow) = 0;
    virtual uint32_t getFramebufferHeight(IAppWindow* appWindow) = 0;
    virtual carb::Format getFramebufferFormat(IAppWindow* appWindow) = 0;
    virtual CommandListHandle getFrameCommandList(IAppWindow* appWindow) = 0;

    virtual TextureGpuReference getInvalidTextureReference() = 0;

    static inline IGpuFoundation* getGpuFoundation();

    virtual bool attachAppWindow(IAppWindow* appWindow) = 0;
    virtual void detachAppWindow(IAppWindow* appWindow) = 0;
    virtual bool isAppWindowAttached(IAppWindow* appWindow) = 0;

    virtual void waitIdle(IAppWindow* appWindow) = 0;
    virtual void forceRenderFrame(float dt) = 0;

    virtual rtx::resourcemanager::SyncScopeId getSyncScopeId() = 0;

    /**
     * Get a reference to the current ResourceManager and it's Context.
     * IAppWindow* can currently be null.
     */
    struct ResourceManagerState
    {
        rtx::resourcemanager::ResourceManager* manager;
        rtx::resourcemanager::Context* context;
    };
    virtual ResourceManagerState getResourceManagerState(IAppWindow* appWindow) = 0;

    virtual void setClearColor(IAppWindow* appWindow, const carb::Float4& clearColor) = 0;
    virtual carb::Float4 getClearColor(IAppWindow* appWindow) = 0;


    virtual bool setRenderQueue(IAppWindow* appWindow, size_t renderQueueIdx) = 0;
    virtual size_t getRenderQueue(IAppWindow* appWindow) = 0;

    /**
     * Functions that deallocate id to use in shaders, given the created/imported GPU data.
     */
    virtual void deallocateReferenceForTexture(
        carb::graphics::ResourceBindingSignature* resourceBindingSignature, const TextureGpuReference& textureGpuReference) = 0;
    virtual void deallocateReferenceForDrawable(
        carb::graphics::ResourceBindingSignature* resourceBindingSignature, const TextureGpuReference& textureGpuReference) = 0;

    /**
     * Present thread streams.
     *
     * Called from the separate thread (present thread) to render the saved UI
     * with bigger framerate than it's generated.
     */
    virtual carb::events::IEventStream* getPresentRenderFrameEventStream(IAppWindow* appWindow) = 0;

    /**
     * @brief Should be called when the new viewport frame is ready
     *
     * @param dt time to render this frame in ns
     * @param viewportId the id of the viewport that should be replaced to the
     *                   given textureId
     * @param subFrameCount the number of subframes
     * @param subFrames the viewport subframes
     */
    virtual void notifyRenderingIsReady(float dtNs,
                                        uint32_t viewportId,
                                        size_t subFrameCount,
                                        const uint32_t* subFrames) = 0;

    /**
     * Functions to create managed GPU data with optional multiGpu and GPU source support.
     *
     * @param multiGpu data will be uploaded or copied to all GPUs in device group
     * @param fromGpu data is a gpu memory address, not host memory
     */
    virtual TextureGpuData* createGpuResourcesForTexture(carb::Format format,
                                                         size_t mipMapCount,
                                                         size_t* mipWidths,
                                                         size_t* mipHeights,
                                                         const uint8_t** mipDatas,
                                                         size_t* mipStrides,
                                                         bool multiGpu,
                                                         bool cudaBytes) = 0;
    virtual void updateGpuResourcesForTexture(TextureGpuData* textureGpuData,
                                              carb::Format format,
                                              size_t mipMapCount,
                                              size_t* mipWidths,
                                              size_t* mipHeights,
                                              const uint8_t** mipDatas,
                                              size_t* mipStrides,
                                              bool cudaBytes) = 0;

    /**
     * Add a callback to be run in the IRenderer / Graphics initialization
     *
     * @param eventHandler An IEventListener to be called during IRenderer/Graphics setup.
     * This handler is currently invoked with a null carb::event::IEvent pointer.
     *
     */
    virtual void addToInitPipeline(carb::events::IEventListenerPtr eventHandler) = 0;

    /**
     * Block current thread until renderer (& GPU Foundation) is fully setup
     *
     */
    virtual void waitForInit() = 0;

    /**
     * @brief Called immediately after the renderer presents the current frame
     * to the frame buffer. This function allows the developer to perform any
     * additional operations or processing on the frame after it has been
     * presented, such as sending the frame data to a CUDA memory frame for
     * further processing. Additionally, the callback can be used to trigger any
     * additional actions that need to be taken after the frame has been
     * presented, such as updating a status or sending a notification.
     */
    virtual carb::events::IEventStream* getPostPresentFrameBufferEventStream(IAppWindow* appWindow) = 0;

    /**
     * @brief Frozen window is skipped from the run loop and it's not updated.
     */
    virtual void freezeAppWindow(IAppWindow* appWindow, bool freeze) = 0;

    /*
     * @brief State of the present thread
     *
     * @return true The present thread is enabled
     * @return false The present thread is disabled
     */
    virtual bool isPresentThreadEnabled() const = 0;

    /**
     * @brief Create or update a GPU resource.
     *
     * @param format The format of the texture being created
     * @param mipMapResolutions An array of resolutions to use for mip-mapped data
     * @param mipMapDatas An array of bytes to upload, for each mip-level
     * @param mipMapStrides An array of strides on the bytes to upload, for each mip-level
     * @param mipMapCount The number of elements in the mipResolutions, mipDatas, and mipStrides arguments
     * @param multiGpu Data will be uploaded or copied to all GPUs in device group
     * @param createShared Whether to create the texture as shareable for later use in CUDA or OpenGL
     * @param cudaBytes Whether the pointers in mipDatas refer to CPU memory or CUDA memory
     * @param deviceMask The devices to create the texture for
     */
    virtual TextureGpuData* createGpuResourcesForTexture(carb::Format format,
                                                         const carb::Uint2* mipMapResolutions,
                                                         const uint8_t* const* mipMapDatas,
                                                         const size_t* mipMapStrides,
                                                         size_t mipMapCount,
                                                         const IRendererTextureOptions& options = {}) = 0;


    virtual void updateGpuResourcesForTexture(TextureGpuData* textureGpuData,
                                              carb::Format format,
                                              const carb::Uint2* mipMapResolutions,
                                              const uint8_t* const* mipMapDatas,
                                              const size_t* mipMapStrides,
                                              size_t mipMapCount,
                                              const IRendererTextureOptions& options = {}) = 0;

    /**
     * Ensures that the resource exists on the target display device, potentially performing an asynchronous GPU copy.
     * @return The pointer to the resource on the target display device. If the resource is already on the target
     * device, then this pointer will be the same as the input. This function will return nullptr on failure.
     * @param resource The resource to potentially copy to the target display device
     * @param timeout The timeout (in milliseconds) for the copy operation. The future will return eTimeout if the
     * timeout is exceeded.
     * @param outTextureHandle Optional: If not nullptr, the GPU handle for the resource on the target device.
     * @param outTextureDesc Optional: If not nullptr, the TextureDesc of the returned RpResource will be returned in
     * outTextureDesc.
     * @param outFuture Optional: If not nullptr, the receiver for a future that can be waited on for the completion of
     * the potential GPU copy.
     */
    virtual rtx::resourcemanager::RpResource* ensureRpResourceOnTargetDevice(
        rtx::resourcemanager::RpResource* resource,
        uint32_t timeout,
        carb::renderer::TextureHandle* outTextureHandle,
        const carb::graphics::TextureDesc** outTextureDesc,
        carb::tasking::SharedFuture<gpu::GfResult>* outFuture) = 0;

    /**
     * @brief Draw-frozen window is not skipped from the run loop and but the
     * draw list of this window is not updated.
     */
    virtual void drawFreezeAppWindow(IAppWindow* appWindow, bool freeze) = 0;
};

inline IGpuFoundation* IRenderer::getGpuFoundation()
{
    return carb::getCachedInterface<IGpuFoundation>();
}

}
}
}
