#include <log.h>
#include <mem.h>
#include <pager.h>
#include <stdint.h>
#include <sysconfig.h>
#include <tag.h>
#include <translation_table.h>

// System registers

#define SCTLR_EL1_WXN (1 << 19)
#define SCTLR_EL1_I (1 << 12)
#define SCTLR_EL1_C (1 << 2)
#define SCTLR_EL1_A (1 << 1)
#define SCTLR_EL1_M (1 << 0)

#define TCR_EL1_IPS__SHIFT 32
#define TCR_EL1_IPS__MASK (7UL << 32)
#define TCR_EL1_IPS_4GB 0UL
#define TCR_EL1_IPS_64GB 1UL
#define TCR_EL1_IPS_1TB 2UL

#define TCR_EL1_EPD1 (1 << 23)
#define TCR_EL1_A1 (1 << 22)

#define TCR_EL1_TG0__SHIFT 14
#define TCR_EL1_TG0__MASK (3 << 14)
#define TCR_EL1_TG0_64KB 1

#define TCR_EL1_SH0__SHIFT 12
#define TCR_EL1_SH0__MASK (3 << 12)
#define TCR_EL1_SH0_INNER 3

#define TCR_EL1_ORGN0__SHIFT 10
#define TCR_EL1_ORGN0__MASK (3 << 10)
#define TCR_EL1_ORGN0_WR_BACK_ALLOCATE 1
#define TCR_EL1_ORGN0_WR_BACK 3

#define TCR_EL1_IRGN0__SHIFT 8
#define TCR_EL1_IRGN0__MASK (3 << 8)
#define TCR_EL1_IRGN0_WR_BACK_ALLOCATE 1
#define TCR_EL1_IRGN0_WR_BACK 3

#define TCR_EL1_EPD0 (1 << 7)

#define TCR_EL1_T0SZ__SHIFT 0
#define TCR_EL1_T0SZ__MASK (0x3f << 0)
#define TCR_EL1_T0SZ_4GB 32
#define TCR_EL1_T0SZ_64GB 28
#define TCR_EL1_T0SZ_128GB 27

static void enable_mmu(void) {
  uint64_t mair_el1 =
      0xff << ATTR_INDX_NORMAL *
                  8 // inner/outer write-back non-transient, allocating
      | 0x04 << ATTR_INDX_DEVICE * 8    // Device-nGnRE
      | 0x00 << ATTR_INDX_COHERENT * 8; // Device-nGnRnE

  __asm__ __volatile("msr mair_el1, %0" : : "r"(mair_el1));
  __asm__ __volatile("msr ttbr0_el1, %0"
                     :
                     : "r"(gilbraltar_translation_table_base()));

  uint64_t tcr_el1;
  __asm__ __volatile("mrs %0, tcr_el1" : "=r"(tcr_el1));

  tcr_el1 &= ~(TCR_EL1_IPS__MASK | TCR_EL1_A1 | TCR_EL1_TG0__MASK |
               TCR_EL1_SH0__MASK | TCR_EL1_ORGN0__MASK | TCR_EL1_IRGN0__MASK |
               TCR_EL1_EPD0 | TCR_EL1_T0SZ__MASK);
  tcr_el1 |= TCR_EL1_EPD1 | TCR_EL1_TG0_64KB << TCR_EL1_TG0__SHIFT |
             TCR_EL1_SH0_INNER << TCR_EL1_SH0__SHIFT |
             TCR_EL1_ORGN0_WR_BACK_ALLOCATE << TCR_EL1_ORGN0__SHIFT |
             TCR_EL1_IRGN0_WR_BACK_ALLOCATE << TCR_EL1_IRGN0__SHIFT |
             TCR_EL1_IPS_1TB << TCR_EL1_IPS__SHIFT |
             TCR_EL1_T0SZ_128GB << TCR_EL1_T0SZ__SHIFT;

  __asm__ __volatile("msr tcr_el1, %0" : : "r"(tcr_el1));

  uint64_t sctlr_el1;
  __asm__ __volatile("mrs %0, sctlr_el1" : "=r"(sctlr_el1));

  sctlr_el1 &= ~(SCTLR_EL1_WXN | SCTLR_EL1_A);
  sctlr_el1 |= SCTLR_EL1_I | SCTLR_EL1_C | SCTLR_EL1_M;
  __asm__ __volatile("msr sctlr_el1, %0" : : "r"(sctlr_el1) : "memory");
}

void gilbraltar_memory_init(void) {
  size_t rsv = ARM_MEM_SIZE - MEM_HEAP_START - PAGE_RESERVE;

  gilbraltar_pager_init(MEM_HEAP_START + rsv, PAGE_RESERVE);
  gilbraltar_log(INFO, "          pager @ 0x%08lx - 0x%08lx\r\n",
                 MEM_HEAP_START + rsv, MEM_HEAP_START + rsv + PAGE_RESERVE);

  gilbraltar_translation_table_init(ARM_MEM_SIZE);
  enable_mmu();
  instruction_sync_barrier();
}