/*
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/.
 */

#define _WIN32_DCOM

#include "Host/SystemInfo.hpp"
#include "Base/Logger.hpp"
#include <iostream>
#include <utf8.h>
#include <Windows.h>
#include <Psapi.h>
#include <Lmcons.h>
#include <Winsock.h>
#include <Winbase.h>
#include <powersetting.h>
#include <guiddef.h>
#include <comdef.h>
#include <Wbemidl.h>

// NOTE(vb): Windows defines macros like GetUserName that are used to automatically select the appropriate function version (GetUserNameA for ANSI and GetUserNameW for Unicode)
// based on whether the _UNICODE macro is defined, so we manually undefine these macros to avoid naming collisions. 
#undef GetUserName

// link this for power settings
#pragma comment(lib, "PowrProf.lib")
#pragma comment(lib, "wbemuuid.lib")

namespace OpenVulkano
{
	namespace
	{
		enum class SYS_MEM_TYPE { TOTAL_PHYS, AVAIL_PHYS };
		
		enum BatteryChargeStatus
		{
			MEDIUM = 0,   // [low, high]
			HIGH = 1,     // > 66%
			LOW = 2,      // < 33%
			CRITICAL = 4, // <5% 
			CHARGING = 8,
			NO_SYSTEM_BATTERY = 128,
			UNKNOWN = 255
		};

		size_t ReadSystemMemInfo(SYS_MEM_TYPE type)
		{
			MEMORYSTATUSEX status;
			status.dwLength = sizeof(status);
			if (GlobalMemoryStatusEx(&status))
			{
				switch (type)
				{
					case SYS_MEM_TYPE::TOTAL_PHYS: return status.ullTotalPhys;
					case SYS_MEM_TYPE::AVAIL_PHYS: return status.ullAvailPhys;
				}
			}
			Logger::PERF->error("Failed to get system RAM info");
			return 0;
		}

		enum class APP_MEM_TYPE { VM_MAX, USED };

		size_t ReadAppMemInfo(APP_MEM_TYPE type)
		{
			PROCESS_MEMORY_COUNTERS_EX counters;
			if (GetProcessMemoryInfo(GetCurrentProcess(), reinterpret_cast<PROCESS_MEMORY_COUNTERS*>(&counters), sizeof(counters)))
			{
				switch(type)
				{
					case APP_MEM_TYPE::VM_MAX: return counters.PeakPagefileUsage;
					case APP_MEM_TYPE::USED: return counters.PrivateUsage;
				}
			}
			Logger::PERF->error("Failed to get process memory info");
			return 0;
		}

		size_t GetCommitLimit()
		{
			MEMORYSTATUSEX status;
			status.dwLength = sizeof(status);
			if(GlobalMemoryStatusEx(&status))
			{
				return status.ullTotalPageFile;
			}
			Logger::PERF->error("Failed to get system commit limit");
			return 0;
		}

		std::optional<std::string> GetWMIProperty(const std::string& propName)
		{
			std::optional<std::string> res;

			// https://learn.microsoft.com/en-us/windows/win32/wmisdk/example-creating-a-wmi-application?redirectedfrom=MSDN
			HRESULT hres;

			// Initialize COM.
			hres = CoInitializeEx(0, COINIT_MULTITHREADED);
			if (FAILED(hres))
			{
				Logger::APP->error("Failed to initialize COM library. Error code = 0x{}", hres);
				return {};
			}

			// Initialize
			hres = CoInitializeSecurity(NULL,
			                            -1, // COM negotiates service
			                            NULL, // Authentication services
			                            NULL, // Reserved
			                            RPC_C_AUTHN_LEVEL_DEFAULT, // authentication
			                            RPC_C_IMP_LEVEL_IMPERSONATE, // Impersonation
			                            NULL, // Authentication info
			                            EOAC_NONE, // Additional capabilities
			                            NULL // Reserved
			);


			if (FAILED(hres))
			{
				Logger::APP->error("Failed to initialize security. Error code = 0x{}", hres);
				CoUninitialize();
				return {};
			}

			// Obtain the initial locator to Windows Management
			// on a particular host computer.
			IWbemLocator* pLoc = 0;

			hres = CoCreateInstance(CLSID_WbemLocator, 0, CLSCTX_INPROC_SERVER, IID_IWbemLocator, (LPVOID*) &pLoc);

			if (FAILED(hres))
			{
				Logger::APP->error("Failed to create IWbemLocator object. Error code = 0x{}", hres);
				CoUninitialize();
				return {};
			}

			IWbemServices* pSvc = 0;

			// Connect to the root\cimv2 namespace with the
			// current user and obtain pointer pSvc
			// to make IWbemServices calls.

			hres = pLoc->ConnectServer(
				_bstr_t(L"ROOT\\CIMV2"), // WMI namespace
				NULL, // User name
				NULL, // User password
				0, // Locale
				NULL, // Security flags
				0, // Authority
				0, // Context object
				&pSvc // IWbemServices proxy
			);

			if (FAILED(hres))
			{
				Logger::APP->error("Could not connect. Error code = 0x{}", hres);
				pLoc->Release();
				CoUninitialize();
				return {};
			}

			// Set the IWbemServices proxy so that impersonation
			// of the user (client) occurs.
			hres = CoSetProxyBlanket(
				pSvc, // the proxy to set
				RPC_C_AUTHN_WINNT, // authentication service
				RPC_C_AUTHZ_NONE, // authorization service
				NULL, // Server principal name
				RPC_C_AUTHN_LEVEL_CALL, // authentication level
				RPC_C_IMP_LEVEL_IMPERSONATE, // impersonation level
				NULL, // client identity
				EOAC_NONE // proxy capabilities
			);

			if (FAILED(hres))
			{
				Logger::APP->error("Could not set proxy blanket. Error code = 0x{}", hres);
				pSvc->Release();
				pLoc->Release();
				CoUninitialize();
				return {};
			}

			// Use the IWbemServices pointer to make requests of WMI.
			// Make requests here:

			// For example, query for all the running processes
			IEnumWbemClassObject* pEnumerator = NULL;
			hres = pSvc->ExecQuery(bstr_t("WQL"), bstr_t("SELECT * FROM Win32_ComputerSystem"),
			                       WBEM_FLAG_FORWARD_ONLY | WBEM_FLAG_RETURN_IMMEDIATELY, NULL, &pEnumerator);

			if (FAILED(hres))
			{
				Logger::APP->error("Query for processes failed. Error code = 0x{}", hres);
				pSvc->Release();
				pLoc->Release();
				CoUninitialize();
				return {};
			}
			else
			{
				IWbemClassObject* pclsObj;
				ULONG uReturn = 0;

				while (pEnumerator)
				{
					hres = pEnumerator->Next(WBEM_INFINITE, 1, &pclsObj, &uReturn);

					if (0 == uReturn)
					{
						break;
					}

					VARIANT vtProp;
					ZeroMemory(&vtProp, sizeof(vtProp));

					std::wstring propNameUtf16;
					utf8::utf8to16(propName.begin(), propName.end(), std::back_inserter(propNameUtf16));

					// Get the value of the specified property
					hres = pclsObj->Get(propNameUtf16.data(), 0, &vtProp, 0, 0);

					if (FAILED(hres))
					{
						Logger::APP->error("Failed to get property {}. Error code = 0x{}", propName, hres);
					}
					else
					{
						// init optional
						res = "";
						std::wstring propValue = vtProp.bstrVal;
						utf8::utf16to8(propValue.begin(), propValue.end(), std::back_inserter(*res));
					}
					VariantClear(&vtProp);
					pclsObj->Release();
					pclsObj = NULL;
				}
			}

			// Cleanup
			// ========

			pSvc->Release();
			pLoc->Release();
			pEnumerator->Release();

			CoUninitialize();

			return res;
		}

		std::string GetHumanReadableOSName()
		{
			NTSTATUS(WINAPI *RtlGetVersion)(LPOSVERSIONINFOEXW);
			OSVERSIONINFOEXW info;
			*(FARPROC*)&RtlGetVersion = GetProcAddress(GetModuleHandleA("ntdll"), "RtlGetVersion");
			info.dwOSVersionInfoSize = sizeof(info);
			RtlGetVersion(&info);
			if (info.wProductType == VER_NT_WORKSTATION)
			{
				if (info.dwMajorVersion == 10)
				{
					if (info.dwBuildNumber >= 22000)
						return "Windows 11";
					return "Windows 10";
				}
				else if (info.dwMajorVersion == 6)
				{
					switch (info.dwMinorVersion)
					{
						case 0: return "Windows Vista";
						case 1: return "Windows 7";
						case 2: return "Windows 8";
						case 3: return "Windows 8.1";
					}
				}

				return "Windows " + std::to_string(info.dwMajorVersion) + "." + std::to_string(info.dwMinorVersion);
			}
			else
			{
				if (info.dwMajorVersion == 10)
				{
					switch (info.dwBuildNumber)
					{
						case 14393:
							return "Windows Server 2016";
						case 16299:
							return "Windows Server 1709";
						case 17134:
							return "Windows Server 1803";
						case 17763: 
							return "Windows Server 2019";
						case 18362:
							return "Windows Server 1903";
						case 18363:
							return "Windows Server 1909";
						case 19041:
							return "Windows Server 2004";
						case 19042:
							return "Windows Server 20H2";
						case 20348:
							return "Windows Server 2022";
						case 25398:
							return "Windows Server 23H2";
						case 26100:
							return "Windows Server 2025";
						default:
							return "Windows Server";
					}
				}
				else if (info.dwMajorVersion == 6)
				{
					switch (info.dwMinorVersion)
					{
						case 0: return "Windows Server 2008";
						case 1: return "Windows Server 2008 R2";
						case 2: return "Windows Server 2012";
						case 3: return "Windows Server 2012 R2";
					}
				}

				return "Windows Server " + std::to_string(info.dwMajorVersion) + "." + std::to_string(info.dwMinorVersion);
			}
		}
	}

	size_t SystemInfo::GetSystemRam()
	{
		return ReadSystemMemInfo(SYS_MEM_TYPE::TOTAL_PHYS);
	}

	size_t SystemInfo::GetSystemRamAvailable()
	{
		return ReadSystemMemInfo(SYS_MEM_TYPE::AVAIL_PHYS);
	}

	size_t SystemInfo::GetAppRamMax()
	{
		return std::min(GetSystemRam(), GetCommitLimit());
	}

	size_t SystemInfo::GetAppRamAvailable()
	{
		return std::min(GetSystemRamAvailable(), GetCommitLimit() - GetAppRamUsed());
	}

	size_t SystemInfo::GetAppRamUsed()
	{
		return ReadAppMemInfo(APP_MEM_TYPE::USED);
	}

	size_t SystemInfo::GetAppVirtualMemoryMax()
	{
		return ReadAppMemInfo(APP_MEM_TYPE::VM_MAX);
	}

	const std::string& SystemInfo::GetUserName()
	{
		static std::string userName;
		if (userName.empty())
		{
			char username[UNLEN + 1];
			DWORD username_len = UNLEN + 1;
			::GetUserNameA(username, &username_len);
			userName = username;
		}
		return userName;
	}
	
	const std::string& SystemInfo::GetHostName()
	{
		static std::string hostName;
		if (hostName.empty())
		{
			char hostname[UNLEN+1];
			gethostname(hostname, UNLEN+1);
			hostName = hostname;
		}
		return hostName;
	}
	
	const std::string& SystemInfo::GetDeviceName()
	{
		static std::string devName;
		if (devName.empty())
		{
			char computerName[UNLEN+1];
			DWORD computerName_len = UNLEN+1;
			GetComputerNameA(computerName, &computerName_len);
			devName = computerName;
		}
		return devName;
	}
	
	const std::string& SystemInfo::GetDeviceVendorName()
	{
		static std::string vendorName;
		if (vendorName.empty())
		{
			std::optional<std::string> res = GetWMIProperty("Manufacturer");
			vendorName = res ? *res : "Unknown";
		}
		return vendorName;
	}
	
	const std::string& SystemInfo::GetDeviceModelName()
	{ 
		static std::string deviceModelName;
		if (deviceModelName.empty())
		{
			std::optional<std::string> res = GetWMIProperty("Model");
			deviceModelName = res ? *res : "Unknown";
		}
		return deviceModelName;
	}
	
	const std::string& SystemInfo::GetOsName()
	{
		static const std::string osName = "Windows";
		return osName;
	}

	OsVersion SystemInfo::GetOsVersion()
	{
		static OsVersion osVersion = {};
		if (osVersion.major == 0 && osVersion.minor == 0)
		{
			NTSTATUS(WINAPI * RtlGetVersion)(LPOSVERSIONINFOEXW);
			OSVERSIONINFOEXW info;
			*(FARPROC*) &RtlGetVersion = GetProcAddress(GetModuleHandleA("ntdll"), "RtlGetVersion");
			info.dwOSVersionInfoSize = sizeof(info);
			RtlGetVersion(&info);
			osVersion = { (int)info.dwMajorVersion, (int)info.dwMinorVersion, 0, (int)info.dwBuildNumber };
		}
		return osVersion;
	}
	
	const std::string& SystemInfo::GetOsNameHumanReadable()
	{
		static const std::string osName = GetHumanReadableOSName();
		return osName;
	}
	
	DeviceType SystemInfo::GetDeviceType()
	{
		return DeviceType::PC;
	}

	size_t SystemInfo::GetCpuCoreCount()
	{
		static size_t coreCount = 0;
		if (coreCount == 0)
		{
			DWORD bufferSize = 0;
			GetLogicalProcessorInformationEx(LOGICAL_PROCESSOR_RELATIONSHIP::RelationProcessorCore, nullptr, &bufferSize);

			std::vector<BYTE> buf(bufferSize);
			SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX* info = reinterpret_cast<SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX*>(buf.data());
			GetLogicalProcessorInformationEx(LOGICAL_PROCESSOR_RELATIONSHIP::RelationProcessorCore, info, &bufferSize);
			
			size_t physProcessorCount = 0;
			BYTE* start = buf.data();
			BYTE* end = buf.data() + bufferSize;
			while (start < end)
			{
				SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX* current = reinterpret_cast<SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX*>(start);
				physProcessorCount++;
				start += current->Size;
			}
			coreCount = physProcessorCount;
		}
		return coreCount;
	}
	
	size_t SystemInfo::GetCpuThreadCount()
	{
		return std::thread::hardware_concurrency();
	}
	
	int32_t SystemInfo::GetCpuTemperature()
	{
		return 0; //TODO
	}
	
	CpuThermalState SystemInfo::GetCpuThermalState()
	{
		return CpuThermalState::Normal; // TODO
	}
	
	bool SystemInfo::IsDeviceInLowPowerMode()
	{
		GUID* activePowerMode;
		bool res = false;
		if(PowerGetActiveScheme(0, &activePowerMode) == 0)
		{
			res = IsEqualGUID(*activePowerMode, GUID_MAX_POWER_SAVINGS);
		}
		LocalFree(activePowerMode);
		return res;
	}
	
	BatteryState SystemInfo::GetDeviceBatteryState()
	{
		SYSTEM_POWER_STATUS info;
		GetSystemPowerStatus(&info);
		BYTE batteryFlag = info.BatteryFlag;
		if (batteryFlag & BatteryChargeStatus::CHARGING)
		{
			if (info.BatteryLifePercent == 100)
				return BatteryState::ChargingFull;
			return BatteryState::Charging;
		}
		else if (batteryFlag < BatteryChargeStatus::CHARGING)
		{
			return BatteryState::Unplugged;
		}
		// BatteryState::NotCharging is impossible to distinguish
		// NO_SYSTEM_BATTERY or UNKNOWN
		return BatteryState::Unavailable;
	}
	
	float SystemInfo::GetDeviceBatteryLevel()
	{
		SYSTEM_POWER_STATUS info;
		GetSystemPowerStatus(&info);
		return info.BatteryLifePercent != 255 ? (static_cast<float>(info.BatteryLifePercent / 100.f)) : -1.f;
	}
	
	void SystemInfo::EnableEnergyEvents()
	{
		//TODO
	}
	
	bool SystemInfo::IsMultitaskingSupported()
	{
		return true;
	}
	
	DeviceOrientation SystemInfo::GetDeviceOrientation()
	{
		DEVMODEA devmode;
		ZeroMemory(&devmode, sizeof(devmode));
		devmode.dmSize = sizeof(DEVMODE);
		if (EnumDisplaySettingsA(0, ENUM_CURRENT_SETTINGS, &devmode))
		{
			DWORD width = devmode.dmPelsWidth;
			DWORD height = devmode.dmPelsHeight;
			if (width > height)
			{
				// landscape
				return devmode.dmDisplayOrientation == DMDO_180 ? DeviceOrientation::LandscapeLeft :
				                                                  DeviceOrientation::LandscapeRight;
			}
			// portrait
			return devmode.dmDisplayOrientation == DMDO_180 ? DeviceOrientation::PortraitUpsideDown :
			                                                  DeviceOrientation::Portrait;
		}
		return DeviceOrientation::Unknown;
	}
	
	void SystemInfo::EnableDeviceOrientationEvents()
	{
		//TODO
	}
	
	InterfaceOrientation SystemInfo::GetInterfaceOrientation()
	{
		DeviceOrientation devOrient = GetDeviceOrientation();
		switch (devOrient)
		{
			case OpenVulkano::DeviceOrientation::Portrait:
				return InterfaceOrientation::Portrait;
			case OpenVulkano::DeviceOrientation::PortraitUpsideDown:
				return InterfaceOrientation::PortraitUpsideDown;
			case OpenVulkano::DeviceOrientation::LandscapeLeft:
				return InterfaceOrientation::LandscapeRight;
			case OpenVulkano::DeviceOrientation::LandscapeRight:
				return InterfaceOrientation::LandscapeLeft;
			default:
				return InterfaceOrientation::Landscape;
		}
	}
}