#include <algorithm>
#include <memory>
#include <mutex>
#include <utility>
#include <string>
#include <stdexcept>
#include <vector>
#include <stdio.h>
#include <assert.h>

#include "VapourSynth4.h"
#include "VSHelper4.h"

#ifndef _WIN32
#error "Unsupported platform"
#else
static std::vector<std::string> autoDllErrors;
#define CUDA_DLL L"nvcuda.dll","nvcuda.dll",autoDllErrors
#endif
#include "../ngx/cuda.h"

#include "nvvfx/include/nvCVStatus.h"
#include "nvvfx/include/nvCVImage.h"
#include "nvvfx/include/nvVideoEffects.h"

#define CK_VFX(x) do { \
    NvCV_Status r = (x); \
    if (r != NVCV_SUCCESS) { \
        fprintf(stderr, "failed VFX call %s: %x (%s)\n", #x, r, NvCV_GetErrorStringFromCode(r)); \
        abort(); \
    } \
} while (0)
#define CK_CUDA(x) do { \
    int r = (x); \
    if (r != CUDA_SUCCESS) { \
        fprintf(stderr, "failed cuda call %s: %d\n", #x, r); \
        abort(); \
    } \
} while (0)

struct VfxData {
    std::mutex lock;

    int num_streams;

    VSNode *node;
    VSVideoInfo vi;
    double scale;
    double strength;
    int output_depth;

    int in_width, in_height;

    NvVFX_Handle vfx;
    CUstream stream;
    CUdeviceptr state;

    float srcTransferFactor, dstTransferFactor;
    NvCVImage srcGpuImg;
    NvCVImage dstGpuImg;
    NvCVImage srcTmpImg, dstTmpImg;
    void *srcCpuBuf, *dstCpuBuf;

    typedef float T;
    uint64_t in_image_width() const   { return in_width; }
    uint64_t out_image_width() const  { return vi.width; }
    uint64_t in_image_height() const  { return in_height; }
    uint64_t out_image_height() const { return vi.height; }

    VfxData() : node(nullptr), vi(), scale(0), strength(0), vfx(nullptr), stream(nullptr), state(nullptr) {}
    ~VfxData() {
        if (vfx) NvVFX_DestroyEffect(vfx);
        if (stream) NvVFX_CudaStreamDestroy(stream);
        if (state) cuMemFree_v2(state);
        NvCVImage_Dealloc(&srcGpuImg);
        NvCVImage_Dealloc(&dstGpuImg);
        NvCVImage_Dealloc(&srcTmpImg);
        NvCVImage_Dealloc(&dstTmpImg);
        cuMemFreeHost(srcCpuBuf);
        cuMemFreeHost(dstCpuBuf);
    }
};

static const VSFrame *VS_CC vfxGetFrame(int n, int activationReason, void *instanceData, void **frameData, VSFrameContext *frameCtx, VSCore *core, const VSAPI *vsapi) {
    VfxData *ds = static_cast<VfxData *>(instanceData);

    if (activationReason == arInitial) {
        vsapi->requestFrameFilter(n, ds->node, frameCtx);
    } else if (activationReason == arAllFramesReady) {
        for (int i = 0; i <= ds->num_streams; ++i) {
            auto d = i < ds->num_streams ? ds + i : ds + rand() % ds->num_streams;
            std::unique_lock<std::mutex> lock(d->lock, std::defer_lock);
            if (i < ds->num_streams) {
                if (!lock.try_lock()) continue;
            } else {
                lock.lock();
            }

            const VSFrame *src = vsapi->getFrameFilter(n, d->node, frameCtx);

            // TODO: can be refactored out to function creation.
            VSVideoFormat fi;
            vsapi->queryVideoFormat(&fi, cfRGB, d->output_depth == 32 ? stFloat : stInteger, d->output_depth, 0, 0, core);

            assert(vsapi->getFrameHeight(src, 0) == (int)d->in_image_height());
            assert(vsapi->getFrameWidth(src, 0) == (int)d->in_image_width());
            int planes[3] = { 0, 1, 2 };
            const VSFrame *srcf[3] = { nullptr, nullptr, nullptr };
            VSFrame *dst = vsapi->newVideoFrame2(&fi, d->out_image_width(), d->out_image_height(), srcf, planes, src, core);

            auto host = static_cast<char*>(d->srcCpuBuf);
            for (int plane = 0; plane < 3; plane++) {
                const auto stride = vsapi->getStride(src, plane);
                const auto *ptr = vsapi->getReadPtr(src, plane);
                const size_t w = d->in_image_width(), h = d->in_image_height();
                const auto pitch = d->srcTmpImg.pitch;
                vsh::bitblt(host + pitch * h * plane, pitch, ptr, stride, w * d->vi.format.bytesPerSample, h);
            }

            {
                CUDA_MEMCPY2D mcp2d {};
                mcp2d.srcMemoryType = CU_MEMORYTYPE_HOST;
                mcp2d.srcHost = host;
                mcp2d.srcPitch = (size_t)d->srcTmpImg.pitch;
                mcp2d.dstMemoryType = CU_MEMORYTYPE_DEVICE;
                mcp2d.dstDevice = (CUdeviceptr)d->srcTmpImg.pixels;
                mcp2d.dstPitch = (size_t)d->srcTmpImg.pitch;
                mcp2d.WidthInBytes = (size_t)d->in_image_width() * d->vi.format.bytesPerSample;
                mcp2d.Height = d->in_image_height() * 3;
                CK_CUDA(cuMemcpy2DAsync_v2(&mcp2d, d->stream));
            }

            CK_VFX(NvCVImage_Transfer(&d->srcTmpImg, &d->srcGpuImg, d->srcTransferFactor, d->stream, nullptr));
            CK_VFX(NvVFX_Run(d->vfx, 1));
            CK_VFX(NvCVImage_Transfer(&d->dstGpuImg, &d->dstTmpImg, d->dstTransferFactor, d->stream, nullptr));

            host = static_cast<char*>(d->dstCpuBuf);
            {
                CUDA_MEMCPY2D mcp2d {};
                mcp2d.srcMemoryType = CU_MEMORYTYPE_DEVICE;
                mcp2d.srcDevice = (CUdeviceptr)d->dstTmpImg.pixels;
                mcp2d.srcPitch = (size_t)d->dstTmpImg.pitch;
                mcp2d.dstMemoryType = CU_MEMORYTYPE_HOST;
                mcp2d.dstHost = host;
                mcp2d.dstPitch = (size_t)d->dstTmpImg.pitch;
                mcp2d.WidthInBytes = (size_t)d->out_image_width() * d->output_depth / 8;
                mcp2d.Height = d->out_image_height() * 3;
                CK_CUDA(cuMemcpy2DAsync_v2(&mcp2d, d->stream));
            }

            CK_CUDA(cuStreamSynchronize(d->stream));

            for (int plane = 0; plane < 3; plane++) {
                const auto stride = vsapi->getStride(dst, plane);
                auto *ptr = vsapi->getWritePtr(dst, plane);
                const size_t w = d->out_image_width(), h = d->out_image_height();
                const auto pitch = d->dstTmpImg.pitch;
                vsh::bitblt(ptr, stride, host + pitch * h * plane, pitch, w * d->output_depth / 8, h);
            }

            vsapi->freeFrame(src);
            lock.unlock();
            return dst;
        }
    }

    return nullptr;
}

static void VS_CC vfxFree(void *instanceData, VSCore *core, const VSAPI *vsapi) {
    VfxData *ds = static_cast<VfxData *>(instanceData);
    for (int i = 0; i < ds->num_streams; ++i) {
        auto d = ds + i;
        vsapi->freeNode(d->node);
    }

    delete[] ds;
}

static void VS_CC vfxCreate(const VSMap *in, VSMap *out, void *userData, VSCore *core, const VSAPI *vsapi) {
    unsigned int version = 0;
    if (NvVFX_GetVersion(&version) != NVCV_SUCCESS) {
            vsapi->mapSetError(out, "DLVFX: unable to load NvVFX.");
            return;
    }
    int err;
    auto num_streams = vsh::int64ToIntS(vsapi->mapGetInt(in, "num_streams", 0, &err));
    if (err) num_streams = 1;

    std::unique_ptr<VfxData[]> ds(new VfxData[num_streams]);

    for (int i = 0; i < num_streams; ++i) {
        auto d = &ds[i];
        d->num_streams = num_streams;
        size_t op = ~0U;
        enum { OP_AR, OP_SUPERRES, OP_DENOISE };
        const NvVFX_EffectSelector selectors[] = { NVVFX_FX_ARTIFACT_REDUCTION, NVVFX_FX_SUPER_RES, NVVFX_FX_DENOISING };
        try {
            if (autoDllErrors.size() > 0) {
                std::string error, last;
                for (const auto &s: autoDllErrors) {
                    if (error.size()) {
                        if (last != s)
                            error += "; " + s;
                    } else
                        error = s;
                    last = s;
                }
                throw std::runtime_error(error);
            }
            d->node = vsapi->mapGetNode(in, "clip", 0, &err);
            d->vi = *vsapi->getVideoInfo(d->node);

            if (!vsh::isConstantVideoFormat(&d->vi)) {
                throw std::runtime_error("Only clips with constant format and dimensions allowed");
            }
            if (d->vi.format.numPlanes != 3 || d->vi.format.colorFamily != cfRGB)
                throw std::runtime_error("input clip must be RGB format");

            op = vsh::int64ToIntS(vsapi->mapGetInt(in, "op", 0, &err));
            if (err) throw std::runtime_error("op is required argument");
            if (op >= sizeof selectors / sizeof selectors[0])
                throw std::runtime_error("op is out of range.");

            if (op != OP_SUPERRES)
                d->scale = 1;
            else {
                double scale = vsapi->mapGetFloat(in, "scale", 0, &err);
                if (err) scale = 1;
                if (scale < 1)
                    throw std::runtime_error("invalid scale parameter");
                d->scale = scale;
            }

            double strength = vsapi->mapGetFloat(in, "strength", 0, &err);
            if (err) strength = 0;
            d->strength = strength;

            const char *modelDir = vsapi->mapGetData(in, "model_dir", 0, &err);
            if (modelDir == nullptr)
                modelDir = getenv("MODEL_DIR");
            if (modelDir == nullptr)
                modelDir = "C:\\Program Files\\NVIDIA Corporation\\NVIDIA Video Effects\\models";
            fprintf(stderr, "MODEL_DIR = %s\n", modelDir);

            NvCV_Status r = NvVFX_CreateEffect(selectors[op], &d->vfx);
            if (r != NVCV_SUCCESS) {
                const char *err = NvCV_GetErrorStringFromCode(r);
                fprintf(stderr, "NvVFX_CreateEffect failed: %x (%s)\n", r, err);
                throw std::runtime_error("unable to create effect: " + std::string(err));
            }

            CK_VFX(NvVFX_CudaStreamCreate(&d->stream));
            CK_VFX(NvVFX_SetCudaStream(d->vfx, NVVFX_CUDA_STREAM, d->stream));

            if (op == OP_AR || op == OP_SUPERRES) {
                r = NvVFX_SetU32(d->vfx, NVVFX_STRENGTH, int(d->strength));
                if (r != NVCV_SUCCESS)
                    r = NvVFX_SetU32(d->vfx, NVVFX_MODE, int(d->strength));
            } else if (op == OP_DENOISE)
                r = NvVFX_SetF32(d->vfx, NVVFX_STRENGTH, d->strength);
            else
                throw std::runtime_error("unknown op " + std::to_string(op));
            if (r != NVCV_SUCCESS) {
                const char *err = NvCV_GetErrorStringFromCode(r);
                fprintf(stderr, "NvVFX set strength failed: %x (%s)\n", r, err);
                throw std::runtime_error("failed to set strength: " + std::string(err));
            }

            r = NvVFX_SetString(d->vfx, NVVFX_MODEL_DIRECTORY, modelDir);
            if (r != NVCV_SUCCESS) {
                fprintf(stderr, "NvVFX set model directory to %s failed: %x (%s)\n", modelDir, r, NvCV_GetErrorStringFromCode(r));
                throw std::runtime_error("unable to set model directory " + std::string(modelDir));
            }
        } catch (std::runtime_error &e) {
            if (d->node)
                vsapi->freeNode(d->node);
            vsapi->mapSetError(out, (std::string{ "DLVFX: " } + e.what()).c_str());
            return;
        }

        d->in_width = d->vi.width;
        d->in_height = d->vi.height;
        d->vi.width *= d->scale;
        d->vi.height *= d->scale;

        NvCVImage_ComponentType src_ct, dst_ct;
        if (auto bps = d->vi.format.bitsPerSample, st = d->vi.format.sampleType; bps == 32 && st == stFloat) {
            src_ct = NVCV_F32;
	    d->srcTransferFactor = 1.0f;
        } else if (bps == 8 && st == stInteger) {
            src_ct = NVCV_U8;
	    d->srcTransferFactor = 1.0f/255.0f;
        } else {
            throw std::runtime_error("unsupported clip format");
        }

        int output_depth = vsh::int64ToIntS(vsapi->mapGetInt(in, "output_depth", 0, &err));
        if (err) output_depth = d->vi.format.bitsPerSample;
        d->output_depth = output_depth;
        if (output_depth == 32) {
            dst_ct = NVCV_F32;
	    d->dstTransferFactor = 1.0f;
        } else if (output_depth == 8) {
            dst_ct = NVCV_U8;
	    d->dstTransferFactor = 255.0f;;
        } else {
            throw std::runtime_error("unsupported output_depth: only 8 (RGB24) or 32 (RGBS) are supported");
	}

        CK_VFX(NvCVImage_Alloc(&d->srcTmpImg, d->in_image_width(), d->in_image_height(), NVCV_RGB, src_ct, NVCV_PLANAR, NVCV_GPU, 0));
        CK_VFX(NvCVImage_Alloc(&d->srcGpuImg, d->in_image_width(), d->in_image_height(), NVCV_BGR, NVCV_F32, NVCV_PLANAR, NVCV_GPU, 0));
        CK_VFX(NvCVImage_Alloc(&d->dstTmpImg, d->out_image_width(), d->out_image_height(), NVCV_RGB, dst_ct, NVCV_PLANAR, NVCV_GPU, 0));
        CK_VFX(NvCVImage_Alloc(&d->dstGpuImg, d->out_image_width(), d->out_image_height(), NVCV_BGR, NVCV_F32, NVCV_PLANAR, NVCV_GPU, 0));

        CK_CUDA(cuMemHostAlloc(&d->srcCpuBuf, d->srcTmpImg.pitch * d->in_image_height() * 3, CU_MEMHOSTALLOC_WRITECOMBINED));
        CK_CUDA(cuMemHostAlloc(&d->dstCpuBuf, d->dstTmpImg.pitch * d->out_image_height() * 3, 0));

        CK_VFX(NvVFX_SetImage(d->vfx, NVVFX_INPUT_IMAGE, &d->srcGpuImg));
        CK_VFX(NvVFX_SetImage(d->vfx, NVVFX_OUTPUT_IMAGE, &d->dstGpuImg));

        if (op == OP_DENOISE) {
            unsigned int stateSizeInBytes = 0;
            CK_VFX(NvVFX_GetU32(d->vfx, NVVFX_STATE_SIZE, &stateSizeInBytes));
            CK_CUDA(cuMemAlloc_v2(&d->state, stateSizeInBytes));
            CK_CUDA(cuMemsetD8Async(d->state, 0, stateSizeInBytes, d->stream));
            void *stateArray[1] = { d->state };
            CK_VFX(NvVFX_SetObject(d->vfx, NVVFX_STATE, (void*)stateArray));
        }

        CK_VFX(NvVFX_Load(d->vfx));
    }

    // Copy a video info object and set its format to the expected output format.
    VSVideoInfo outputVideoInfo = ds[0].vi;
    vsapi->queryVideoFormat(&outputVideoInfo.format, cfRGB, ds[0].output_depth == 32 ? stFloat : stInteger, ds[0].output_depth, 0, 0, core);

    std::vector<VSFilterDependency> deps;
    deps.reserve(num_streams);
    for (auto i = 0; i < num_streams; i++) {
        deps.emplace_back(ds[i].node, rpStrictSpatial);
    }

    vsapi->createVideoFilter(out, "DLVFX", &outputVideoInfo, vfxGetFrame, vfxFree, fmParallel, deps.data(), deps.size(), ds.release(), core);
}

//////////////////////////////////////////
// Init

#ifndef STANDALONE_VFX
void VS_CC vfxInitialize(VSPlugin *plugin, const VSPLUGINAPI *vsapi) {
#else
VS_EXTERNAL_API(void) VapourSynthPluginInit(VSPlugin *plugin, const VSPLUGINAPI *vsapi) {
    vsapi->configPlugin("info.akarin.plugin", "akarin2", "Experimental Nvidia Maxine plugin", VAPOURSYNTH_API_VERSION, 1, plugin);
#endif
    vsapi->registerFunction("DLVFX", "clip:vnode;op:int;scale:float:opt;strength:float:opt;output_depth:int:opt;num_streams:int:opt;model_dir:data:opt;", "clip:vnode", vfxCreate, nullptr, plugin);
}
