#include <Uefi.h>

#include <Library/UefiBootServicesTableLib.h>
#include <Library/UefiDriverEntryPoint.h>
#include <Library/BaseLib.h>
#include <Library/IoLib.h>

#include <Protocol/MpService.h>

#define IA32_TSC_MSR 0x10
#define POST_PORT 0x80


//#define DEBUG_TSCFIX
#define PERFECT 80


STATIC volatile UINT32 ReadyCount;
STATIC volatile UINT32 GoFlag;
STATIC volatile UINT32 DoneFlag;

STATIC volatile UINT64 Cpu0Tsc;
STATIC volatile UINT64 Cpu1Tsc;

STATIC VOID EFIAPI ReadTscWorker(IN VOID *Buffer)
{
    __sync_fetch_and_add(&ReadyCount, 1);
    while (__sync_fetch_and_add(&GoFlag, 0) == 0)
        CpuPause();
    MemoryFence();
    Cpu1Tsc = AsmReadTsc();
    MemoryFence();
    __sync_fetch_and_add(&DoneFlag, 1);
}

EFI_STATUS EFIAPI ModuleEntryPoint(
    IN EFI_HANDLE ImageHandle,
    IN EFI_SYSTEM_TABLE *SystemTable)
{
    EFI_STATUS               Status;
    EFI_MP_SERVICES_PROTOCOL *MpServices;
    UINTN NumProcessors;
    UINTN NumEnabledProcessors;
    
    Status = gBS->LocateProtocol(&gEfiMpServiceProtocolGuid, NULL, (VOID **)&MpServices);
    if (EFI_ERROR(Status)) {
        return Status;
    }
    
    Status = MpServices->GetNumberOfProcessors(MpServices, &NumProcessors, &NumEnabledProcessors);
    if (EFI_ERROR(Status)) {
        return Status;
    }
    
    if (NumEnabledProcessors < 2) {
        return EFI_UNSUPPORTED;
    }

	INT64 MinDelta;
	INT64 MaxDelta;
	UINT64 Correct = 25000;
	UINTN i;
	STATIC EFI_EVENT ApDoneEvent;
	Status = gBS->CreateEvent(
		0,
		TPL_CALLBACK,
		NULL,
		NULL,
		&ApDoneEvent
	);
	if (EFI_ERROR(Status)) {
#ifdef DEBUG_TSCFIX
		IoWrite8(POST_PORT, 0xE0);
		gBS->Stall(500000);
#endif
		return Status;
	}
	for (i = 0; i < 2550; i++) {
		ReadyCount = 0;
		GoFlag = 0;
		DoneFlag = 0;
		Cpu0Tsc = 0;
		Cpu1Tsc = 0;
		gBS->Stall(10);
		Status = MpServices->StartupThisAP(
			MpServices,
			ReadTscWorker,
			1,
			ApDoneEvent,
			0,
			NULL,
			NULL
		);
		if (EFI_ERROR(Status)) {
#ifdef DEBUG_TSCFIX
			IoWrite8(POST_PORT, 0xE1);
			gBS->Stall(500000);
#endif
			return Status;
		}
		while (__sync_fetch_and_add(&ReadyCount, 0) == 0)
			CpuPause();
		MemoryFence();
		__sync_fetch_and_add(&GoFlag, 1);
		MemoryFence();
		Cpu0Tsc = AsmReadTsc();
		while (__sync_fetch_and_add(&DoneFlag, 0) == 0)
			CpuPause();
		if((i % 10) == 0) {
			AsmWriteMsr64(IA32_TSC_MSR, Cpu1Tsc + Correct);
			MinDelta = MAX_INT64;
			MaxDelta = MIN_INT64;
			continue;
		}
		INT64 Delta = (INT64)Cpu1Tsc - (INT64)Cpu0Tsc;
		if(Delta < MinDelta)
			MinDelta = Delta;
		if(Delta > MaxDelta)
			MaxDelta = Delta;
		if((i % 10) == 9) {
			if(MinDelta < PERFECT && MaxDelta > PERFECT)
				break;
			Correct += MinDelta - PERFECT;
		}
	}
	gBS->CloseEvent(ApDoneEvent);

#ifdef DEBUG_TSCFIX
	i /= 10;
	UINT64 x = 128 + MinDelta - PERFECT;
	UINT64 y = 128 + MaxDelta - PERFECT;

    IoWrite8(POST_PORT, i & 0xFF);
    gBS->Stall(500000);
    IoWrite8(POST_PORT, x & 0xFF);
    gBS->Stall(500000);
    IoWrite8(POST_PORT, y & 0xFF);
    gBS->Stall(500000);
#endif

    return EFI_SUCCESS;
}