// Copyright Epic Games, Inc. All Rights Reserved.

#include "zencore/system.h"

#include <zencore/compactbinarybuilder.h>
#include <zencore/except.h>
#include <zencore/memory/memory.h>
#include <zencore/string.h>

#if ZEN_PLATFORM_WINDOWS
#	include <zencore/windows.h>

ZEN_THIRD_PARTY_INCLUDES_START
#	include <iphlpapi.h>
#	include <winsock2.h>
#	include <pdh.h>
#	pragma comment(lib, "pdh.lib")
ZEN_THIRD_PARTY_INCLUDES_END
#elif ZEN_PLATFORM_LINUX
#	include <sys/utsname.h>
#	include <unistd.h>
#elif ZEN_PLATFORM_MAC
#	include <unistd.h>
#	include <mach/mach.h>
#	include <mach/mach_host.h>
#	include <sys/types.h>
#	include <sys/sysctl.h>
#endif

namespace zen {

using namespace std::literals;

int g_FakeCpuCount = 0;

void
SetCpuCountForReporting(int FakeCpuCount)
{
	g_FakeCpuCount = FakeCpuCount;
}

#if ZEN_PLATFORM_WINDOWS
std::string
GetMachineName()
{
	WCHAR NameBuffer[256];
	DWORD dwNameSize = 255;
	BOOL  Success	 = GetComputerNameW(NameBuffer, &dwNameSize);

	if (Success)
	{
		return WideToUtf8(NameBuffer);
	}

	return {};
}

SystemMetrics
GetSystemMetrics()
{
	SYSTEM_INFO SysInfo{};
	GetSystemInfo(&SysInfo);

	SystemMetrics Metrics;

	Metrics.LogicalProcessorCount = SysInfo.dwNumberOfProcessors;

	// Determine physical core count

	{
		DWORD BufferSize = 0;
		BOOL  Result	 = GetLogicalProcessorInformationEx(RelationAll, nullptr, &BufferSize);
		if (int32_t Error = GetLastError(); Error != ERROR_INSUFFICIENT_BUFFER)
		{
			ThrowSystemError(Error, "Failed to get buffer size for logical processor information");
		}

		PSYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX Buffer = (PSYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX)Memory::Alloc(BufferSize);

		Result = GetLogicalProcessorInformationEx(RelationAll, Buffer, &BufferSize);
		if (!Result)
		{
			Memory::Free(Buffer);
			throw std::runtime_error("Failed to get logical processor information");
		}

		DWORD ProcessorPkgCount		= 0;
		DWORD ProcessorCoreCount	= 0;
		DWORD LogicalProcessorCount = 0;

		BYTE*		Ptr = reinterpret_cast<BYTE*>(Buffer);
		BYTE* const End = Ptr + BufferSize;
		while (Ptr < End)
		{
			const SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX& Slpi = *reinterpret_cast<const SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX*>(Ptr);
			if (Slpi.Relationship == RelationProcessorCore)
			{
				++ProcessorCoreCount;

				// Count logical processors (threads) across all processor groups for this core.
				// Each core entry lists one GROUP_AFFINITY per group it spans; each set bit
				// in the Mask represents one logical processor (HyperThreading sibling).
				for (WORD g = 0; g < Slpi.Processor.GroupCount; ++g)
				{
					LogicalProcessorCount += static_cast<DWORD>(__popcnt64(Slpi.Processor.GroupMask[g].Mask));
				}
			}
			else if (Slpi.Relationship == RelationProcessorPackage)
			{
				++ProcessorPkgCount;
			}
			Ptr += Slpi.Size;
		}

		Metrics.CoreCount			  = ProcessorCoreCount;
		Metrics.CpuCount			  = ProcessorPkgCount;
		Metrics.LogicalProcessorCount = LogicalProcessorCount;

		Memory::Free(Buffer);
	}

	// Query memory status

	{
		MEMORYSTATUSEX MemStatus{.dwLength = sizeof(MEMORYSTATUSEX)};
		GlobalMemoryStatusEx(&MemStatus);

		Metrics.SystemMemoryMiB		  = MemStatus.ullTotalPhys / 1024 / 1024;
		Metrics.AvailSystemMemoryMiB  = MemStatus.ullAvailPhys / 1024 / 1024;
		Metrics.VirtualMemoryMiB	  = MemStatus.ullTotalVirtual / 1024 / 1024;
		Metrics.AvailVirtualMemoryMiB = MemStatus.ullAvailVirtual / 1024 / 1024;
		Metrics.PageFileMiB			  = MemStatus.ullTotalPageFile / 1024 / 1024;
		Metrics.AvailPageFileMiB	  = MemStatus.ullAvailPageFile / 1024 / 1024;
	}

	// Query CPU usage using PDH
	//
	// TODO: This should be changed to not require a Sleep, perhaps by using some
	// background metrics gathering mechanism.

	{
		PDH_HQUERY	 QueryHandle   = nullptr;
		PDH_HCOUNTER CounterHandle = nullptr;

		if (PdhOpenQueryW(nullptr, 0, &QueryHandle) == ERROR_SUCCESS)
		{
			if (PdhAddEnglishCounterW(QueryHandle, L"\\Processor(_Total)\\% Processor Time", 0, &CounterHandle) == ERROR_SUCCESS)
			{
				PdhCollectQueryData(QueryHandle);
				Sleep(100);
				PdhCollectQueryData(QueryHandle);

				PDH_FMT_COUNTERVALUE CounterValue;
				if (PdhGetFormattedCounterValue(CounterHandle, PDH_FMT_DOUBLE, nullptr, &CounterValue) == ERROR_SUCCESS)
				{
					Metrics.CpuUsagePercent = static_cast<float>(CounterValue.doubleValue);
				}
			}
			PdhCloseQuery(QueryHandle);
		}
	}

	return Metrics;
}
#elif ZEN_PLATFORM_LINUX
std::string
GetMachineName()
{
	static char Result[256] = "";
	if (!Result[0])
	{
		struct utsname name;
		const char*	   SysName = name.nodename;
		if (uname(&name))
		{
			SysName = "Unix Computer";
		}

		strcpy(Result, SysName);
	}
	return Result;
}

SystemMetrics
GetSystemMetrics()
{
	SystemMetrics Metrics;

	// Get processor information
	long NumProcessors = sysconf(_SC_NPROCESSORS_ONLN);
	if (NumProcessors > 0)
	{
		Metrics.LogicalProcessorCount = static_cast<uint32_t>(NumProcessors);
	}

	// On Linux, we approximate core count from logical processor count
	// A more accurate implementation would parse /proc/cpuinfo
	Metrics.CoreCount = Metrics.LogicalProcessorCount;
	Metrics.CpuCount  = 1;	// Default to 1 CPU package

	// Parse /proc/cpuinfo for more accurate core/CPU info
	if (FILE* CpuInfo = fopen("/proc/cpuinfo", "r"))
	{
		char Line[1024];
		int	 PhysicalIds = 0;
		int	 CoreIds	 = 0;

		while (fgets(Line, sizeof(Line), CpuInfo))
		{
			if (strncmp(Line, "physical id", 11) == 0)
			{
				int Id;
				if (sscanf(Line, "physical id\t: %d", &Id) == 1)
				{
					if (Id + 1 > PhysicalIds)
					{
						PhysicalIds = Id + 1;
					}
				}
			}
			else if (strncmp(Line, "cpu cores", 9) == 0)
			{
				sscanf(Line, "cpu cores\t: %d", &CoreIds);
			}
		}
		fclose(CpuInfo);

		if (PhysicalIds > 0)
		{
			Metrics.CpuCount = PhysicalIds;
		}
		if (CoreIds > 0)
		{
			Metrics.CoreCount = CoreIds * Metrics.CpuCount;
		}
	}

	// Query CPU usage
	Metrics.CpuUsagePercent = 0.0f;
	if (FILE* Stat = fopen("/proc/stat", "r"))
	{
		char				 Line[256];
		unsigned long		 User, Nice, System, Idle, IoWait, Irq, SoftIrq;
		static unsigned long PrevUser = 0, PrevNice = 0, PrevSystem = 0, PrevIdle = 0, PrevIoWait = 0, PrevIrq = 0, PrevSoftIrq = 0;

		if (fgets(Line, sizeof(Line), Stat))
		{
			if (sscanf(Line, "cpu %lu %lu %lu %lu %lu %lu %lu", &User, &Nice, &System, &Idle, &IoWait, &Irq, &SoftIrq) == 7)
			{
				unsigned long TotalDelta = (User + Nice + System + Idle + IoWait + Irq + SoftIrq) -
										   (PrevUser + PrevNice + PrevSystem + PrevIdle + PrevIoWait + PrevIrq + PrevSoftIrq);
				unsigned long IdleDelta = Idle - PrevIdle;

				if (TotalDelta > 0)
				{
					Metrics.CpuUsagePercent = 100.0f * (TotalDelta - IdleDelta) / TotalDelta;
				}

				PrevUser	= User;
				PrevNice	= Nice;
				PrevSystem	= System;
				PrevIdle	= Idle;
				PrevIoWait	= IoWait;
				PrevIrq		= Irq;
				PrevSoftIrq = SoftIrq;
			}
		}
		fclose(Stat);
	}

	// Get memory information
	long Pages		= sysconf(_SC_PHYS_PAGES);
	long PageSize	= sysconf(_SC_PAGE_SIZE);
	long AvailPages = sysconf(_SC_AVPHYS_PAGES);

	if (Pages > 0 && PageSize > 0)
	{
		Metrics.SystemMemoryMiB		 = (Pages * PageSize) / 1024 / 1024;
		Metrics.AvailSystemMemoryMiB = (AvailPages * PageSize) / 1024 / 1024;
	}

	// Virtual memory is not directly available via sysconf
	// Set to system memory as a reasonable default
	Metrics.VirtualMemoryMiB	  = Metrics.SystemMemoryMiB;
	Metrics.AvailVirtualMemoryMiB = Metrics.AvailSystemMemoryMiB;

	// Parse /proc/meminfo for swap/page file information
	Metrics.PageFileMiB		 = 0;
	Metrics.AvailPageFileMiB = 0;

	if (FILE* MemInfo = fopen("/proc/meminfo", "r"))
	{
		char Line[256];
		long SwapTotal = 0;
		long SwapFree  = 0;

		while (fgets(Line, sizeof(Line), MemInfo))
		{
			if (strncmp(Line, "SwapTotal:", 10) == 0)
			{
				sscanf(Line, "SwapTotal: %ld kB", &SwapTotal);
			}
			else if (strncmp(Line, "SwapFree:", 9) == 0)
			{
				sscanf(Line, "SwapFree: %ld kB", &SwapFree);
			}
		}
		fclose(MemInfo);

		if (SwapTotal > 0)
		{
			Metrics.PageFileMiB		 = SwapTotal / 1024;
			Metrics.AvailPageFileMiB = SwapFree / 1024;
		}
	}

	return Metrics;
}
#elif ZEN_PLATFORM_MAC
std::string
GetMachineName()
{
	static char Result[256] = "";

	if (!Result[0])
	{
		gethostname(Result, sizeof(Result));
	}
	return Result;
}

SystemMetrics
GetSystemMetrics()
{
	SystemMetrics Metrics;

	// Get processor information
	size_t Size = sizeof(Metrics.LogicalProcessorCount);
	sysctlbyname("hw.logicalcpu", &Metrics.LogicalProcessorCount, &Size, nullptr, 0);

	uint32_t PhysicalCpu = 0;
	Size				 = sizeof(PhysicalCpu);
	sysctlbyname("hw.physicalcpu", &PhysicalCpu, &Size, nullptr, 0);
	Metrics.CoreCount = PhysicalCpu;

	uint32_t Packages = 0;
	Size			  = sizeof(Packages);
	sysctlbyname("hw.packages", &Packages, &Size, nullptr, 0);
	Metrics.CpuCount = Packages > 0 ? Packages : 1;

	// Query CPU usage using host_statistics64
	Metrics.CpuUsagePercent = 0.0f;
	host_cpu_load_info_data_t CpuLoad;
	mach_msg_type_number_t	  CpuCount = sizeof(CpuLoad) / sizeof(natural_t);
	if (host_statistics(mach_host_self(), HOST_CPU_LOAD_INFO, (host_info_t)&CpuLoad, &CpuCount) == KERN_SUCCESS)
	{
		unsigned long TotalTicks = 0;
		for (int i = 0; i < CPU_STATE_MAX; ++i)
		{
			TotalTicks += CpuLoad.cpu_ticks[i];
		}

		if (TotalTicks > 0)
		{
			unsigned long IdleTicks = CpuLoad.cpu_ticks[CPU_STATE_IDLE];
			Metrics.CpuUsagePercent = 100.0f * (TotalTicks - IdleTicks) / TotalTicks;
		}
	}

	// Get memory information
	uint64_t MemSize = 0;
	Size			 = sizeof(MemSize);
	sysctlbyname("hw.memsize", &MemSize, &Size, nullptr, 0);
	Metrics.SystemMemoryMiB = MemSize / 1024 / 1024;

	// Get available memory using vm_stat
	vm_size_t PageSize = 0;
	host_page_size(mach_host_self(), &PageSize);

	vm_statistics64_data_t VmStats;
	mach_msg_type_number_t InfoCount = sizeof(VmStats) / sizeof(natural_t);
	host_statistics64(mach_host_self(), HOST_VM_INFO64, (host_info64_t)&VmStats, &InfoCount);

	uint64_t FreeMemory			 = (uint64_t)(VmStats.free_count + VmStats.inactive_count) * PageSize;
	Metrics.AvailSystemMemoryMiB = FreeMemory / 1024 / 1024;

	// Virtual memory on macOS is essentially unlimited, set to a large value
	Metrics.VirtualMemoryMiB	  = Metrics.SystemMemoryMiB * 16;
	Metrics.AvailVirtualMemoryMiB = Metrics.VirtualMemoryMiB;

	// macOS doesn't use traditional page files, use swap space
	xsw_usage SwapUsage;
	Size = sizeof(SwapUsage);
	sysctlbyname("vm.swapusage", &SwapUsage, &Size, nullptr, 0);
	Metrics.PageFileMiB		 = SwapUsage.xsu_total / 1024 / 1024;
	Metrics.AvailPageFileMiB = (SwapUsage.xsu_total - SwapUsage.xsu_used) / 1024 / 1024;

	return Metrics;
}
#else
#	error "Unknown platform"
#endif

SystemMetrics
GetSystemMetricsForReporting()
{
	SystemMetrics Sm = GetSystemMetrics();

	if (g_FakeCpuCount)
	{
		Sm.CoreCount			 = g_FakeCpuCount;
		Sm.LogicalProcessorCount = g_FakeCpuCount;
	}

	return Sm;
}

std::string_view
GetOperatingSystemName()
{
	return ZEN_PLATFORM_NAME;
}

std::string_view
GetCpuName()
{
#if ZEN_ARCH_X64
	return "x64"sv;
#elif ZEN_ARCH_ARM64
	return "arm64"sv;
#endif
}

void
Describe(const SystemMetrics& Metrics, CbWriter& Writer)
{
	Writer << "cpu_count" << Metrics.CpuCount << "core_count" << Metrics.CoreCount << "lp_count" << Metrics.LogicalProcessorCount
		   << "total_memory_mb" << Metrics.SystemMemoryMiB << "avail_memory_mb" << Metrics.AvailSystemMemoryMiB << "total_virtual_mb"
		   << Metrics.VirtualMemoryMiB << "avail_virtual_mb" << Metrics.AvailVirtualMemoryMiB << "total_pagefile_mb" << Metrics.PageFileMiB
		   << "avail_pagefile_mb" << Metrics.AvailPageFileMiB;
}

}  // namespace zen