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

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

#ifndef _WIN32
#error "Unsupported platform"
#else
static std::vector<std::string> autoDllErrors;
static const wchar_t *dllPath(const wchar_t *suffix);
#define NGX_DLL dllPath(L"dlisr.dll"),".dlisr.dll",autoDllErrors
#define CUDA_DLL L"nvcuda.dll","nvcuda.dll",autoDllErrors
#endif
#include "cuda.h"
#include "ngx.h"

#define CK_NGX(x) do { \
    int r = (x); \
    if (r != NVSDK_NGX_Result_Success) { \
        fprintf(stderr, "failed NGX call %s: %x at line %d\n", #x, r, __LINE__); \
        abort(); \
    } \
} while (0)
#define CK_CUDA(x) do { \
    int r = (x); \
    if (r != CUDA_SUCCESS) { \
        fprintf(stderr, "failed cuda call %s: %d at line %d\n", #x, r, __LINE__); \
        abort(); \
    } \
} while (0)

#ifdef _WIN32
static const wchar_t *dllPath(const wchar_t *suffix) {
    static const std::wstring res = [suffix]() -> std::wstring {
        HMODULE mod = 0;
        if (GetModuleHandleExA(GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS | GET_MODULE_HANDLE_EX_FLAG_UNCHANGED_REFCOUNT, (char *)dllPath, &mod)) {
            std::vector<wchar_t> buf;
            size_t n = 0;
            do {
                buf.resize(buf.size() + MAX_PATH);
                n = GetModuleFileNameW(mod, buf.data(), buf.size());
            } while (n >= buf.size());
            buf.resize(n);
            std::wstring stem(buf.begin(), buf.end() - 4);
            std::wstring path = stem + L"." + suffix;
            if (std::filesystem::exists(path)) return path;
            path = stem + L"/" + suffix;
            return path;
        }
        throw std::runtime_error("unable to locate myself");
    }();
    return res.c_str();
}
#endif

static void *cudaMalloc(size_t size) {
    void *ptr = nullptr;
    CK_CUDA(cuMemAlloc_v2(&ptr, size));
    return ptr;
}

struct NgxData {
    std::mutex lock;

    VSNode *node;
    VSVideoInfo vi;
    int scale;

    typedef float T;
    uint64_t pixel_size() const { return 3 * sizeof(T); }
    uint64_t in_image_width() const   { return vi.width / scale; }
    uint64_t out_image_width() const  { return vi.width; }
    uint64_t in_image_height() const  { return vi.height / scale; }
    uint64_t out_image_height() const { return vi.height; }
    uint64_t in_image_row_bytes() const  { return pixel_size() * in_image_width(); }
    uint64_t out_image_row_bytes() const { return pixel_size() * out_image_width(); }
    uint64_t in_size() const  { return in_image_height()  * in_image_row_bytes(); }
    uint64_t out_size() const { return out_image_height() * out_image_row_bytes(); }

    NVSDK_NGX_Parameter *param;
    NVSDK_NGX_Handle *DUHandle;
    CUcontext ctx;

    std::vector<uint8_t> in_host, out_host;
    CUdeviceptr inp, outp;
    void allocate() {
        in_host.resize(in_size());
        out_host.resize(out_size());
        inp = cudaMalloc(in_size());
        outp = cudaMalloc(out_size());
    }

    NgxData() : node(nullptr), vi(), scale(0), param(nullptr), DUHandle(nullptr), ctx(nullptr), inp(nullptr), outp(nullptr) {}
    ~NgxData() {
        if (ctx) {
            CK_CUDA(cuCtxPushCurrent_v2(ctx));
            if (inp) CK_CUDA(cuMemFree_v2(inp));
            if (outp) CK_CUDA(cuMemFree_v2(outp));
            if (DUHandle) CK_NGX(NVSDK_NGX_CUDA_ReleaseFeature(DUHandle));
            cuCtxPopCurrent_v2(nullptr);
        }
    }
};

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

    if (activationReason == arInitial) {
        vsapi->requestFrameFilter(n, d->node, frameCtx);
    } else if (activationReason == arAllFramesReady) {
        const VSFrame *src = vsapi->getFrameFilter(n, d->node, frameCtx);

        const VSVideoFormat fi = d->vi.format;
        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);

        // The NGX API is not thread safe.
        std::lock_guard<std::mutex> lock(d->lock);
        CK_CUDA(cuCtxPushCurrent_v2(d->ctx));

        auto params = d->param;
        params->Set(NVSDK_NGX_Parameter_Width, (uint64_t)d->in_image_width());
        params->Set(NVSDK_NGX_Parameter_Height, (uint64_t)d->in_image_height());
        params->Set(NVSDK_NGX_Parameter_Scale, d->scale);

        // Create the feature
        //CK_NGX(NVSDK_NGX_CUDA_CreateFeature(NVSDK_NGX_Feature_ImageSuperResolution, params, &d->DUHandle));

        void *in_image_dev_ptr = d->inp;
        void *out_image_dev_ptr = d->outp;;

        uint8_t *host = d->in_host.data();
        typedef float T;
        const T factor = 255.0f;
        for (int plane = 0; plane < 3; plane++) {
            const size_t stride = vsapi->getStride(src, plane);
            const uint8_t *ptr = (uint8_t*)vsapi->getReadPtr(src, plane);
            for (size_t i = 0; i < d->in_image_height(); i++)
                for (size_t j = 0; j < d->in_image_width(); j++)
                    *(T*)&host[i * d->in_image_row_bytes() + j * d->pixel_size() + plane * sizeof(T)] = *(T*)&ptr[i * stride + j * sizeof(T)] * factor;
        }
        CK_CUDA(cuMemcpyHtoD_v2(in_image_dev_ptr, host, d->in_size()));

        // Pass the pointers to the GPU allocations to the
        // parameter block along with the format and size.
        params->Set(NVSDK_NGX_Parameter_Color_SizeInBytes, d->in_size());
        params->Set(NVSDK_NGX_Parameter_Color_Format, NVSDK_NGX_Buffer_Format_RGB32F);
        params->Set(NVSDK_NGX_Parameter_Color, in_image_dev_ptr);
        params->Set(NVSDK_NGX_Parameter_Output_SizeInBytes, d->out_size());
        params->Set(NVSDK_NGX_Parameter_Output_Format, NVSDK_NGX_Buffer_Format_RGB32F);
        params->Set(NVSDK_NGX_Parameter_Output, out_image_dev_ptr);

        // Execute the feature.
        CK_NGX(NVSDK_NGX_CUDA_EvaluateFeature(d->DUHandle, params, nullptr));

        host = d->out_host.data();
        CK_CUDA(cuMemcpyDtoH_v2(host, out_image_dev_ptr, d->out_size()));
        for (int plane = 0; plane < 3; plane++) {
            const size_t stride = vsapi->getStride(dst, plane);
            uint8_t *ptr = (uint8_t*)vsapi->getWritePtr(dst, plane);
            for (size_t i = 0; i < d->out_image_height(); i++)
                for (size_t j = 0; j < d->out_image_width(); j++)
                    *(T*)&ptr[i * stride + j * sizeof(T)] = *(T*)&host[i * d->out_image_row_bytes() + j * d->pixel_size() + plane * sizeof(T)] / factor;
        }

        cuCtxPopCurrent_v2(nullptr);

        vsapi->freeFrame(src);
        return dst;
    }

    return nullptr;
}

static void VS_CC ngxFree(void *instanceData, VSCore *core, const VSAPI *vsapi) {
    NgxData *d = static_cast<NgxData *>(instanceData);
    vsapi->freeNode(d->node);

    delete d;
}

static void VS_CC ngxCreate(const VSMap *in, VSMap *out, void *userData, VSCore *core, const VSAPI *vsapi) {
    std::unique_ptr<NgxData> d(new NgxData);
    int err;

    int devid = 0;
    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");
        if (d->vi.format.sampleType != stFloat || d->vi.format.bitsPerSample != 32)
            throw std::runtime_error("input clip must be 32-bit float format");

        int scale = vsh::int64ToIntS(vsapi->mapGetInt(in, "scale", 0, &err));
        if (err) scale = 2;
        if (scale != 2 && scale != 4 && scale != 8)
            throw std::runtime_error("scale must be 2/4/8");
        d->scale = scale;

        devid = vsh::int64ToIntS(vsapi->mapGetInt(in, "device_id", 0, &err));
        if (err) devid = 0;
    } catch (std::runtime_error &e) {
        if (d->node)
            vsapi->freeNode(d->node);
        vsapi->mapSetError(out, (std::string{ "DLISR: " } + e.what()).c_str());
        return;
    }

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

    static bool inited = []() -> bool {
        CUcontext ctx;
        bool hasCtx = cuCtxGetCurrent(&ctx) == CUDA_SUCCESS;
        CK_NGX(NVSDK_NGX_CUDA_Init(0, L"./", NVSDK_NGX_Version_API));
        // We don't expect NVSDK_NGX_CUDA_Init to create a context, but if it did, we need to
        // switch to save a global CUDA context, instead of the pre-filter context.
        if (!hasCtx && cuCtxGetCurrent(&ctx) == CUDA_SUCCESS) {
            fprintf(stderr, "invariant violated: NVSDK_NGX_CUDA_Init created CUDA context: %p\n", ctx);
            abort();
        }
        return true;
    }();
    (void) inited;
    NV_new_Parameter(&d->param);

    d->param->Set(NVSDK_NGX_Parameter_Width, d->in_image_width());
    d->param->Set(NVSDK_NGX_Parameter_Height, d->in_image_height());
    d->param->Set(NVSDK_NGX_Parameter_Scale, d->scale);

    // Get the scratch buffer size and create the scratch allocation.
    size_t byteSize{ 0u };
    CK_NGX(NVSDK_NGX_CUDA_GetScratchBufferSize(NVSDK_NGX_Feature_ImageSuperResolution, d->param, &byteSize));
    if (byteSize != 0) // should request none.
        abort();

    // Create the feature
    CUdevice dev = 0;
    CK_CUDA(cuInit(0));
    CK_CUDA(cuDeviceGet(&dev, devid));
    CK_CUDA(cuCtxCreate_v2(&d->ctx, 0, dev));
    CK_NGX(NVSDK_NGX_CUDA_CreateFeature(NVSDK_NGX_Feature_ImageSuperResolution, d->param, &d->DUHandle));
    CK_CUDA(cuCtxGetCurrent(&d->ctx));
    d->allocate();
    CK_CUDA(cuCtxPopCurrent_v2(nullptr));

    VSFilterDependency deps[] = {{d->node, rpStrictSpatial}};

    const VSVideoInfo *vi = &d->vi;
    vsapi->createVideoFilter(out, "DLISR", vi, ngxGetFrame, ngxFree, fmParallel, deps, 1, d.release(), core);
}

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

#ifndef STANDALONE_NGX
void VS_CC ngxInitialize(VSPlugin *plugin, const VSPLUGINAPI *vsapi) {
#else
VS_EXTERNAL_API(void) VS_CC VapourSynthPluginInit(VSPlugin *plugin, const VSPLUGINAPI *vsapi) {
    vsapi->configPlugin("info.akarin.plugin", "akarin2", "Experimental Nvidia DLISR plugin", VAPOURSYNTH_API_VERSION, 1, plugin);
#endif
    vsapi->registerFunction("DLISR", "clip:vnode;scale:int:opt;device_id:int:opt;", "clip:vnode", ngxCreate, nullptr, plugin);
}
