#include "mmu.h"

#include "cpu/exception.h"

#include <stdio.h>
#include <stdlib.h>

extern uint8_t* memory;

// Memory Managment Unit implementation
// We only support Sv32

#define PAGE_SIZE (4 * 1024) // 4KiB, 2^12B
#define LEVELS 2
#define PTE_SIZE 4 // sizeof(uint32_t)

// SATP CSR Register: [MODE(1bit) ASID(9bits) PPN(22bits)]
#define SATP_MODE (1 << 31)
#define SATP_MODE_BARE (0)
#define SATP_MODE_SV32 (1 << 31)
#define SATP_ASID (0x1FF << 22)
#define SATP_PPN (0x3FFFFF)

// Page Table entry: [PPN[1](12bits) PPN[0](10bits) RSW(2bits) D A G U X W R V]
#define PTE_PPN_1(pte) ((pte & 0xFFF00000) >> 20)
#define PTE_PPN_0(pte) ((pte & 0x000FFC00) >> 10)
#define PTE_PPN(pte) ((pte & 0xFFFFFC00) >> 10)
#define PTE_RSW (0b11 << 8)
#define PTE_D (1 << 7)
#define PTE_A (1 << 6)
#define PTE_G (1 << 5)
#define PTE_U (1 << 4)
#define PTE_X (1 << 3)
#define PTE_W (1 << 2)
#define PTE_R (1 << 1)
#define PTE_V (1 << 0)

// Physical address: 34 bits [PPN[1](12bits, ) PPN[0](10bits) Offset(12bits)]
// We only use 32-bits addresses, so the top 2 are always 0
#define PADDR_PAGE_OFFSET (0x00000FFF)

// Virtual address: [VPN[1](10bits) VPN[0](10bits) Offset(12bits)]
#define VADDR_VPN_1 (0xFFC00000)
#define VADDR_VPN_0 (0x003FF000)
#define VADDR_PAGE_OFFSET (0x00000FFF)

uint32_t mmu_scause_from_access(memory_access_type_t access_type)
{
	switch(access_type)
	{
		case READ:
			return SCAUSE_LOAD_PAGE_FAULT;
		case WRITE:
			return SCAUSE_STORE_AMO_PAGE_FAULT;
		case INSTRUCTION_FETCH:
			return SCAUSE_INSTRUCTION_PAGE_FAULT;
		default:
			fprintf(stderr, "mmu_scause_from_access: invalid parameter\n");
			exit(EXIT_FAILURE);
			break;
	}
}

uint32_t mmu_resolve(rv32_cpu_t* cpu, memory_access_type_t access_type, uint32_t vaddr)
{
	// TODO: Make sure we are in S-mode or U-mode

	// Check if MODE field is 'bare', meaning no mmu
	if((cpu->csr[CSR_SATP] & SATP_MODE) == SATP_MODE_BARE)
		return vaddr;

	// fprintf(stderr, "MMU enabled on (virtual) address 0x%x resolution\n", vaddr);

	uint32_t page_table = (cpu->csr[CSR_SATP] & SATP_PPN) * PAGE_SIZE;

	// Resolve first-level page table entry
	uint32_t vpn_1 = (vaddr & VADDR_VPN_1) >> 22;
	uint32_t pte_address = page_table + vpn_1 * PTE_SIZE;
	uint32_t pte = *((uint32_t*) (&memory[pte_address]));

	if(!(pte & PTE_V))
	{
		// Invalid PTE
		exception_trigger(cpu, mmu_scause_from_access(access_type), vaddr);
	}

	if((pte & PTE_R) || (pte & PTE_W) || (pte & PTE_X))
	{
		// Leaf PTE, we are ready to resolve the mapping
		// This is a 4 MiB megapage

		// For an execute, check if we are allowed to execute
		if(access_type == INSTRUCTION_FETCH && !(pte & PTE_X))
			exception_trigger(cpu, mmu_scause_from_access(access_type), vaddr);

		// For a write, check if we are allowed to write
		if(access_type == WRITE && !(pte & PTE_W))
			exception_trigger(cpu, mmu_scause_from_access(access_type), vaddr);

		// For a read, check if we are allowed to read
		if(access_type == READ && !(pte & PTE_R))
			exception_trigger(cpu, mmu_scause_from_access(access_type), vaddr);

		// Physical Address: [PPN[1] = pte.PPN[1], PPN[0] = vaddr.VPN[0], offset]
		uint32_t paddr = 0;
		paddr |= (PTE_PPN_1(pte) << 22);
		paddr |= (vaddr & VADDR_VPN_0);
		paddr |= vaddr & VADDR_PAGE_OFFSET;

		return paddr;
	}

	// PTE is a pointer to next level of page table
	page_table = PTE_PPN(pte) * PAGE_SIZE;

	// Resolve second-level page table entry
	uint32_t vpn_0 = (vaddr & VADDR_VPN_0) >> 12;
	pte_address = page_table + vpn_0 * PTE_SIZE;
	pte = *((uint32_t*) (&memory[pte_address]));

	if(!(pte & PTE_V))
	{
		// Invalid PTE
		exception_trigger(cpu, mmu_scause_from_access(access_type), vaddr);
	}

	// This must be a leaf PTE, as Sv32 only supports 2-level mappings
	// This is a 4 KiB page
	if(!((pte & PTE_R) || (pte & PTE_W) || (pte & PTE_X)))
	{
		fprintf(stderr, "Error: Pointer second-level Page Table Entry 0x%x at 0x%x while resolving virtual address 0x%x\n", pte, pte_address, vaddr);
		exit(EXIT_FAILURE);
	}

	// For an execute, check if we are allowed to execute
	if(access_type == INSTRUCTION_FETCH && !(pte & PTE_X))
		exception_trigger(cpu, mmu_scause_from_access(access_type), vaddr);

	// For a write, check if we are allowed to write
	if(access_type == WRITE && !(pte & PTE_W))
		exception_trigger(cpu, mmu_scause_from_access(access_type), vaddr);

	// For a read, check if we are allowed to read
	if(access_type == READ && !(pte & PTE_R))
		exception_trigger(cpu, mmu_scause_from_access(access_type), vaddr);

	// Physical Address: [PPN[1] = pte.PPN[1], PPN[0] = pte.PPN[0], offset]
	uint32_t paddr = 0;
	paddr |= (PTE_PPN_1(pte) << 22);
	paddr |= (PTE_PPN_0(pte) << 12);
	paddr |= vaddr & VADDR_PAGE_OFFSET;

	return paddr;
}