// KomCoreDemo.cpp : A simple application demonstrating the use of
// KomCore IOCTLs.
//

#include <stdio.h>
#include <string.h>
#include <intrin.h>
#include <conio.h>

#include <stdlib.h>
#include <windows.h>
#include <winioctl.h>
#include <KomCoreIoctl.h>


void CloseIfOpen(void);
void ReadAllCores(bool continuous);

// Handle to device opened in driver.
//
HANDLE	hDevice = INVALID_HANDLE_VALUE;

// Name used to open device
//
wchar_t *sLinkName = L"\\\\.\\KomCoreDevice0";

int procs=0;	//numer of logical processors in the system

////////////////////////////////////////////////////////////////////////
// OpenByName
//
//		Open a handle to the requested device
//
HANDLE OpenByName(void)
{	
	// Create a handle to the driver
	return CreateFile(sLinkName,
					  GENERIC_READ | GENERIC_WRITE,
					  FILE_SHARE_READ,
					  NULL,
					  OPEN_EXISTING,
					  0,
					  NULL);
}

////////////////////////////////////////////////////////////////////////
// Exit
//
//		Print a message and exit
//
   void Exit(int res)
{
	printf("Exiting...\n\n");
	CloseIfOpen();
	exit(res);
}


////////////////////////////////////////////////////////////////////////
// Main entry point
//
//

int main(int argc, char* argv[])
{	
	DWORD procAff,sysAff;	//affinity	
	DWORD dw;
    int CPUInfo[4] = {-1};
    int nSteppingID = 0;
    int nModel = 0;
    int nFamily = 0;

	printf_s("KomCoreDemo - KomCore driver demo application\n");
	printf_s(" USAGE: KomCoreDemo [n]\n  n - logical processor number to read core temperature (zero-based)\n");
	printf_s("  if n='c' - read for all cores continuously\n");
	printf_s("  if n is not specified - read for all cores once\n\n");

    __cpuid(CPUInfo, 1);

    // Interpret CPU feature information.
    nSteppingID = CPUInfo[0] & 0xf;
    nModel = (CPUInfo[0] >> 4) & 0xf;
    nFamily = (CPUInfo[0] >> 8) & 0xf;

    if  (nSteppingID)
        printf_s("Stepping ID = %x\n", nSteppingID);
    if  (nModel)
        printf_s("Model = %x\n", nModel);
    if  (nFamily)
        printf_s("Family = %x\n", nFamily);

	//find out system affinity
	GetProcessAffinityMask(GetCurrentProcess(),(PDWORD_PTR)&procAff,(PDWORD_PTR)&sysAff);
	
	//calculating number of logical processors
	dw=sysAff;
	while (dw)
	{
		dw>>=1;
		procs++;
	}

	printf_s("Logical processors in the system: %d\n",procs);

	hDevice = OpenByName();
	if (hDevice == INVALID_HANDLE_VALUE)
	{
		printf_s("ERROR opening device: (%0x) returned from CreateFile\n", GetLastError());
		Exit(1);
	}
	else
	{
		printf_s("KomCore device found, handle open.\n");
	}

    //parsing command line
    if (argc==1)
	{
		ReadAllCores(false);
	}
	else if (argc==2)
	{
		if ((argv[1][0]=='c') || (argv[1][0]=='C'))
			ReadAllCores(true);
		else if ((argv[1][0]>='0') && (argv[1][0]<='9'))
		{	int proc=(int)argv[1][0]-'0';

			//reading specified processor:
			if (proc>=procs)
				printf_s("Processor %d not present in the system. Max processor index is %d.\n",proc,procs-1);
			else
			{	KomCoreInfo info;
				ULONG	nOutput;	// Count written to bufOutput
				int temperature;

				procAff=1<<proc;
				SetThreadAffinityMask(GetCurrentThread(),procAff);
				Sleep(100);

				// Call device IO Control interface (IOCTL_KOMCORE_GET_INFO) in driver
				printf_s("Issuing IOCTL_KOMCORE_GET_INFO to device - ");
				if (!DeviceIoControl(hDevice,IOCTL_KOMCORE_GET_INFO,
						 NULL,0,
						 &info,sizeof(info),
						 &nOutput,NULL))
				{
					printf_s("ERROR: DeviceIoControl returns %0x.", GetLastError());
					Exit(1);
				}

				printf_s(" maxCores: %d, max logical processors per core: %d,\n",info.maxCores,info.maxLPC);
				printf_s(" current core ID: %d, current core temperature resolution: %d C.\n",info.currCoreId,info.currCoreTempResol);


				// Call device IO Control interface (IOCTL_KOMCORE_GET_TEMP) in driver
				printf_s("Issuing IOCTL_KOMCORE_GET_TEMP to device - ");
				if (!DeviceIoControl(hDevice,IOCTL_KOMCORE_GET_TEMP,
						 NULL,0,
						 &temperature,sizeof(temperature),
						 &nOutput,NULL))
				{
					printf_s("ERROR: DeviceIoControl returns %0x.", GetLastError());
					Exit(1);
				}

				if (temperature==TEMP_INVALID)
					printf_s(" TEMPERATURE: INVALID VALUE!!\n\n");
				else
					printf_s(" TEMPERATURE: %d C\n\n",temperature);
			}
		}
		else
			printf_s("Command line error, see Usage...\n");
	}
	else //argc>2
	{
		printf_s("Command line error, see Usage...\n");
	}

	CloseIfOpen();

	return 0;
}


////////////////////////////////////////////////////////////////////////
// CloseIfOpen
//
//		Close the device if we previously opened a handle to it.
//
void CloseIfOpen(void)
{
	if (hDevice != INVALID_HANDLE_VALUE)
	{
		// Close the handle to the driver
		if (!CloseHandle(hDevice))
		{
			printf("ERROR: CloseHandle returns %0x.\n", GetLastError());
		}
		hDevice = INVALID_HANDLE_VALUE;
	}
}


void ReadAllCores(bool continuous)
{	KomCoreInfo info;
	ULONG	nOutput;	// Count written to bufOutput
	int i;
	int cores=0;	//number of idtentified cores
	UCHAR *coreIDs,*procIdx;	//tables for coreID and processor index (first found) for each core

	//setting up tables at max size: procs
	coreIDs=new UCHAR[procs];
	procIdx=new UCHAR[procs];

	// Call device IO Control interface (IOCTL_KOMCORE_GET_INFO) for every
	// processor to find out if there are any belonging to the same core 
	// (hyperthreaded)

	printf("Core IDs of each logical processor:\n");
	for (i=0; i<procs; i++)
	{	
		DWORD affinity=1<<i;
		int j;
		bool coreFound=false;

		SetThreadAffinityMask(GetCurrentThread(),affinity);
		Sleep(100);

		if (!DeviceIoControl(hDevice,IOCTL_KOMCORE_GET_INFO,
				 NULL,0,
				 &info,sizeof(info),
				 &nOutput,NULL))
		{
			printf("\n\nERROR: IOCTL_KOMCORE_GET_INFO returns %0x.", GetLastError());
			delete[] coreIDs;
			delete[] procIdx;
			Exit(1);
		}

		printf("Proc. #%d=0x%02d, ",i,info.currCoreId);
		
		if (i==3) 
			printf("\n");

		//seek this core ID in the table:
		for (j=0; j<cores; j++)
			if (coreIDs[j]==info.currCoreId)
			{
				coreFound=true;
				break;
			}

		if (!coreFound)
		{	//core not found, add it to the table
			coreIDs[cores]=info.currCoreId;
			procIdx[cores]=i;
			cores++;
		}
	}
	printf("\n...found %d cores.\n\nReading temperatures:\n",cores);

	do
	{
		for (i=0; i<cores; i++)
		{	DWORD affinity=1<<(procIdx[i]);
			int temperature;

			SetThreadAffinityMask(GetCurrentThread(),affinity);
			Sleep(100);

			// Call device IO Control interface (IOCTL_KOMCORE_GET_TEMP) in driver
			if (!DeviceIoControl(hDevice,IOCTL_KOMCORE_GET_TEMP,
						 NULL,0,
						 &temperature,sizeof(temperature),
						 &nOutput,NULL))
			{
				printf("\n\nERROR: IOCTL_KOMCORE_GET_TEMP returns %0x.", GetLastError());
				delete[] coreIDs;
				delete[] procIdx;
				Exit(1);
			}
	
			if (temperature==TEMP_INVALID)
				printf("Core 0x%02x: INV ,  ",coreIDs[i]);
			else
				printf("Core 0x%02x: %d C,  ",coreIDs[i],temperature);
		}
		printf("\n");
		Sleep(1000);
	} while (continuous && !_kbhit());

	delete[] coreIDs;
	delete[] procIdx;
}
