//
// Copyright (c) 2021 Advanced Micro Devices, Inc. All rights reserved.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.
//

#include <iostream>
#include <dxgi1_4.h>
#include <d3d11_3.h>

#pragma comment( lib, "dxgi" )
#pragma comment( lib, "d3d11" )

bool isUMA(ID3D11Device* pDevice)
{
	bool result = false;
	ID3D11Device3* pD3D11Device3 = nullptr;
	if (S_OK == pDevice->QueryInterface(IID_PPV_ARGS(&pD3D11Device3)) && pD3D11Device3)
	{
		D3D11_FEATURE_DATA_D3D11_OPTIONS2 data = {};
		if (S_OK == pD3D11Device3->CheckFeatureSupport(
			D3D11_FEATURE_D3D11_OPTIONS2,
			&data,
			sizeof(data)))
		{
			result = data.UnifiedMemoryArchitecture;
		}
		pD3D11Device3->Release();
	}
	return result;
}

int main()
{
	UINT flags = NULL; // D3D11_CREATE_DEVICE_SINGLETHREADED;		
	D3D_FEATURE_LEVEL    featureLevels[] = { D3D_FEATURE_LEVEL_11_0 };
	UINT                 numFeatureLevels = ARRAYSIZE(featureLevels);
	D3D_FEATURE_LEVEL    featureLevel;
	ID3D11Device* pDevice = nullptr;
	ID3D11DeviceContext* pImmediateContext = nullptr;
	if SUCCEEDED(D3D11CreateDevice(
		NULL,
		D3D_DRIVER_TYPE_HARDWARE,
		NULL,
		flags,
		featureLevels,
		numFeatureLevels,
		D3D11_SDK_VERSION,
		&pDevice,
		&featureLevel,
		&pImmediateContext))
	{
		IDXGIDevice* pDXGIDevice = nullptr;
		IDXGIAdapter* pAdapter = nullptr;
		DXGI_ADAPTER_DESC desc;
		if (SUCCEEDED(pDevice->QueryInterface(__uuidof(IDXGIDevice), (void**)&pDXGIDevice))
			&& SUCCEEDED(pDXGIDevice->GetAdapter(&pAdapter))
			&& SUCCEEDED(pAdapter->GetDesc(&desc)))
		{
			printf("DedicatedVideoMemory %I64u\n", desc.DedicatedVideoMemory);
			printf("DedicatedSystemMemory %I64u\n", desc.DedicatedSystemMemory);
			printf("SharedSystemMemory %I64u\n", desc.SharedSystemMemory);
			printf("isUMA %i\n", isUMA(pDevice));
			SIZE_T budget = desc.DedicatedVideoMemory;
			if (isUMA(pDevice))
			{
				budget += desc.DedicatedSystemMemory + desc.SharedSystemMemory;
			}
			IDXGIAdapter3* pAdapter3 = nullptr;
			DXGI_QUERY_VIDEO_MEMORY_INFO info = {};
			if (SUCCEEDED(pAdapter->QueryInterface(__uuidof(IDXGIAdapter3), (void**)&pAdapter3))
				&& SUCCEEDED(pAdapter3->QueryVideoMemoryInfo(0, DXGI_MEMORY_SEGMENT_GROUP_LOCAL, &info)))
			{
				budget = info.Budget;
			}
			printf("budget %I64u\n", budget);
		}
	}
}