#include <efi.h>
#include <efilib.h>
#include <efiapi.h>

#include <cpuid.h>
#include <x86/msr.h>
#include <x86/cr.h>
#include <x86/vmx.h>

#include "virt.h"
#include "std.h"

static struct VMXON *vmxon_region;
static struct VMCS  *vmcs_region;

static bool
vintel_has_bios_support(void)
{
	uint64_t value = readmsr64(IA32_FEATURE_CONTROL);
	return (value & 0x5) == 0x5;
}

static bool
vintel_has_cpu_support(void)
{
	unsigned info[4];
	__get_cpuid(0x1, &info[0], &info[1], &info[2], &info[3]);
	return (info[2] & (1 << 5)) != 0;
}

bool
vintel_has_support(char *err, size_t errmax)
{
	if (!vintel_has_cpu_support()) {
		strlcpy(err, "CPU does not support VT-x.", errmax);
		return false;
	}

	if (!vintel_has_bios_support()) {
		strlcpy(err, "VT-x support is not enabled in BIOS.", errmax);
		return false;
	}

	return true;
}

static void
vintel_fix_cr_bits(void)
{
	uint64_t cr0_fixed0 = readmsr64(IA32_VMX_CR0_FIXED0);
	uint64_t cr0_fixed1 = readmsr64(IA32_VMX_CR0_FIXED1);
	uint64_t cr4_fixed0 = readmsr64(IA32_VMX_CR4_FIXED0);
	uint64_t cr4_fixed1 = readmsr64(IA32_VMX_CR4_FIXED1);

	writecr0((readcr0() | cr0_fixed0) & cr0_fixed1);
	writecr4((readcr4() | cr4_fixed0) & cr4_fixed1);
}

#define GUEST_CR0 0x6800
#define GUEST_CR3 0x6802
#define GUEST_CR4 0x6804
#define GUEST_DR7 0x681A
#define GUEST_RSP 0x681C
#define GUEST_RIP 0x681E
#define GUEST_RFLAGS 0x6820
#define VMCS_LINK_POINTER 0x2800
#define GUEST_IA32_DEBUGCTL 0x2802
#define GUEST_IA32_PAT 0x2804
#define GUEST_IA32_EFER 0x2806
#define GUEST_FS_BASE 0x680E
#define GUEST_GS_BASE 0x6810
#define GUEST_SYSENTER_CS  0x482A
#define GUEST_SYSENTER_ESP 0x6824
#define GUEST_SYSENTER_EIP 0x6826

static void
vintel_init_guest(void)
{
	uint64_t guestSP = 0x0;
	uint64_t guestIP = 0x0;

	vmwrite(GUEST_CR0,           readcr0());
	vmwrite(GUEST_CR3,           readcr3());
	vmwrite(GUEST_CR4,           readcr4());
	vmwrite(GUEST_DR7,           readdr7());
	vmwrite(GUEST_RFLAGS,        readflags());
	vmwrite(GUEST_RSP,           guestSP);
	vmwrite(GUEST_RIP,           guestIP);
	vmwrite(VMCS_LINK_POINTER,   -1LL);
	vmwrite(GUEST_IA32_DEBUGCTL, readmsr64(IA32_DEBUGCTL));
	vmwrite(GUEST_IA32_PAT,      readmsr64(IA32_PAT));
	vmwrite(GUEST_IA32_EFER,     readmsr64(IA32_EFER));
	vmwrite(GUEST_FS_BASE,       readmsr64(IA32_FS_BASE));
	vmwrite(GUEST_GS_BASE,       readmsr64(IA32_GS_BASE));
	vmwrite(GUEST_SYSENTER_CS,   readmsr64(0x174));
	vmwrite(GUEST_SYSENTER_ESP,  readmsr64(0x175));
	vmwrite(GUEST_SYSENTER_EIP,  readmsr64(0x176));
}

#if 0
static void
vintel_init_host(void)
{
	// Read TR
	trBase.LowPart = ((trItem[0] >> 16) & 0xFFFF) | ((trItem[1] & 0xFF) << 16) | ((trItem[1] & 0xFF000000) >> 8);
	trBase.HighPart = trItem[2];

	// Set TR
	vmwrite(HOST_TR_BASE, trBase.QuadPart);
	vmwrite(HOST_TR_SELECTOR, trSelector);

	// Set segment selectors
	vmwrite(HOST_ES_SELECTOR, AsmReadES() & 0xfff8);
	vmwrite(HOST_CS_SELECTOR, AsmReadCS() & 0xfff8);
	vmwrite(HOST_SS_SELECTOR, AsmReadSS() & 0xfff8);
	vmwrite(HOST_DS_SELECTOR, AsmReadDS() & 0xfff8);
	vmwrite(HOST_FS_SELECTOR, AsmReadFS() & 0xfff8);
	vmwrite(HOST_GS_SELECTOR, AsmReadGS() & 0xfff8);

	// Set control registers
	vmwrite(HOST_CR0, readcr0());
	vmwrite(HOST_CR3, readcr3());
	vmwrite(HOST_CR4, readcr4());

	// Set RSP and RIP
	vmwrite(HOST_RSP, (ULONG64)pVcpu->VmxHostStackBase);
	vmwrite(HOST_RIP, HostEip);

	// Set MSRs
	vmwrite(HOST_IA32_PAT, __readmsr(IA32_MSR_PAT));
	vmwrite(HOST_IA32_EFER, __readmsr(IA32_MSR_EFER));
	vmwrite(HOST_FS_BASE, __readmsr(IA32_FS_BASE));
	vmwrite(HOST_GS_BASE, __readmsr(IA32_GS_KERNEL_BASE));
	vmwrite(HOST_IA32_SYSENTER_CS, __readmsr(0x174));
	vmwrite(HOST_IA32_SYSENTER_ESP, __readmsr(0x175));
	vmwrite(HOST_IA32_SYSENTER_EIP, __readmsr(0x176));

	// Set GDT and IDT
	GdtTable idtTable;
	__sidt(&idtTable);
	vmwrite(HOST_GDTR_BASE, gdtTable.Base);
	vmwrite(HOST_IDTR_BASE, idtTable.Base);
}
#endif

void
vintel_enable(void)
{
	Print(L"Set CR4 Bit\n");
	writecr4(readcr4() | CR4_VMXE);

#if 0
	Print(L"Enable in Feature Control\n");
	uint64_t fctl = rdmsr64(IA32_FEATURE_CONTROL);
	fctl |= (1 << 0);
	wrmsr64(IA32_FEATURE_CONTROL, fctl);
#endif

	vintel_fix_cr_bits();

	Print(L"Allocating VMXON Page\n");
	uefi_call_wrapper(BS->AllocatePages, 4, AllocateAnyPages, EfiRuntimeServicesData, 1, &vmxon_region);

	Print(L"VMXON: %p\n", vmxon_region);
	vmxon_region->revisionID = readmsr32(IA32_VMX_BASIC);

	uint64_t status = vmxon(vmxon_region);
	Print(L"VMXON Status: %p\n", (void *)status);
	if (status & (1 << 0)) {
		Print(L"Invalid VMCS Pointer\n");
	}
	if (status & (1 << 6)) {
		Print(L"Extended VMX Error\n");
	}

	Print(L"Allocating VMCS Page\n");
	uefi_call_wrapper(BS->AllocatePages, 4, AllocateAnyPages, EfiRuntimeServicesData, 1, &vmcs_region);

	Print(L"VMCS: %p\n", vmcs_region);
	vmcs_region->revisionID = readmsr32(IA32_VMX_BASIC);

	status = vmclear(vmcs_region);
	Print(L"VMCLEAR Status: %p\n", (void *)status);

	status = vmptrld(vmcs_region);
	Print(L"VMPTRLD Status: %p\n", (void *)status);

	vintel_init_guest();
}

struct virt_vtable virt_vtable_intel = {
	.has_support = vintel_has_support,
	.enable      = vintel_enable,
};