#include "stdafx.h"
#include "iostream"
#include "windows.h"
#include "D3D9.h"
#include <Wincodec.h>
#include <chrono>

#pragma comment(lib, "D3d9.lib")

using namespace std;

#define EXIT(hr) { if (FAILED(hr)) { printf("Incorrect hr! Exit..."); goto cleanup; } }
#define RELEASE(p) { if (p != nullptr) p->Release(); }


void Direct3D9TakeScreenshots()
{
    HRESULT hr;
    IDirect3D9 *d3d;
    IDirect3DDevice9 *device;
    IDirect3DSurface9 *surface;
    D3DPRESENT_PARAMETERS parameters = { 0 };
    D3DDISPLAYMODE mode;
    D3DLOCKED_RECT rc;
    LPBYTE *shots;

    // Init D3D and get screen size.
    d3d = Direct3DCreate9(D3D_SDK_VERSION);
    hr = d3d->GetAdapterDisplayMode(D3DADAPTER_DEFAULT, &mode);
    EXIT(hr);

    parameters.Windowed = TRUE;
    parameters.BackBufferCount = 1;
    parameters.BackBufferHeight = mode.Height;
    parameters.BackBufferWidth = mode.Width;
    parameters.SwapEffect = D3DSWAPEFFECT_DISCARD;
    parameters.hDeviceWindow = NULL;

    // Create device.
    hr = d3d->CreateDevice(D3DADAPTER_DEFAULT,
        D3DDEVTYPE_HAL, NULL,
        D3DCREATE_SOFTWARE_VERTEXPROCESSING,
        &parameters, &device);

    EXIT(hr);

    // Capture surface.
    hr = device->CreateOffscreenPlainSurface(mode.Width,
        mode.Height,
        D3DFMT_A8R8G8B8, D3DPOOL_SYSTEMMEM,
        &surface, nullptr);

    EXIT(hr);

    // Compute the required buffer size.

    hr = surface->LockRect(&rc, NULL, 0);
    EXIT(hr);

    UINT pitch = rc.Pitch;

    hr = surface->UnlockRect();
    EXIT(hr);

    for (int i = 0; i < 20; i++)
    {
        auto start = std::chrono::high_resolution_clock::now();

        // Allocate screenshot buffer.
        LPBYTE shot = new BYTE[pitch * mode.Height];

        // Get the data.
        hr = device->GetFrontBufferData(0, surface);
        EXIT(hr);

        // Copy data into the buffers.

        hr = surface->LockRect(&rc, NULL, 0);
        EXIT(hr);

        CopyMemory(shot, rc.pBits, rc.Pitch * mode.Height);

        hr = surface->UnlockRect();
        EXIT(hr);

        auto finish = std::chrono::high_resolution_clock::now();
        std::cout << "Time: " << std::chrono::duration_cast<std::chrono::microseconds>(finish - start).count() << std::endl;
    }

cleanup:

    RELEASE(surface);
    RELEASE(device);
    RELEASE(d3d);
}

void main()
{
    Direct3D9TakeScreenshots();
    system("pause");
}