#include "dxgi_capturer.h"
#include <windows.h>

#include "basic/basic.h"
#include "basic/frame.h"
#include "d3d/buffer_filler.h"
#include "d3d/convert.h"
#include "d3d/gen_frame.h"
#include <d3d11.h>
#include <dxgi1_2.h>

class DxgiCapturerPrivate
{
public:
    DxgiCapturerPrivate() {}

    bool _bInit = false;
    bool _isCaptureSuccess = false;

    ID3D11Device* _hDevice = nullptr;
    ID3D11DeviceContext* _hContext = nullptr;
    IDXGIOutputDuplication* _hDeskDupl = nullptr;
    IDXGISurface1* _hStagingSurf = nullptr;
    ID3D11Texture2D* _gdiImage = nullptr;
    D3D11_TEXTURE2D_DESC _desc;
    bool _isAttached = false;
    AVFrame* _xrgbFrame = nullptr;
    AVFrame* _nv12Frame = nullptr;
    BufferFiller _xrgbBuffers;
    BufferFiller _nv12Buffers;
    D3dConverter _rgbToNv12;
};

DxgiCapturer::DxgiCapturer()
    : d(new DxgiCapturerPrivate)
{
    ZeroMemory(&d->_desc, sizeof(d->_desc));
}

DxgiCapturer::~DxgiCapturer()
{
    Close();
    delete d;
}

bool DxgiCapturer::Open(int left, int top, int width, int height)
{
    Close();
    HRESULT hr = S_OK;
    d->_isAttached = false;

    if (d->_bInit) {
        return false;
    }

    // Driver types supported
    D3D_DRIVER_TYPE DriverTypes[] = {
        D3D_DRIVER_TYPE_HARDWARE,
        D3D_DRIVER_TYPE_WARP,
        D3D_DRIVER_TYPE_REFERENCE,
    };
    UINT NumDriverTypes = ARRAYSIZE(DriverTypes);

    // Feature levels supported
    D3D_FEATURE_LEVEL FeatureLevels[] = {
        D3D_FEATURE_LEVEL_11_0,
        D3D_FEATURE_LEVEL_10_1,
        D3D_FEATURE_LEVEL_10_0,
        D3D_FEATURE_LEVEL_9_1};
    UINT NumFeatureLevels = ARRAYSIZE(FeatureLevels);

    D3D_FEATURE_LEVEL FeatureLevel;

    // Create D3D device
    for (UINT DriverTypeIndex = 0; DriverTypeIndex < NumDriverTypes; ++DriverTypeIndex) {
        hr = D3D11CreateDevice(nullptr,
                               DriverTypes[DriverTypeIndex],
                               nullptr,
                               0,
                               FeatureLevels,
                               NumFeatureLevels,
                               D3D11_SDK_VERSION,
                               &d->_hDevice,
                               &FeatureLevel,
                               &d->_hContext);
        if (SUCCEEDED(hr)) {
            break;
        }
    }
    __CheckBool(SUCCEEDED(hr));

    // Get DXGI device
    IDXGIDevice* hDxgiDevice = nullptr;
    __CheckBool(SUCCEEDED(d->_hDevice->QueryInterface(__uuidof(IDXGIDevice),
                                                      reinterpret_cast<void**>(&hDxgiDevice))));

    // Get DXGI adapter
    IDXGIAdapter* hDxgiAdapter = nullptr;
    hr = hDxgiDevice->GetParent(__uuidof(IDXGIAdapter), reinterpret_cast<void**>(&hDxgiAdapter));
    Free(hDxgiDevice, [=] { hDxgiDevice->Release(); });
    __CheckBool(SUCCEEDED(hr));

    // Get output
    INT nOutput = 0;
    IDXGIOutput* hDxgiOutput = nullptr;
    DXGI_OUTPUT_DESC dxgiOutDesc;
    ZeroMemory(&dxgiOutDesc, sizeof(dxgiOutDesc));

    for (int idx = 0; SUCCEEDED(hr = hDxgiAdapter->EnumOutputs(idx, &hDxgiOutput)); ++idx) {
        // get output description struct
        hDxgiOutput->GetDesc(&dxgiOutDesc);
        if (dxgiOutDesc.DesktopCoordinates.left == left
            && dxgiOutDesc.DesktopCoordinates.top == top) { // 寻找显示器
            break;
        }
    }
    Free(hDxgiAdapter, [=] { hDxgiAdapter->Release(); });
    __CheckBool(SUCCEEDED(hr));

    // QI for Output 1
    IDXGIOutput1* hDxgiOutput1 = nullptr;
    hr = hDxgiOutput->QueryInterface(__uuidof(hDxgiOutput1), reinterpret_cast<void**>(&hDxgiOutput1));
    Free(hDxgiOutput, [=] { hDxgiOutput->Release(); });
    __CheckBool(SUCCEEDED(hr));

    // Create desktop duplication
    hr = hDxgiOutput1->DuplicateOutput(d->_hDevice, &d->_hDeskDupl);
    Free(hDxgiOutput1, [=] { hDxgiOutput1->Release(); });
    __CheckBool(SUCCEEDED(hr));

    // Set ColorSpace
    D3D11_VIDEO_PROCESSOR_COLOR_SPACE inputColorSpace;
    inputColorSpace.Usage = 1;
    inputColorSpace.RGB_Range = 0;
    inputColorSpace.YCbCr_Matrix = 1;
    inputColorSpace.YCbCr_xvYCC = 0;
    inputColorSpace.Nominal_Range = D3D11_VIDEO_PROCESSOR_NOMINAL_RANGE_0_255;

    D3D11_VIDEO_PROCESSOR_COLOR_SPACE outputColorSpace;
    outputColorSpace.Usage = 0;
    outputColorSpace.RGB_Range = 0;
    outputColorSpace.YCbCr_Matrix = 1;
    outputColorSpace.YCbCr_xvYCC = 0;
    outputColorSpace.Nominal_Range = D3D11_VIDEO_PROCESSOR_NOMINAL_RANGE_16_235;
    d->_rgbToNv12.Open(d->_hDevice, d->_hContext, inputColorSpace, outputColorSpace);
    d->_nv12Frame = Frame<MediaType::VIDEO>::Alloc(AV_PIX_FMT_NV12, width, height);
    d->_xrgbFrame = Frame<MediaType::VIDEO>::Alloc(AV_PIX_FMT_BGR0, width, height);
    __CheckBool(d->_nv12Frame);
    __CheckBool(d->_xrgbFrame);
    // 初始化成功
    d->_bInit = true;
    return true;
}
void DxgiCapturer::Close()
{
    if (!d->_bInit) {
        return;
    }

    d->_bInit = false;
    d->_nv12Buffers.Clear();
    d->_xrgbBuffers.Clear();
    d->_rgbToNv12.Close();
    Free(d->_nv12Frame, [this] { av_frame_free(&d->_nv12Frame); });
    Free(d->_xrgbFrame, [this] { av_frame_free(&d->_xrgbFrame); });
    Free(d->_hDeskDupl, [this] { d->_hDeskDupl->Release(); });
    Free(d->_hDevice, [this] { d->_hDevice->Release(); });
    Free(d->_hContext, [this] { d->_hContext->Release(); });
}

HDC DxgiCapturer::GetHdc()
{
    d->_isCaptureSuccess = false;
    if (!d->_bInit) {
        return nullptr;
    }

    IDXGIResource* hDesktopResource = nullptr;
    DXGI_OUTDUPL_FRAME_INFO FrameInfo;
    HRESULT hr = d->_hDeskDupl->AcquireNextFrame(0, &FrameInfo, &hDesktopResource);
    if (FAILED(hr)) {
        if (hr == DXGI_ERROR_WAIT_TIMEOUT) { // 这里是因为当桌面没有动画更新时就会有一个错误值，不进行错误打印
            return nullptr;
        }
        return nullptr;
    }

    // query next frame staging buffer
    ID3D11Texture2D* srcImage = nullptr;
    hr = hDesktopResource->QueryInterface(__uuidof(ID3D11Texture2D), reinterpret_cast<void**>(&srcImage));
    Free(hDesktopResource, [=] { hDesktopResource->Release(); });
    __CheckNullptr(SUCCEEDED(hr));

    srcImage->GetDesc(&d->_desc);

    // create a new staging buffer for fill frame image
    auto desc = d->_desc;
    desc.ArraySize = 1;
    desc.BindFlags = D3D11_BIND_FLAG::D3D11_BIND_RENDER_TARGET;
    desc.MiscFlags = D3D11_RESOURCE_MISC_GDI_COMPATIBLE;
    desc.SampleDesc.Count = 1;
    desc.SampleDesc.Quality = 0;
    desc.MipLevels = 1;
    desc.CPUAccessFlags = 0;
    desc.Usage = D3D11_USAGE_DEFAULT;
    hr = d->_hDevice->CreateTexture2D(&desc, nullptr, &d->_gdiImage);
    if (FAILED(hr)) {
        __DebugPrint("Create _gdiImage failed");
        Free(srcImage, [=] { srcImage->Release(); });
        Free(d->_hDeskDupl, [this] { d->_hDeskDupl->ReleaseFrame(); });
        return nullptr;
    }

    // copy next staging buffer to new staging buffer
    d->_hContext->CopyResource(d->_gdiImage, srcImage);
    Free(srcImage, [=] { srcImage->Release(); });
    d->_hDeskDupl->ReleaseFrame();

    // create staging buffer for map bits
    d->_hStagingSurf = nullptr;
    hr = d->_gdiImage->QueryInterface(__uuidof(IDXGISurface), (void**) (&d->_hStagingSurf));
    if (FAILED(hr)) {
        __DebugPrint("_gdiImage->QueryInterface failed");
        Free(d->_gdiImage, [this] { d->_gdiImage->Release(); });
        return nullptr;
    }

    d->_isCaptureSuccess = true;
    HDC hdc = nullptr;
    // if GetDc is failed, the hdc is nullptr
    d->_hStagingSurf->GetDC(FALSE, &hdc);
    return hdc;
}

AVFrame* DxgiCapturer::GetFrame()
{
    if (!d->_isCaptureSuccess) {
        return nullptr;
    }
    d->_isCaptureSuccess = false;
    d->_hStagingSurf->ReleaseDC(nullptr);

    // 创建一个临时的纹理
    ID3D11Texture2D* tmpImage = nullptr;
    d->_desc.MiscFlags = 2050;
    __CheckNullptr(SUCCEEDED(d->_hDevice->CreateTexture2D(&d->_desc, nullptr, &tmpImage)));
    d->_hContext->CopyResource(tmpImage, d->_gdiImage);

    // 首先尝试创建 NV12 纹理
    AVFrame* frame = nullptr;
    auto tmpFormat = d->_desc.Format;
    d->_desc.Format = DXGI_FORMAT_NV12;
    if (GenNv12Frame(d->_hDevice,
                     d->_hContext,
                     d->_desc,
                     tmpImage,
                     d->_nv12Buffers,
                     d->_nv12Frame,
                     d->_rgbToNv12)) {
        frame = d->_nv12Frame;
    } else {
        d->_desc.Format = tmpFormat;
        GenRgbFrame(d->_hDevice,
                    d->_hContext,
                    d->_desc,
                    d->_gdiImage,
                    d->_xrgbBuffers,
                    d->_xrgbFrame);
        frame = d->_xrgbFrame;
    }
    Free(d->_hStagingSurf, [this] { d->_hStagingSurf->Release(); });
    Free(tmpImage, [&tmpImage] { tmpImage->Release(); });
    Free(d->_gdiImage, [this] { d->_gdiImage->Release(); });

    return frame;
}