#include "DXGIDuplicator.h"
#include "Utils/Utils.h"
DXGIDuplicator::DXGIDuplicator()
{
}

DXGIDuplicator::~DXGIDuplicator()
{
    if (duplication_)
    {
        duplication_->Release();
    }
    if (device_)
    {
        device_->Release();
    }
    if (deviceContext_)
    {
        deviceContext_->Release();
    }
}
bool DXGIDuplicator::InitD3D11Device(ID3D11Device* g_pd3dDevice, ID3D11DeviceContext* g_pImmediateContext)
{
    device_ = g_pd3dDevice;
    deviceContext_ = g_pImmediateContext;
    return true;
}
bool DXGIDuplicator::InitDuplication()
{
    HRESULT hr = S_OK;

    IDXGIDevice* dxgiDevice = nullptr;
    hr = device_->QueryInterface(__uuidof(IDXGIDevice), reinterpret_cast<void**>(&dxgiDevice));
    if (FAILED(hr))
    {
        return false;
    }

    IDXGIAdapter* dxgiAdapter = nullptr;
    hr = dxgiDevice->GetAdapter(&dxgiAdapter);
    dxgiDevice->Release();
    if (FAILED(hr))
    {
        return false;
    }

    UINT output = 0;
    IDXGIOutput* dxgiOutput = nullptr;
    while (true)
    {
        hr = dxgiAdapter->EnumOutputs(output++, &dxgiOutput);
        if (hr == DXGI_ERROR_NOT_FOUND)
        {
            return false;
        }
        else
        {
            DXGI_OUTPUT_DESC desc;
            dxgiOutput->GetDesc(&desc);
            int width = desc.DesktopCoordinates.right - desc.DesktopCoordinates.left;
            int height = desc.DesktopCoordinates.bottom - desc.DesktopCoordinates.top;
            break;
        }
    }
    dxgiAdapter->Release();

    IDXGIOutput1* dxgiOutput1 = nullptr;
    hr = dxgiOutput->QueryInterface(__uuidof(IDXGIOutput1), reinterpret_cast<void**>(&dxgiOutput1));
    dxgiOutput->Release();
    if (FAILED(hr))
    {
        return false;
    }

    hr = dxgiOutput1->DuplicateOutput(device_, &duplication_);
    dxgiOutput1->Release();
    if (FAILED(hr))
    {
        return false;
    }

    return true;
}
#include <dwmapi.h>
HRESULT DXGIDuplicator::GetDesktopFrame(ID3D11Texture2D*& texture)
{
    HRESULT hr = S_OK;
    DXGI_OUTDUPL_FRAME_INFO frameInfo;
    IDXGIResource* resource = NULL;
    ID3D11Texture2D* pSrcTexture = NULL;
    FreeFrame();
    hr = duplication_->AcquireNextFrame(INFINITE, &frameInfo, &resource);
    if (FAILED(hr))
        return hr;
    frameRef += 1;
    hr = resource->QueryInterface(IID_PPV_ARGS(&texture));
    resource->Release();
    if (FAILED(hr))
        return hr;
    if (frameInfo.LastPresentTime.QuadPart == 0)
        return DXGI_ERROR_WAIT_TIMEOUT;
    return hr;
}
ID3D11Texture2D* DXGIDuplicator::ConvertFormat(ID3D11Texture2D* pSrcTexture, DXGI_FORMAT fmt)
{
    HRESULT hr = S_OK;

    ID3D11Texture2D* pDestTexture = NULL;

    D3D11_TEXTURE2D_DESC desc;
    pSrcTexture->GetDesc(&desc);
    desc.Format = fmt;
    desc.Usage = D3D11_USAGE_STAGING;
    desc.CPUAccessFlags = D3D11_CPU_ACCESS_READ | D3D11_CPU_ACCESS_WRITE;
    desc.BindFlags = 0;
    desc.MiscFlags = 0;
    desc.MipLevels = 1;
    desc.ArraySize = 1;
    desc.SampleDesc.Count = 1;
    desc.SampleDesc.Quality = 0;
    device_->CreateTexture2D(&desc, NULL, &pDestTexture);
    if (FAILED(hr))
        return NULL;
    deviceContext_->CopyResource(pDestTexture, pSrcTexture);
    return pDestTexture;
}
void* DXGIDuplicator::MapBuffer(ID3D11Texture2D* texture2D)
{
    D3D11_MAPPED_SUBRESOURCE mappedResource{};
    HRESULT hr = this->deviceContext_->Map(texture2D, 0, D3D11_MAP_READ, 0, &mappedResource);
    if (FAILED(hr))
        return NULL;
    return mappedResource.pData;
}
void DXGIDuplicator::UnMapBuffer(ID3D11Texture2D* texture2D)
{
    this->deviceContext_->Unmap(texture2D, 0);
}
std::vector<uint8_t> DXGIDuplicator::GenerateBitmapFile(ID3D11Texture2D* texture2D)
{
    D3D11_TEXTURE2D_DESC desc{};
    texture2D->GetDesc(&desc);
    D3D11_MAPPED_SUBRESOURCE mappedResource;
    HRESULT hr = this->deviceContext_->Map(texture2D, 0, D3D11_MAP_READ, 0, &mappedResource);
    if (FAILED(hr))
        return std::vector<uint8_t>();

    size_t bmpSize = desc.Width * desc.Height * 4;
    std::vector<uint8_t> result(sizeof(BITMAPFILEHEADER) + sizeof(BITMAPINFOHEADER) + bmpSize, 0);

    BITMAPFILEHEADER* bmpHeader = (BITMAPFILEHEADER*)(&result[0]);
    BITMAPINFOHEADER* bmiHeader = (BITMAPINFOHEADER*)(&result[0] + sizeof(BITMAPFILEHEADER));

    uint8_t* rgba = (uint8_t*)(&result[0] + sizeof(BITMAPFILEHEADER) + sizeof(BITMAPINFOHEADER));

    bmpHeader->bfSize = sizeof(BITMAPFILEHEADER) + sizeof(BITMAPINFOHEADER) + bmpSize;
    bmpHeader->bfType = 0x4D42;
    bmpHeader->bfOffBits = sizeof(BITMAPFILEHEADER) + sizeof(BITMAPINFOHEADER);
    bmpHeader->bfReserved1 = 0;
    bmpHeader->bfReserved2 = 0;

    bmiHeader->biSize = sizeof(BITMAPINFOHEADER);
    bmiHeader->biWidth = desc.Width;
    bmiHeader->biHeight = 0 - desc.Height;
    bmiHeader->biPlanes = 1;
    bmiHeader->biBitCount = 32;
    bmiHeader->biCompression = 0;
    bmiHeader->biSizeImage = bmpSize;
    bmiHeader->biXPelsPerMeter = 0;
    bmiHeader->biYPelsPerMeter = 0;
    bmiHeader->biClrUsed = 0;
    bmiHeader->biClrImportant = 0;

    memcpy(rgba, mappedResource.pData, bmpSize);
    this->deviceContext_->Unmap(texture2D, 0);
    return result;
}

HRESULT DXGIDuplicator::FreeFrame()
{
    if (frameRef > 0)
    {
        frameRef -= 1;
        return duplication_->ReleaseFrame();
    }
    return DXGI_ERROR_NOT_CURRENTLY_AVAILABLE;
}