#include "sw/device/lib/virtual_memory.h"

#include "sw/device/lib/testing/test_framework/check.h"

// Defines to help build page tables for the SV32 scheme.
// More details can be found in the
// RISC-V privileged spec.

// Mask to preserve the top 22-bits of a physical address.
#define SV32_PPN_MASK (0xFFFFFC00U)
// Shift to apply to PPNs when making a PTE.
// This is due to physical addresses being 34-bits under the hood.
// See section 4.3.1 in the privileged spec for more detail.
#define SV32_PPN_SHIFT (2)
// Number of bits to shift right to get the VPN0 of a virtual address.
// VPN1 is used as an index into the L1 page-table.
#define SV32_VPN1_SHIFT (22)
// Number of bits to shift right to get the VPN0 of a virtual address.
// VPN0 is used as an index into the L2 page-table.
#define SV32_VPN0_SHIFT (12)
// Mask to apply after shifting a virtual address, to preserve VPN0 only.
#define SV32_VPN0_MASK (0x3FF)
// PTE status bits, setting the page as valid and RWX.
#define SV32_PAGE_SUPERVISOR (0xCF)
// PTE status bits, setting the page as valid, RWX, and available to user mode.
#define SV32_PAGE_USER (0xDF)
// PTE status bits, to indicate that the PTE is valid. Used alone for non-leaf
// pages.
#define SV32_PTE_VALID (1)

void machine_map_megapage(void *paddr, void *vaddr, uint32_t *pt) {
  uint32_t page = (uint32_t)vaddr >> SV32_VPN1_SHIFT;
  uint32_t leaf = (((uint32_t)paddr & SV32_PPN_MASK) >> SV32_PPN_SHIFT) |
                  SV32_PAGE_SUPERVISOR;
  pt[page] = leaf;
}

void machine_map_region(uint32_t paddr, uint32_t vaddr_start,
                        uint32_t vaddr_end, uint32_t *l1pt, uint32_t *l2pt,
                        bool user) {
  CHECK((vaddr_start % PAGE_SIZE) == 0);
  CHECK((paddr % PAGE_SIZE) == 0);

  uint32_t ppn1 = vaddr_start >> SV32_VPN1_SHIFT;
  uint32_t pte_lvl1 =
      (((uint32_t)l2pt & SV32_PPN_MASK) >> SV32_PPN_SHIFT) | SV32_PTE_VALID;
  l1pt[ppn1] = pte_lvl1;
  uint32_t page_count = vaddr_end - vaddr_start;
  page_count = (page_count + (PAGE_SIZE - 1)) & ~(PAGE_SIZE - 1);
  page_count /= PAGE_SIZE;
  for (int i = 0; i < page_count; ++i) {
    uint32_t paddr_l2 = paddr + (PAGE_SIZE * i);
    uint32_t vaddr_l2 = vaddr_start + (PAGE_SIZE * i);
    uint32_t vpn0_2 = (vaddr_l2 >> SV32_VPN0_SHIFT) & SV32_VPN0_MASK;
    uint32_t pte_lvl2 = ((paddr_l2 & SV32_PPN_MASK) >> SV32_PPN_SHIFT);
    if (user) {
      pte_lvl2 |= SV32_PAGE_USER;
    } else {
      pte_lvl2 |= SV32_PAGE_SUPERVISOR;
    }
    l2pt[vpn0_2] = pte_lvl2;
  }
}

__attribute__((naked)) void supervisor_trampoline(uint32_t a0_smode_fn,
                                                  uint32_t a1_satp,
                                                  uint32_t a2_medeleg,
                                                  uint32_t a3_stvec,
                                                  uint32_t a4_mideleg) {
  asm volatile(
      ".option push\n"
      ".option norelax\n"

      "csrw mideleg, a4\n"
      "csrw medeleg, a2\n"
      "csrw stvec, a3\n"
      "csrw satp, a1\n"
      "sfence.vma\n"

      "li x1, 0x0000080a\n"
      "csrw mstatus, x1\n"
      "csrw mepc, a0\n"
      "mret\n"
      ".option pop\n"
      :
      :
      : "memory");
}

__attribute__((naked)) void user_trampoline(uint32_t a0_umode_fn) {
  asm volatile(
      ".option push\n"
      ".option norelax\n"

      "csrw sepc, a0\n"
      "sret\n"
      ".option pop\n"
      :
      :
      : "memory");
}