#include <lcrash/mm/virt.h>

#include <lcrash/mm/kmalloc.h>
#include <lcrash/debug/debug.h>

u64 (*VmemOriginalPML4)[512] = 0;
u64 (*VmemPML4)[512] = 0;

void VmemInitialize() {
	// Copy control registers
	u64 CR3;
	asm ("mov %%cr3, %0" : "=r" (CR3));

	u64 CR4;
	asm ("mov %%cr4, %0" : "=r" (CR4));

	if (CR4 & 0x800) Panic("5 Level paging not supported");

	u64 (*OriginalPML4)[512] = (u64(*)[512])(CR3 & 0x0007fffffffff000);
	u64 (*NewPML4)[512] = KernelHeapAlloc(sizeof(u64) * 512, 0x1000, 0);

	for (int i = 0; i < 512; i++) {
		if ((*OriginalPML4)[i] & 1) {
			u64 (*OriginalPDPT)[512] = (u64(*)[512])(((*OriginalPML4)[i] & 0x0007fffffffff000));
			u64 (*NewPDPT)[512] = KernelHeapAlloc(sizeof(u64) * 512, 0x1000, 0);

			for (int j = 0; j < 512; j++) {
				if ((*OriginalPDPT)[j] & 1) {
					u64 (*OriginalPD)[512] = (u64(*)[512])(((*OriginalPDPT)[j] & 0x000ffffffffff000));
					u64 (*NewPD)[512] = KernelHeapAlloc(sizeof(u64) * 512, 0x1000, 0);

					for (int k = 0; k < 512; k++) {
						if ((*OriginalPD)[k] & 1) {
							u64 (*OriginalPT)[512] = (u64(*)[512])(((*OriginalPD)[k] & 0x000ffffffffff000));
							u64 (*NewPT)[512] = KernelHeapAlloc(sizeof(u64) * 512, 0x1000, 0);

							if ((s64)NewPT <= 0) Panic("Failed to allocate memory");

							for (int l = 0; l < 512; l++) {
								//if ((*OriginalPT)[l] & 1) {
								//	(*NewPT)[l] = (*OriginalPT)[l];
								//} else (*NewPT)[l] = 0;*
								(*NewPT)[l] = (*OriginalPT)[l] | 0x87;
							}

							NewPT = (void*)((u64)OriginalPT | 0x87);

							(*NewPD)[k] = (u64)NewPT;
						} else (*NewPD)[k] = 0;
					}

					NewPD = (void*)((uptr)NewPD | (((*OriginalPDPT)[j]) & ~0x0007fffffffff000));

					(*NewPDPT)[j] = (u64)NewPD;
				} else (*NewPDPT)[j] = 0;
			}

			NewPDPT = (void*)((uptr)NewPDPT | (((*OriginalPML4)[i]) & ~0x0007fffffffff000));

			(*NewPML4)[i] = (u64)NewPDPT;
		} else (*NewPML4)[i] = 0;
	}

	// Save these
	VmemOriginalPML4 = OriginalPML4;
	VmemPML4 = NewPML4;

	// Load our new page table
	CR3 ^= (u64)VmemOriginalPML4;
	CR3 |= ((u64)VmemPML4) & 0x0007fffffffff000;
	asm volatile ("mov %q0, %%cr3" : : "r" (CR3));
	
	Panic("FUCK");
}