use efi;
use log;

let L0: [512]uintptr;
let L1_ident: [512]uintptr;
let L1_kernel: [512]uintptr;
let L2_kernel: [512]uintptr;
let L3_kernel: [512]uintptr;

def PAGESIZE: size = 4096;
def UPAGESIZE: uintptr = 4096;

// Memory attributes
def PT_NORMAL: uintptr = 0 << 2;
def PT_DEVICE: uintptr = 1 << 2;

def PT_TABLE: uintptr = 0b11;
def PT_BLOCK: uintptr = 0b01;
def PT_OFFS: uintptr = 48;

def PT_XN: uintptr     = (1 << 54); // execute never
def PT_PXN: uintptr    = (1 << 53); // privileged execute never
def PT_CONT: uintptr   = (1 << 52); // continuous bit

def PT_NG: uintptr     = (1 << 11); // non-global
def PT_AF: uintptr     = (1 << 10); // access
def PT_ISM: uintptr    = (3 << 8);  // inner sharable
def PT_OSM: uintptr    = (2 << 8);  // outer sharable
def PT_RW: uintptr     = (0 << 7);  // Read/write
def PT_RO: uintptr     = (1 << 7);  // Read-only

// Initializes the memory map and returns the physical address of the largest
// area of conventional memory, in lieu of a better memory allocator.
fn init_mmap() uintptr = {
	const iter = efi::iter_mmap()!;
	let maxphys: uintptr = 0, maxpages = 0u64;
	for (true) {
		const desc = match (efi::mmap_next(&iter)) {
		case let desc: *efi::MEMORY_DESCRIPTOR =>
			yield desc;
		case void =>
			break;
		};
		if (desc.DescriptorType != efi::MEMORY_TYPE::CONVENTIONAL) {
			continue;
		};
		if (desc.NumberOfPages > maxpages) {
			maxphys = desc.PhysicalStart;
			maxpages = desc.NumberOfPages;
		};
	};
	assert(maxphys != 0, "No suitable memory area found for kernel loader");
	return maxphys;
};

fn init_pagetables() void = {
	// 0xFFFF0000xxxxxxxx - 0xFFFF0200xxxxxxxx: identity map
	// 0xFFFF8000xxxxxxxx - 0xFFFF8000xxxxxxxx: kernel image
	//
	// L0[0x00] => L1_ident
	// L0[0x80] => L1_kernel
	// L1_ident => 1 GiB identity mappings
	// L1_kernel[0] => L2_kernel
	// L2_kernel[0] => L3_kernel
	// L3_kernel[*] => 4 KiB kernel pages
	L0[0] = PT_TABLE | &L1_ident: uintptr;
	L0[0x80] = PT_TABLE | &L2_kernel: uintptr;
	L2_kernel[0] = PT_TABLE | &L3_kernel: uintptr;
	for (let i = 0u64; i < len(L1_ident): u64; i += 1) {
		L1_ident[i] = PT_BLOCK | (i * 0x40000000): uintptr |
			PT_NORMAL | PT_AF | PT_ISM | PT_RW;
	};
};

// Maps a physical page into the kernel's virtual address space.
fn kmmap(phys: uintptr, virt: uintptr, flags: uintptr) void = {
	return; // TODO
};

def SCTLR_EL1_MMU: u64 = 1 << 0;

def TCR_EL1_DS: u64     = (1 << 59);
def TCR_EL1_TCMA1: u64  = (1 << 58);
def TCR_EL1_TCMA0: u64  = (1 << 57);
def TCR_EL1_E0PD1: u64  = (1 << 56);
def TCR_EL1_E0PD2: u64  = (1 << 55);
def TCR_EL1_NFD0: u64   = (1 << 54);
def TCR_EL1_NFD1: u64   = (1 << 53);
def TCR_EL1_TBID0: u64  = (1 << 52);
def TCR_EL1_TBID1: u64  = (1 << 51);
def TCR_EL1_HWU162: u64 = (1 << 50);
def TCR_EL1_HWU161: u64 = (1 << 49);
def TCR_EL1_HWU160: u64 = (1 << 48);
def TCR_EL1_HWU159: u64 = (1 << 47);
def TCR_EL1_HWU062: u64 = (1 << 46);
def TCR_EL1_HWU061: u64 = (1 << 45);
def TCR_EL1_HWU060: u64 = (1 << 44);
def TCR_EL1_HWU059: u64 = (1 << 43);
def TCR_EL1_HPD1: u64   = (1 << 42);
def TCR_EL1_HPD0: u64   = (1 << 41);
def TCR_EL1_HD: u64     = (1 << 40);
def TCR_EL1_HA: u64     = (1 << 39);
def TCR_EL1_TBI1: u64   = (1 << 38);
def TCR_EL1_TBI0: u64   = (1 << 37);
def TCR_EL1_AS: u64     = (1 << 36);

def TCR_EL1_IPS_MASK: u64   = (0b111 << 32);

def TCR_EL1_IPS_32B_4G: u64   = (0b000 << 32);
def TCR_EL1_IPS_36B_64G: u64  = (0b001 << 32);
def TCR_EL1_IPS_40B_1T: u64   = (0b010 << 32);
def TCR_EL1_IPS_42B_4T: u64   = (0b011 << 32);
def TCR_EL1_IPS_44B_16T: u64  = (0b100 << 32);
def TCR_EL1_IPS_48B_256T: u64 = (0b101 << 32);
def TCR_EL1_IPS_52B_5P: u64   = (0b111 << 32);

def TCR_EL1_TG1_16K: u64 = (0b01 << 30);
def TCR_EL1_TG1_4K: u64  = (0b10 << 30);
def TCR_EL1_TG1_64K: u64 = (0b11 << 30);

def TCR_EL1_SH1_NS: u64 = (0b00 << 28);
def TCR_EL1_SH1_OS: u64 = (0b10 << 28);
def TCR_EL1_SH1_IS: u64 = (0b11 << 28);

def TCR_EL1_ORGN1_NC: u64 = (0b00 << 26);
def TCR_EL1_ORGN1_WB: u64 = (0b01 << 26);
def TCR_EL1_ORGN1_WT: u64 = (0b10 << 26);
def TCR_EL1_ORGN1_NA: u64 = (0b11 << 26);

def TCR_EL1_IRGN1_NC: u64 = (0b00 << 24);
def TCR_EL1_IRGN1_WB: u64 = (0b01 << 24);
def TCR_EL1_IRGN1_WT: u64 = (0b10 << 24);
def TCR_EL1_IRGN1_NA: u64 = (0b11 << 24);

def TCR_EL1_EPD1: u64 = (1 << 23);
def TCR_EL1_A1: u64 = (1 << 22);

def TCR_EL1_T1SZ: u64 = 16;

def TCR_EL1_TG0_4K: u64  = (0b01 << 14);
def TCR_EL1_TG0_64K: u64 = (0b10 << 14);
def TCR_EL1_TG0_16K: u64 = (0b11 << 14);

def TCR_EL1_SH0_NS: u64 = (0b00 << 12);
def TCR_EL1_SH0_OS: u64 = (0b10 << 12);
def TCR_EL1_SH0_IS: u64 = (0b11 << 12);

def TCR_EL1_ORGN0_NC: u64 = (0b00 << 10);
def TCR_EL1_ORGN0_WB: u64 = (0b01 << 10);
def TCR_EL1_ORGN0_WT: u64 = (0b10 << 10);
def TCR_EL1_ORGN0_NA: u64 = (0b11 << 10);

def TCR_EL1_IRGN0_NC: u64 = (0b00 << 8);
def TCR_EL1_IRGN0_WB: u64 = (0b01 << 8);
def TCR_EL1_IRGN0_WT: u64 = (0b10 << 8);
def TCR_EL1_IRGN0_NA: u64 = (0b11 << 8);

def TCR_EL1_EPD0: u64 = (1 << 7);

def TCR_EL1_T0SZ: u64 = 0;

// Initializes the ARM MMU to our desired specifications. This should take place
// *after* EFI boot services have exited because we're going to mess up the MMU
// configuration that it depends on.
fn init_mmu() void = {
	// Disable MMU
	const sctlr_el1 = rdsctlr_el1();
	wrsctlr_el1(sctlr_el1 & ~SCTLR_EL1_MMU);

	// Configure MAIR
	const mair: u64 =
		(0xFF << 0) | // Attr0: Normal memory; IWBWA, OWBWA, NTR
		(0x00 << 8);  // Attr1: Device memory; nGnRnE, OSH
	wrmair_el1(mair);

	const ips = rdtcr_el1() & TCR_EL1_IPS_MASK;
	log::printfln("old TCR_EL1: {:b}", rdtcr_el1());
	const tcr_el1: u64 =
		ips |			// IPS from EFI
		TCR_EL1_TG1_4K |	// Higher half: 4K granule size
		TCR_EL1_SH1_IS |	// Higher half: inner shareable
		TCR_EL1_ORGN1_WB |	// Higher half: outer write-back
		TCR_EL1_IRGN1_WB |	// Higher half: inner write-back
		TCR_EL1_EPD1 |		// Higher half enable
		(24 << TCR_EL1_T1SZ) |	// Higher half: 3 levels (1024G)
		TCR_EL1_TG0_4K |	// Lower half: 4K granule size
		TCR_EL1_SH0_IS |	// Lower half: inner sharable
		TCR_EL1_ORGN0_WB |	// Lower half: outer write-back
		TCR_EL1_IRGN0_WB |	// Lower half: inner write-back
		TCR_EL1_EPD0 |		// Lower half enable
		(24 << TCR_EL1_T0SZ);	// Lower half: 3 levels (1024G)
	wrtcr_el1(tcr_el1);
	log::printfln("new TCR_EL1: {:b}", tcr_el1);

	// Load page tables
	wrttbr0_el1(&L0[0]: uintptr | 1);
	wrttbr1_el1(&L0[0]: uintptr | 1);

	// Enable MMU
	wrsctlr_el1(sctlr_el1 | SCTLR_EL1_MMU);

	// Test higher half identity map
	//let x = 1337;
	//let y = (&x: uintptr | 0xFFFF000000000000: uintptr): *int;
	//assert(*y == 1337);
};