// 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>

#include <mutex>

#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;
	}

	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;
		}
	}

	// 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;

	// 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

ExtendedSystemMetrics
ApplyReportingOverrides(ExtendedSystemMetrics Metrics)
{
	if (g_FakeCpuCount)
	{
		Metrics.CoreCount			  = g_FakeCpuCount;
		Metrics.LogicalProcessorCount = g_FakeCpuCount;
	}
	return Metrics;
}

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

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

	return Sm;
}

///////////////////////////////////////////////////////////////////////////
// SystemMetricsTracker
///////////////////////////////////////////////////////////////////////////

// Per-platform CPU sampling helper. Called with m_Mutex held.

#if ZEN_PLATFORM_WINDOWS || ZEN_PLATFORM_LINUX

// Samples CPU usage by reading /proc/stat. Used natively on Linux and as a
// Wine fallback on Windows (where /proc/stat is accessible via the Z: drive).
struct ProcStatCpuSampler
{
	const char*	  Path		  = "/proc/stat";
	unsigned long PrevUser	  = 0;
	unsigned long PrevNice	  = 0;
	unsigned long PrevSystem  = 0;
	unsigned long PrevIdle	  = 0;
	unsigned long PrevIoWait  = 0;
	unsigned long PrevIrq	  = 0;
	unsigned long PrevSoftIrq = 0;

	explicit ProcStatCpuSampler(const char* InPath = "/proc/stat") : Path(InPath) {}

	float Sample()
	{
		float CpuUsage = 0.0f;

		if (FILE* Stat = fopen(Path, "r"))
		{
			char		  Line[256];
			unsigned long User, Nice, System, Idle, IoWait, Irq, SoftIrq;

			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)
					{
						CpuUsage = 100.0f * (TotalDelta - IdleDelta) / TotalDelta;
					}

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

		return CpuUsage;
	}
};

#endif

#if ZEN_PLATFORM_WINDOWS

struct CpuSampler
{
	PDH_HQUERY		   QueryHandle		 = nullptr;
	PDH_HCOUNTER	   CounterHandle	 = nullptr;
	bool			   HasPreviousSample = false;
	bool			   IsWine			 = false;
	ProcStatCpuSampler ProcStat{"Z:\\proc\\stat"};

	CpuSampler()
	{
		IsWine = zen::windows::IsRunningOnWine();

		if (!IsWine)
		{
			if (PdhOpenQueryW(nullptr, 0, &QueryHandle) == ERROR_SUCCESS)
			{
				if (PdhAddEnglishCounterW(QueryHandle, L"\\Processor(_Total)\\% Processor Time", 0, &CounterHandle) != ERROR_SUCCESS)
				{
					CounterHandle = nullptr;
				}
			}
		}
	}

	~CpuSampler()
	{
		if (QueryHandle)
		{
			PdhCloseQuery(QueryHandle);
		}
	}

	float Sample()
	{
		if (IsWine)
		{
			return ProcStat.Sample();
		}

		if (!QueryHandle || !CounterHandle)
		{
			return 0.0f;
		}

		PdhCollectQueryData(QueryHandle);

		if (!HasPreviousSample)
		{
			HasPreviousSample = true;
			return 0.0f;
		}

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

		return 0.0f;
	}
};

#elif ZEN_PLATFORM_LINUX

struct CpuSampler
{
	ProcStatCpuSampler ProcStat;

	float Sample() { return ProcStat.Sample(); }
};

#elif ZEN_PLATFORM_MAC

struct CpuSampler
{
	unsigned long PrevTotalTicks = 0;
	unsigned long PrevIdleTicks	 = 0;

	float Sample()
	{
		float CpuUsage = 0.0f;

		host_cpu_load_info_data_t CpuLoad;
		mach_msg_type_number_t	  Count = sizeof(CpuLoad) / sizeof(natural_t);
		if (host_statistics(mach_host_self(), HOST_CPU_LOAD_INFO, (host_info_t)&CpuLoad, &Count) == KERN_SUCCESS)
		{
			unsigned long TotalTicks = 0;
			for (int i = 0; i < CPU_STATE_MAX; ++i)
			{
				TotalTicks += CpuLoad.cpu_ticks[i];
			}
			unsigned long IdleTicks = CpuLoad.cpu_ticks[CPU_STATE_IDLE];

			unsigned long TotalDelta = TotalTicks - PrevTotalTicks;
			unsigned long IdleDelta	 = IdleTicks - PrevIdleTicks;

			if (TotalDelta > 0 && PrevTotalTicks > 0)
			{
				CpuUsage = 100.0f * (TotalDelta - IdleDelta) / TotalDelta;
			}

			PrevTotalTicks = TotalTicks;
			PrevIdleTicks  = IdleTicks;
		}

		return CpuUsage;
	}
};

#endif

struct SystemMetricsTracker::Impl
{
	using Clock = std::chrono::steady_clock;

	std::mutex				  Mutex;
	CpuSampler				  Sampler;
	float					  CachedCpuPercent = 0.0f;
	Clock::time_point		  NextSampleTime   = Clock::now();
	std::chrono::milliseconds MinInterval;

	explicit Impl(std::chrono::milliseconds InMinInterval) : MinInterval(InMinInterval) {}

	float SampleCpu()
	{
		const auto Now = Clock::now();
		if (Now >= NextSampleTime)
		{
			CachedCpuPercent = Sampler.Sample();
			NextSampleTime	 = Now + MinInterval;
		}
		return CachedCpuPercent;
	}
};

SystemMetricsTracker::SystemMetricsTracker(std::chrono::milliseconds MinInterval) : m_Impl(std::make_unique<Impl>(MinInterval))
{
}

SystemMetricsTracker::~SystemMetricsTracker() = default;

ExtendedSystemMetrics
SystemMetricsTracker::Query()
{
	ExtendedSystemMetrics Metrics;
	static_cast<SystemMetrics&>(Metrics) = GetSystemMetrics();

	std::lock_guard Lock(m_Impl->Mutex);
	Metrics.CpuUsagePercent = m_Impl->SampleCpu();
	return Metrics;
}

///////////////////////////////////////////////////////////////////////////

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;
}

void
Describe(const ExtendedSystemMetrics& Metrics, CbWriter& Writer)
{
	Describe(static_cast<const SystemMetrics&>(Metrics), Writer);
	Writer << "cpu_usage_percent" << Metrics.CpuUsagePercent;
}

}  // namespace zen