// Copyright (c) 2020, 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 <carb/Types.h>

#include <map>
#include <memory>
#include <vector>

using namespace carb;

namespace carb
{
namespace graphics
{
struct Context;
struct Graphics;
struct Device;
struct CommandList;
struct CommandQueue;
struct Fence;
struct Shader;
}
}

namespace omni
{
namespace gpucompute
{
enum class MemcpyKind
{
    hostToHost = 0,
    hostToDevice = 1,
    deviceToHost = 2,
    deviceToDevice = 3
};

struct Context;
struct Shader;

struct Parameter
{
    const char* name;
    bool isConstantBuffer;
};

// API agnostic representation of
// static_cast<char*>(buffer) + byteOffset
// which points to data of size elemSize
struct GpuPointer
{
    void* buffer;
    size_t byteOffset;
    size_t elemSize;
};

struct CpuBuffer
{
    void* data;
    size_t count;
};

struct Args
{
    // TODO: move CPU backend into its own plugin

    // GPU
    std::vector<GpuPointer> gpuBuffers;
    std::vector<size_t> gpuArgToBufferCount;

    // CPU
    std::vector<CpuBuffer> cpuArgs;
};

// A CUDA-style interface for D3D, consisting of GPU malloc, free, memcpy and
// kernel dispatch
struct GpuCompute
{
    CARB_PLUGIN_INTERFACE("omni::gpucompute::GpuCompute", 0, 1)

    // CPU memory allocation
    void(CARB_ABI* hostAlloc)(Context& context, void** ptr, size_t byteCount);
    void(CARB_ABI* freeHost)(Context& context, void* ptr);

    // GPU memory allocation
    void*(CARB_ABI* malloc)(Context& context, size_t byteCount, size_t elemSize);
    void(CARB_ABI* free)(Context& context, void* ptr);

    // GPU async memory allocation (uses stream 0 for CUDA)
    void*(CARB_ABI* mallocAsync)(Context& context, size_t byteCount, size_t elemSize);
    void(CARB_ABI* freeAsync)(Context& context, void* ptr);

    void(CARB_ABI* memcpy)(Context& context, void* dst, const void* src, size_t byteCount, MemcpyKind kind);

    void(CARB_ABI* dispatch)(Context& context, Shader& shader, Args& args, carb::Uint3 gridDim);

    Context&(CARB_ABI* createContext)();
    Context&(CARB_ABI* createContextD3dVk)(carb::graphics::Graphics* graphics,
                                           carb::graphics::Device* device,
                                           carb::graphics::CommandList* commandList,
                                           carb::graphics::CommandQueue* commandQueue,
                                           carb::graphics::Fence* fence);
    void(CARB_ABI* destroyContext)(Context& context);

    uint32_t(CARB_ABI* peekAtLastError)(Context& context);
    uint32_t(CARB_ABI* getLastError)(Context& context);

    void(CARB_ABI* memcpyAsync)(Context& context, void* dst, const void* src, size_t byteCount, MemcpyKind kind);
};

enum class Target
{
    CPU,
    GPU
};

struct ComputeCompiler
{
    CARB_PLUGIN_INTERFACE("omni::gpucompute::ComputeCompiler", 0, 1)

    omni::gpucompute::Shader*(CARB_ABI* compile)(carb::graphics::Device* device, Target target, const char* codeString);
    void(CARB_ABI* getParameters)(Parameter*& parameters, size_t& count, Shader& shader);
    void(CARB_ABI* destroyShader)(Shader& shader);
};

}
}
