#include "pch.h"
#include <iostream>

// CUDA includes
#include <cuda.h>
#include <nvcuvid.h>
#include <cuviddec.h>

void LogCUDAError(CUresult result, const std::string& operation) {
    const char* errorName;
    const char* errorString;
    cuGetErrorName(result, &errorName);
    cuGetErrorString(result, &errorString);
    std::cout << "[NVDEC ERROR] " << operation << " failed: " << errorName << " (" << errorString << ")" << std::endl;
}

int main() {
    std::cout << "=== NVDEC Debug Test ===" << std::endl;

    // Step 1: Initialize CUDA driver
    std::cout << "Step 1: Initializing CUDA driver..." << std::endl;
    CUresult result = cuInit(0);
    if (result != CUDA_SUCCESS) {
        LogCUDAError(result, "cuInit");
        return 1;
    }
    std::cout << "✓ CUDA driver initialized successfully" << std::endl;

    // Step 2: Get device count
    std::cout << "Step 2: Getting CUDA device count..." << std::endl;
    int deviceCount = 0;
    result = cuDeviceGetCount(&deviceCount);
    if (result != CUDA_SUCCESS) {
        LogCUDAError(result, "cuDeviceGetCount");
        return 1;
    }
    std::cout << "✓ Found " << deviceCount << " CUDA device(s)" << std::endl;

    if (deviceCount == 0) {
        std::cout << "✗ No CUDA devices found" << std::endl;
        return 1;
    }

    // Step 3: Get device 0
    std::cout << "Step 3: Getting CUDA device 0..." << std::endl;
    CUdevice device;
    result = cuDeviceGet(&device, 0);
    if (result != CUDA_SUCCESS) {
        LogCUDAError(result, "cuDeviceGet");
        return 1;
    }
    std::cout << "✓ Got CUDA device 0" << std::endl;

    // Step 4: Get device name
    std::cout << "Step 4: Getting device name..." << std::endl;
    char deviceName[256];
    result = cuDeviceGetName(deviceName, sizeof(deviceName), device);
    if (result != CUDA_SUCCESS) {
        LogCUDAError(result, "cuDeviceGetName");
        return 1;
    }
    std::cout << "✓ Device name: " << deviceName << std::endl;

    // Step 5: Create context
    std::cout << "Step 5: Creating CUDA context..." << std::endl;
    CUcontext context;
    CUctxCreateParams createParams = {};
    createParams.execAffinityParams = nullptr;
    result = cuCtxCreate_v4(&context, &createParams, 0, device);
    if (result != CUDA_SUCCESS) {
        LogCUDAError(result, "cuCtxCreate_v4");
        return 1;
    }
    std::cout << "✓ CUDA context created successfully" << std::endl;

    // Step 6: Check NVDEC decode capabilities for AV1
    std::cout << "Step 6: Checking NVDEC AV1 decode capabilities..." << std::endl;
    CUVIDDECODECAPS decode_caps = {};
    decode_caps.eCodecType = cudaVideoCodec_AV1;
    decode_caps.eChromaFormat = cudaVideoChromaFormat_420;
    decode_caps.nBitDepthMinus8 = 0;

    result = cuvidGetDecoderCaps(&decode_caps);
    if (result != CUDA_SUCCESS) {
        LogCUDAError(result, "cuvidGetDecoderCaps");

        // Try alternative approach - check if CUVID library is available
        std::cout << "Checking alternative NVDEC detection..." << std::endl;

        // Cleanup context first
        cuCtxDestroy(context);
        return 1;
    }

    std::cout << "✓ NVDEC AV1 capabilities:" << std::endl;
    std::cout << "  - Supported: " << (decode_caps.bIsSupported ? "YES" : "NO") << std::endl;
    std::cout << "  - Max width: " << decode_caps.nMaxWidth << std::endl;
    std::cout << "  - Max height: " << decode_caps.nMaxHeight << std::endl;
    std::cout << "  - Max MB count: " << decode_caps.nMaxMBCount << std::endl;
    std::cout << "  - Min width: " << decode_caps.nMinWidth << std::endl;
    std::cout << "  - Min height: " << decode_caps.nMinHeight << std::endl;

    if (!decode_caps.bIsSupported) {
        std::cout << "✗ AV1 codec is not supported by this NVDEC hardware" << std::endl;
        cuCtxDestroy(context);
        return 1;
    }

    std::cout << "✓ NVDEC AV1 support confirmed!" << std::endl;

    // Cleanup
    cuCtxDestroy(context);

    std::cout << "\n=== NVDEC Debug Test: SUCCESS ===" << std::endl;
    std::cout << "NVDEC AV1 hardware decoding is available on this system!" << std::endl;

    return 0;
}