#include "../Include/OSDef.h"

#include "../Include/Ring0Wrapper.h"
#include "imc.h"
#include <time.h>

#ifdef __WIN_OS__
#include <windows.h>
typedef BOOL (WINAPI *PGNSI)(PULONGLONG);
#else
#include <sys/sysinfo.h>
#include <sys/mman.h>
#endif
#include <stdio.h>

#define CHECK(X) if (!(X)) return false;

bool MSREnabled()
{
	unsigned int EDXBuf = 0;

	__asm
	{
		mov		eax, 1			// Get Version & Feature Info.
		cpuid
		mov		EDXBuf, edx		// Feature Info.
	}

	bool MSR	= ((EDXBuf & 0x00000020) != 0);
	return MSR;
}


bool GetIMCDetectedSpeed(unsigned long long *result)
{
	if (MSREnabled())
	{
		//MSR Registers and Addresses For IMC
		//MSR_UNCORE_PERF_GLOBAL_CTRL 391H 913
		//MSR_UNCORE_FIXED_CTR_CTRL 395H 917
		//MSR_UNCORE_FIXED_CTR0 394H 916

		unsigned long long MSR_UNCORE_PERF_GLOBAL_CTRL
			, MSR_UNCORE_FIXED_CTR_CTRL
			, MSR_UNCORE_FIXED_CTR0;


		unsigned long long ORI_MSR_UNCORE_PERF_GLOBAL_CTRL
			, ORI_MSR_UNCORE_FIXED_CTR_CTRL, ORI_MSR_UNCORE_FIXED_CTR0;
		//Reset counter
		CHECK(RDMSR(913, &MSR_UNCORE_PERF_GLOBAL_CTRL));
		ORI_MSR_UNCORE_PERF_GLOBAL_CTRL = MSR_UNCORE_PERF_GLOBAL_CTRL;
		MSR_UNCORE_PERF_GLOBAL_CTRL = MSR_UNCORE_PERF_GLOBAL_CTRL | 0x100000000;
		CHECK(WRMSR(913, MSR_UNCORE_PERF_GLOBAL_CTRL));

		CHECK(RDMSR(917, &MSR_UNCORE_FIXED_CTR_CTRL));
		ORI_MSR_UNCORE_FIXED_CTR_CTRL = MSR_UNCORE_FIXED_CTR_CTRL;
		MSR_UNCORE_FIXED_CTR_CTRL = MSR_UNCORE_FIXED_CTR_CTRL | 0x1;
		CHECK(RDMSR(916, &ORI_MSR_UNCORE_FIXED_CTR0));
		CHECK(WRMSR(916,0));
		CHECK(WRMSR(917, MSR_UNCORE_FIXED_CTR_CTRL));

		wait(3);

		CHECK(WRMSR(917,ORI_MSR_UNCORE_FIXED_CTR_CTRL));

		CHECK(WRMSR(913, ORI_MSR_UNCORE_PERF_GLOBAL_CTRL));

		CHECK(RDMSR(916, &MSR_UNCORE_FIXED_CTR0));

		CHECK(WRMSR(916, ORI_MSR_UNCORE_FIXED_CTR0));
		MSR_UNCORE_FIXED_CTR0 &= 0x000000FFFFFFFFFF;
		*result = (unsigned long long)MSR_UNCORE_FIXED_CTR0/3;

		return true;
	}
	return false;
}

//bool GetCurrentRatio(unsigned int &result)
//{
//	unsigned long long IA32_PERF_GLOBAL_CTRL, IA32_FIXED_CTR_CTL, CPU_CLK_UNHALTED_CORE, CPU_CLK_UNHALTED_REF, Max_Ratio;
//	unsigned long long ORI_IA32_PERF_GLOBAL_CTRL, ORI_IA32_FIXED_CTR_CTL, ORI_CPU_CLK_UNHALTED_CORE, ORI_CPU_CLK_UNHALTED_REF;
//	int currentRatio = 0;
//	double currentRatioD = 0;
//
//	CHECK(RDMSR(0x30A, &CPU_CLK_UNHALTED_CORE));
//	ORI_CPU_CLK_UNHALTED_CORE = CPU_CLK_UNHALTED_CORE;
//	CHECK(WRMSR(0x30A, 0));
//
//	CHECK(RDMSR(0x30B, &CPU_CLK_UNHALTED_REF));
//	ORI_CPU_CLK_UNHALTED_REF = CPU_CLK_UNHALTED_REF;
//	CHECK(WRMSR(0x30B, 0));
//
//	CHECK(RDMSR(0x38F, &IA32_PERF_GLOBAL_CTRL));
//	ORI_IA32_PERF_GLOBAL_CTRL = IA32_PERF_GLOBAL_CTRL;
//	IA32_PERF_GLOBAL_CTRL = IA32_PERF_GLOBAL_CTRL | 0x600000000;
//	CHECK(WRMSR(0x38F, IA32_PERF_GLOBAL_CTRL));
//
//	CHECK(RDMSR(0x38D, &IA32_FIXED_CTR_CTL));
//	ORI_IA32_FIXED_CTR_CTL = IA32_FIXED_CTR_CTL;
//	IA32_FIXED_CTR_CTL = IA32_FIXED_CTR_CTL | 0x330;
//	CHECK(WRMSR(0x38D, IA32_FIXED_CTR_CTL));
//
//	Wait(1);
//
//	
//	CHECK(WRMSR(0x38F, ORI_IA32_PERF_GLOBAL_CTRL));
//
//	CHECK(WRMSR(0x38D, ORI_IA32_FIXED_CTR_CTL));
//
//	CHECK(RDMSR(0x30A, &CPU_CLK_UNHALTED_CORE));
//
//	CHECK(RDMSR(0x30B, &CPU_CLK_UNHALTED_REF));
//
//	CHECK(WRMSR(0x30A, ORI_CPU_CLK_UNHALTED_CORE));
//
//	CHECK(WRMSR(0x30B, ORI_CPU_CLK_UNHALTED_REF));
//
//	CHECK(RDMSR(0xCE, &Max_Ratio));
//
//	Max_Ratio = Max_Ratio >> 8;
//	Max_Ratio = Max_Ratio & 0xFF;
//
//	currentRatioD =  (float)Max_Ratio * (float)CPU_CLK_UNHALTED_CORE / (float)CPU_CLK_UNHALTED_REF;
//	currentRatio = (int)currentRatioD;
//	if (currentRatioD - currentRatio >= 0.5)
//	{
//		result = currentRatio + 1;
//	}
//	else
//	{	
//		result = currentRatio;
//	}
//	return true;
//}

bool GetIMCExpectedSpeed(unsigned long long *result)
{
	unsigned long IMCRatio = 0;
	unsigned short brandData = 0;
	ReadPCI16(0x3F, 0, 0, 0x00, &brandData);
	if (brandData == 0x8086)
	{
		ReadPCI32(0x3F, 0, 0, 0xC0, &IMCRatio);
	}
	else
	{
		ReadPCI16(0xFF, 0, 0, 0x00, &brandData);
		if (brandData == 0x8086)
		{
			ReadPCI32(0xFF, 0, 0, 0xC0, &IMCRatio);
		}
		else
		{
			return false;
		}
	}
	IMCRatio = IMCRatio & 0xFF;
	//printf("IMCRatio %x\r\n", IMCRatio);
	*result = (unsigned long long)(IMCRatio * 133333333);
	return true;
}

extern bool GetMemorySize(long long *result)
{
#ifdef __WIN_OS__
	OSVERSIONINFOEX osvi;
    BOOL bIsWindowsXPorLater;

    ZeroMemory(&osvi, sizeof(OSVERSIONINFOEX));
    osvi.dwOSVersionInfoSize = sizeof(OSVERSIONINFOEX);

	if (!GetVersionEx((LPOSVERSIONINFO)&osvi)) 
	{
		result = 0;
		return false;
	}

	if ((osvi.dwMajorVersion == 6 && osvi.wServicePackMajor > 0) || (osvi.dwMajorVersion >= 6 && osvi.dwMinorVersion > 0))
	{
		ULONGLONG totalMemoryInKilobytes = 0;
		PGNSI pGNSI;

		pGNSI = (PGNSI) GetProcAddress(
			GetModuleHandle(TEXT("Kernel32.dll")), 
			"GetPhysicallyInstalledSystemMemory");
		if(NULL != pGNSI)
		{
			if (pGNSI(&totalMemoryInKilobytes))
			{
				*result = totalMemoryInKilobytes * 1024;
				return true;
			}
			else 
			{
				return false;
			}
		}		
		else 
		{
			result = 0;
			return false;
		}
	}
	else
	{
		MEMORYSTATUSEX statex;

		statex.dwLength = sizeof (statex);

		if(GlobalMemoryStatusEx (&statex))
		{
			*result = (long long)statex.ullTotalPhys;
			return true;
		}
		else
		{
			result = 0;
			return false;
		}
	}
#else
	struct sysinfo memInfo;
	if (sysinfo (&memInfo) == 0)
	{
		long long totalPhysMem = memInfo.totalram;
		//Multiply in next statement to avoid int overflow on right hand side...
		totalPhysMem *= memInfo.mem_unit;
		*result = totalPhysMem;
		return true;
	}
	return false;
#endif
}

