// Copyright (c) 2023 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.
//
#include <OgnScaleAOVTextureDatabase.h>
#include "OgnTestPostRenderGraphNodeCuda.h"

#include <carb/cuda/CudaRuntime.h>
#include <omni/graph/image/unstable/ComputeParamsBuilder.h>

#include <string>

namespace omni {
namespace graph {
namespace rtxtest {

using namespace omni::graph::image::unstable;

namespace
{
/// A wrapper around a texture object
class ScopedCudaTextureObject final
{
    cudaTextureObject_t _texObj = 0;

public:
    ScopedCudaTextureObject(cudaMipmappedArray_t mmarr, int mipLevel = 0)
    {
        if (!mmarr)
            return;
        cudaArray_t levelArray = 0;
        CUDA_CHECK(cudaGetMipmappedArrayLevel(&levelArray, mmarr, mipLevel));
        if (!levelArray)
            return;
        struct cudaResourceDesc resDesc;
        memset(&resDesc, 0, sizeof(resDesc));
        resDesc.resType = cudaResourceTypeArray;
        resDesc.res.array.array = levelArray;
        struct cudaTextureDesc texDesc;
        memset(&texDesc, 0, sizeof(texDesc));
        texDesc.addressMode[0] = cudaAddressModeClamp;
        texDesc.addressMode[1] = cudaAddressModeClamp;
        texDesc.filterMode = cudaFilterModePoint;
        texDesc.readMode = cudaReadModeElementType;
        texDesc.normalizedCoords = 1;
        CUDA_CHECK(cudaCreateTextureObject(&_texObj, &resDesc, &texDesc, nullptr));
    }

    ~ScopedCudaTextureObject()
    {
        if (_texObj)
        {
            CUDA_CHECK(cudaDestroyTextureObject(_texObj));
        }
    }

    operator cudaTextureObject_t&()
    {
        return _texObj;
    }
};

/// A wrapper around a suface object
class ScopedCudaSurfaceObject final
{
    cudaSurfaceObject_t _surfObj = 0;

public:
    ScopedCudaSurfaceObject(cudaMipmappedArray_t mmarr, int mipLevel = 0)
    {
        if (!mmarr)
            return;
        cudaArray_t levelArray = 0;
        CUDA_CHECK(cudaGetMipmappedArrayLevel(&levelArray, mmarr, 0));
        if (!levelArray)
            return;
        struct cudaResourceDesc surfResDesc;
        memset(&surfResDesc, 0, sizeof(surfResDesc));
        surfResDesc.resType = cudaResourceTypeArray;
        surfResDesc.res.array.array = levelArray;
        CUDA_CHECK(cudaCreateSurfaceObject(&_surfObj, &surfResDesc));
    }

    ~ScopedCudaSurfaceObject()
    {
        if (_surfObj)
        {
            CUDA_CHECK(cudaDestroySurfaceObject(_surfObj));
        }
    }

    operator cudaSurfaceObject_t&()
    {
        return _surfObj;
    }
};
}

class OgnScaleAOVTexture
{
public:
    static bool compute(OgnScaleAOVTextureDatabase& db)
    {
        // The gpu and rp inputs are passed through to the output.
        // They are only present here so that the nodes GpuInteropGpuToCpuCopy and GpuInteropCpuToDisk can be connected
        // after the current test node, to establish the correct topological order in the graph.
        // An alternate way to do this is to use a dummy execution port for this purpose (this is what the SDG nodes do).
        // In the future, the data protocol for post render graph nodes should standardize the inputs of the nodes and
        // prevent the need for passing raw pointers around.
        db.outputs.gpu() = db.inputs.gpu();
        db.outputs.rp() = db.inputs.rp();

        auto gpu = reinterpret_cast<GpuFoundationsInterfaces*>(db.inputs.gpu());
        auto rp = reinterpret_cast<omni::usd::hydra::HydraRenderProduct*>(db.inputs.rp());

        // This node is called in tests and has to fail hard
        CARB_ASSERT(gpu);
        CARB_ASSERT(rp);

        auto result = ComputeParamsBuilder<std::string>{ gpu, rp, db }
            .addValue("multiplier", db.inputs.multiplier())
            .addInputTexture("inputAOV", db.inputs.inputAOV(),
                [&db](cudaMipmappedArray_t cudaPtr, carb::graphics::TextureDesc const* desc, ComputeParams<std::string>& params)
                {
                    // the desc pointer is guaranteed to be valid
                    params.add("width", desc->width);
                    params.add("height", desc->height);
                })
            .addOutputTexture("outputAOV", db.inputs.outputAOV(), db.inputs.width(), db.inputs.height(),
                carb::graphics::Format::eRGBA8_UNORM, "TestTexture",
                [&db](cudaMipmappedArray_t ptr)
                {
                    db.outputs.outputAOVPtr() = reinterpret_cast<uint64_t>(ptr);
                })
            .scheduleCudaTask("TestCudaTask",
                [](ComputeParams<std::string>* data, cudaStream_t stream)
                {
                    auto multiplier = data->get<float>("multiplier");
                    auto inputAOV = data->get<cudaMipmappedArray_t>("inputAOV");
                    auto outputAOV = data->get<cudaMipmappedArray_t>("outputAOV");
                    auto width = data->get<uint32_t>("width");
                    auto height = data->get<uint32_t>("height");

                    CARB_ASSERT(multiplier != 0.0f);
                    CARB_ASSERT(width != 0.0f);
                    CARB_ASSERT(height != 0.0f);

                    ScopedCudaTextureObject inputTex(inputAOV);
                    ScopedCudaSurfaceObject outputTex(outputAOV);

                    cudaScaleImage(inputTex, multiplier, width, height, outputTex, stream);
                },
                [](ComputeParams<std::string> const& params) -> bool
                {
                    return params.hasKey("inputAOV") && params.get<cudaMipmappedArray_t>("inputAOV") != nullptr &&
                        params.hasKey("outputAOV") && params.get<cudaMipmappedArray_t>("outputAOV") != nullptr;
                });
        db.outputs.scheduleResult() = result;
        return result;
    }
};

REGISTER_OGN_NODE()

} // test
} // namespace graph
} // namespace omni
