/*
   * Copyright (C) 2014  René Kijewski  <rene.kijewski@fu-berlin.de>
   *
   * This library is free software; you can redistribute it and/or
   * modify it under the terms of the GNU Lesser General Public
   * License as published by the Free Software Foundation; either
   * version 2.1 of the License, or (at your option) any later version.
   *
   * This library is distributed in the hope that it will be useful,
   * but WITHOUT ANY WARRANTY; without even the implied warranty of
   * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
   * Lesser General Public License for more details.
   *
   * You should have received a copy of the GNU Lesser General Public
   * License along with this library; if not, write to the Free Software
   * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
   */
  
  /**
   * @ingroup     x86
   * @{
   *
   * @file
   * @brief       Virtual memory management.
   *
   * @author      René Kijewski <rene.kijewski@fu-berlin.de>
   *
   * @}
   */
  
  #include "x86_kernel_memory.h"
  #include "x86_interrupts.h"
  #include "x86_memory.h"
  #include "x86_registers.h"
  #include "cpu.h"
  #include "irq.h"
  #include "tlsf-malloc.h"
  
  #include <malloc.h>
  #include <stdint.h>
  #include <stdio.h>
  #include <string.h>
  #include <unistd.h>
  
  /* Compare Figure 44 (p. 99) of "Intel® Quark SoC X1000 Core Developer’s Manual" */
  #define PT_CR3_BITS  (0)
  #define PT_PDPT_BITS (PT_P)
  #define PT_PD_BITS  (PT_P | PT_RW | PT_US)
  
  #ifndef DEBUG_READ_BEFORE_WRITE
  #   define PT_HEAP_BITS (PT_P | PT_RW | PT_US | pt_xd)
  #else
  #   define PT_HEAP_BITS (PT_HEAP_BIT | PT_RW | PT_US | pt_xd)
  #endif
  
  static uint64_t pt_xd = PT_XD;
  
  typedef union page {
      char content[4096];
      uint64_t next_page;
      uint64_t indices[512];
  } __attribute__((aligned(0x1000))) page_t;
  
  static volatile page_t TEMP_PAGE;
  #define TEMP_PAGE_PT (((uintptr_t) &TEMP_PAGE / 0x1000) / 512)
  #define TEMP_PAGE_PTE  (((uintptr_t) &TEMP_PAGE / 0x1000) % 512)
  
  void x86_init_gdt(void)
  {
      static const struct gdt_entry gdt_entries[3] = {
          [0x0000 / 8] = {
              .limit_0_15            = 0,
              .base_0_15             = 0,
              .base_16_23            = 0,
              .access_byte           = 0,
              .limit_16_19_and_flags = 0,
              .base_24_31            = 0,
          },
          [0x0008 / 8] = {
              .limit_0_15            = 0xFFFF,
              .base_0_15             = 0,
              .base_16_23            = 0,
              .access_byte           = GDT_AB_EX | GDT_AB_S | GDT_AB_RING_0 | GDT_AB_PR,
              .limit_16_19_and_flags = 0xF | GDT_FLAG_SZ | GDT_FLAG_GR,
              .base_24_31            = 0,
          },
          [0x0010 / 8] = {
              .limit_0_15            = 0xFFFF,
              .base_0_15             = 0,
              .base_16_23            = 0,
              .access_byte           = GDT_AB_RW | GDT_AB_S | GDT_AB_RING_0 | GDT_AB_PR,
              .limit_16_19_and_flags = 0xF | GDT_FLAG_SZ | GDT_FLAG_GR,
              .base_24_31            = 0,
          },
      };
      static const struct gdtr_t gdt = {
          .size_minus_one = sizeof gdt_entries - 1,
          .offset = (unsigned long) &gdt_entries[0],
      };
  
      __asm__ volatile ("" :: "a"(0x0010));
  
      __asm__ volatile ("lgdt %0" :: "m"(gdt));
      __asm__ volatile ("ljmp $0x0008, $1f\n"
                    "1:");
  
      __asm__ volatile ("mov %ax, %ds");
      __asm__ volatile ("mov %ax, %es");
      __asm__ volatile ("mov %ax, %fs");
      __asm__ volatile ("mov %ax, %gs");
      __asm__ volatile ("mov %ax, %ss");
  }
  
  /* Addresses in PDPT, PD, and PT are linear addresses. */
  /* TEMP_PAGE is used to to access these pages. */
  static pae_page_table_t static_pts[X86_NUM_STATIC_PT];
  static pae_page_directory_t static_pds[X86_NUM_STATIC_PD];
  static pae_page_directory_pointer_table_t pdpt;
  
  static void init_elf_static_section(uint64_t bits, void *start_, void *end_)
  {
      unsigned long start = ((unsigned long) start_) / 0x1000;
      unsigned long end = (((unsigned long) end_) + 0x1000 - 1) / 0x1000;
      for (unsigned i = start; i < end; ++i) {
          unsigned pt_num = i / 512;
          unsigned pte_num = i % 512;
          static_pts[pt_num][pte_num] = (i * 0x1000) | PT_P | PT_G | bits;
      }
  }
  
  static void check_requirements(void)
  {
      uint64_t cpuid = cpuid_caps();
      if ((cpuid & CPUID_PAE) == 0) {
          puts("Your CPU does not support PAE! Halting.");
          x86_hlt();
      }
      if ((cpuid & CPUID_PGE) == 0)  {
          puts("Your CPU does not support PGE! Halting.");
          x86_hlt();
      }
      if ((cpuid & CPUID_MSR) == 0) {
          puts("Warning: Your CPU does not support MSR!\n"
               "         Setting PT_XD = 0.");
          pt_xd = 0;
      }
      else {
          /* enable NX bit (if possible) */
          uint64_t efer = msr_read(MSR_EFER);
          efer |= EFER_NXE;
          msr_set(MSR_EFER, efer);
          if (!(msr_read(MSR_EFER) & EFER_NXE)) {
              puts("Warning: Your hardware does not support the NX bit!\n"
                   "         Setting PT_XD = 0.");
              pt_xd = 0;
          }
      }
  }
  
  static void init_pagetable(void)
  {
      /* identity map tables */
      for (unsigned i = 0; i < X86_NUM_STATIC_PD; ++i) {
          pdpt[i] = ((uintptr_t) &static_pds[i]) | PT_PDPT_BITS;
      }
      for (unsigned i = 0; i < X86_NUM_STATIC_PT; ++i) {
          unsigned pd_num = i / 512;
          unsigned pt_num = i % 512;
          static_pds[pd_num][pt_num] = ((uintptr_t) &static_pts[i]) | PT_PD_BITS;
      }
      init_elf_static_section(PT_RW | pt_xd, (void *) 0, (void *) 0x100000);
      init_elf_static_section(PT_US, &_section_text_start, &_section_text_end);
      init_elf_static_section(PT_US | pt_xd, &_section_rodata_start, &_section_rodata_end);
      init_elf_static_section(PT_US | PT_RW | pt_xd, &_section_data_start, &_section_bss_end);
  
      /* activate PAE */
      /* FIXME: add x86_init_cr4() */
      uint32_t cr4 = cr4_read();
      cr4 |= CR4_PAE | CR4_MCE | CR4_PGE | CR4_PCE | CR4_OSXMMEXCPT;
      cr4 &= ~(CR4_VME | CR4_PVI | CR4_TSD | CR4_DE | CR4_PSE | CR4_OSFXSR | CR4_SMEP);
      cr4_write(cr4);
  
      /* load page table */
      cr3_write((uint32_t) &pdpt | PT_CR3_BITS);
  
      /* activate paging */
      uint32_t cr0 = cr0_read();
      cr0 |= CR0_PE | CR0_NE | CR0_WP | CR0_PG;
      cr0 &= ~(CR0_MP | CR0_EM | CR0_TS | CR0_AM | CR0_NW | CR0_CD);
      cr0_write(cr0);
  }
  
  static void set_temp_page(uint64_t addr)
  {
      static_pts[TEMP_PAGE_PT][TEMP_PAGE_PTE] = addr != -1ull ? addr | PT_P | PT_RW | pt_xd : 0;
      __asm__ volatile ("invlpg (%0)" :: "r"(&TEMP_PAGE));
  }
  
  static inline uint64_t min64(uint64_t a, uint64_t b)
  {
      return a <= b ? a : b;
  }
  
  static inline uint64_t max64(uint64_t a, uint64_t b)
  {
      return a >= b ? a : b;
  }
  
  static uint32_t init_free_pages_heap_position = (uintptr_t) &_heap_start;
  
  static uint64_t init_free_pages_sub(uint64_t table, uint64_t bits, unsigned index, uint64_t *start, uint64_t *pos)
  {
      set_temp_page(table);
      if (TEMP_PAGE.indices[index] & PT_P) {
          return TEMP_PAGE.indices[index] & PT_ADDR_MASK;
      }
  
      TEMP_PAGE.indices[index] = *start | bits;
  
      uint64_t result = *start;
      *start += 0x1000;
      *pos = max64(*start, *pos);
      init_free_pages_heap_position += 0x1000;
      return result;
  }
  
  static bool add_pages_to_pool(uint64_t start, uint64_t end)
  {
      start += 0xFFF;
      start &= ~0xFFF;
      end &= ~0xFFF;
  
      start = max64(start, (uintptr_t) &_kernel_memory_end);
      uint64_t pos = start;
  
      uint32_t addr = init_free_pages_heap_position >> 12;
      unsigned pte_i = addr % 512;
      addr >>= 9;
      unsigned pt_i = addr % 512;
      addr >>= 9;
      unsigned pd_i = addr;
  
      if (pd_i >= 4) {
          return false;
      }
  
      while (pos < end) {
          uint64_t table = (uintptr_t) &pdpt;
  
          table = init_free_pages_sub(table, PT_PDPT_BITS, pd_i, &start, &pos);
          if (pos >= end) {
              break;
          }
  
          table = init_free_pages_sub(table, PT_PD_BITS, pt_i, &start, &pos);
          if (pos >= end) {
              break;
          }
  
          set_temp_page(table);
          TEMP_PAGE.indices[pte_i] = pos | PT_HEAP_BITS;
          pos += 0x1000;
  
          if (++pte_i >= 512) {
              pte_i = 0;
              if (++pt_i >= 512) {
                  pt_i = 0;
                  if (++pd_i >= 4) {
                      break;
                  }
              }
          }
      }
  
      if (start < end) {
          cr3_write((uint32_t) &pdpt | PT_CR3_BITS); /* flush tlb */
          tlsf_add_pool((void *) init_free_pages_heap_position, end - start);
          init_free_pages_heap_position += end - start;
      }
  
      return true;
  }
  
  static void init_free_pages(void)
  {
      printf("Kernel memory: %p - %p\r\n",
             (void *)&_kernel_memory_start, (void *)&_kernel_memory_end);
      printf("  .text:   %p - %p\r\n", (void *)&_section_text_start, (void *)&_section_text_end);
      printf("  .rodata: %p - %p\r\n", (void *)&_section_rodata_start, (void *)&_section_rodata_end);
      printf("  .data:   %p - %p\r\n", (void *)&_section_data_start, (void *)&_section_data_end);
      printf("  .bss:    %p - %p\r\n", (void *)&_section_bss_start, (void *)&_section_bss_end);
      printf("Unmapped memory: %p - %p\r\n", (void *)&_kernel_memory_end, (void *)&_heap_start);
      printf("Heap start: %p\r\n", (void *)&_heap_start);
  
      unsigned long cnt = 0;
      uint64_t start, len;
      while (x86_get_memory_region(&start, &len, &cnt)) {
          uint64_t end = start + len;
          if (!add_pages_to_pool(start, end)) {
              break;
          }
      }
  
      unsigned long free_pages_count = (init_free_pages_heap_position - (uintptr_t) &_heap_start) / 4096;
      float mem_amount = free_pages_count * (4096 / 1024);
      const char *mem_unit = "kB";
      if (mem_amount >= 2 * 1024) {
          mem_amount /= 1024;
          mem_unit = "MB";
      }
      if (mem_amount >= 2 * 1024) {
          mem_amount /= 1024;
          mem_unit = "GB";
      }
      printf("There are %lu free pages (%.3f %s) available for the heap.\n", free_pages_count, mem_amount, mem_unit);
  }
  
  static unsigned handling_pf;
  static void pagefault_handler(uint8_t intr_num, struct x86_pushad *orig_ctx, unsigned long error_code)
  {
      (void) intr_num; /* intr_num == X86_INT_PF */
  
      ++handling_pf;
  
      switch (handling_pf) {
          case 1:
              break; /* first #PF */
  
          case 2: /* pagefault while handing a page fault. */
              puts("A page fault occured while handling a page fault!");
              x86_print_registers(orig_ctx, error_code);
              puts("Halting.");
              /* fall through */
  
          default: /* pagefault while printing #PF message, everything is lost */
              x86_hlt();
      }
  
  #ifdef DEBUG_READ_BEFORE_WRITE
      uint32_t virtual_addr = cr2_read();
      uint64_t pte = x86_get_pte(virtual_addr);
  #endif
  
      if (error_code & PF_EC_I) {
          puts("Page fault while fetching instruction.");
          x86_print_registers(orig_ctx, error_code);
          puts("Halting.");
          x86_hlt();
      }
  #ifdef DEBUG_READ_BEFORE_WRITE
      else if ((pte != NO_PTE) && !(pte & PT_P) && (pte & PT_HEAP_BIT)) {
          /* mark as present */
          TEMP_PAGE.indices[(virtual_addr >> 12) % 512] |= PT_P;
          __asm__ volatile ("invlpg (%0)" :: "r"(virtual_addr));
  
          /* initialize for easier debugging */
          uint32_t *p = (uint32_t *) (virtual_addr & ~0xfff);
          for (unsigned i = 0; i < 0x1000 / 4; ++i) {
              const union {
                  char str_value[4];
                  uint32_t int_value;
              } debug_init = { .str_value = "RIOT" };
              *p++ = debug_init.int_value;
          }
  
          /* print a warning if the page was read before written */
          if (!(error_code & PF_EC_W)) {
              unsigned long *sp = (void *) orig_ctx->sp; /* ip, cs, flags */
              printf("DEBUG: Read before write on heap address 0x%08x (physical: 0x%016llx) at 0x%08lx.\n",
                     virtual_addr, pte & PT_ADDR_MASK, sp[0]);
          }
      }
  #endif
      else if (error_code & PF_EC_P) {
          printf("Page fault: access violation while %s present page.\n", (error_code & PF_EC_W) ? "writing to" : "reading from");
          x86_print_registers(orig_ctx, error_code);
          puts("Halting.");
          x86_hlt();
      }
      else {
          printf("Page fault: access violation while %s non-present page.\n", (error_code & PF_EC_W) ? "writing to" : "reading from");
          x86_print_registers(orig_ctx, error_code);
          puts("Halting.");
          x86_hlt();
      }
  
      --handling_pf;
  }
  
  static void init_pagefault_handler(void)
  {
      x86_interrupt_handler_set(X86_INT_PF, &pagefault_handler);
  }
  
  void x86_init_memory(void)
  {
      check_requirements();
  
      init_pagetable();
      init_pagefault_handler();
      init_free_pages();
  
      puts("Virtual memory initialized");
  }
  
  uint64_t x86_get_pte(uint32_t addr)
  {
      addr >>= 12;
      unsigned pte_i = addr % 512;
      addr >>= 9;
      unsigned pt_i = addr % 512;
      addr >>= 9;
      unsigned pd_i = addr;
  
      if (pdpt[pd_i] & PT_P) {
          set_temp_page(pdpt[pd_i] & PT_ADDR_MASK);
          if (TEMP_PAGE.indices[pt_i] & PT_P) {
              set_temp_page(TEMP_PAGE.indices[pt_i] & PT_ADDR_MASK);
              return TEMP_PAGE.indices[pte_i];
          }
      }
      return NO_PTE;
  }
  
  static void virtual_pages_set_bits(uint32_t virtual_addr, unsigned pages, uint64_t bits)
  {
      while (pages-- > 0) {
          unsigned pte_i = (virtual_addr >> 12) % 512;
  
          uint64_t old_physical_addr = x86_get_pte(virtual_addr) & PT_ADDR_MASK;
          TEMP_PAGE.indices[pte_i] = old_physical_addr | bits;
          __asm__ volatile ("invlpg (%0)" :: "r"(virtual_addr));
  
          virtual_addr += 0x1000;
      }
  }
  
  void *x86_map_physical_pages(uint64_t physical_start, unsigned pages, uint64_t bits)
  {
      if (bits & PT_XD) {
          bits &= ~PT_XD;
          bits |= pt_xd;
      }
  
      /* We use an already set up space, so we are sure that the upper level page tables are allocated. */
      /* We cut out a slice and re-add the physical pages. */
      char *result = memalign(0x1000, pages * 0x1000);
      if (!result) {
          return NULL;
      }
  
      for (unsigned page = 0; page < pages; ++page) {
          uint64_t physical_addr = physical_start + page * 0x1000;
          uint32_t virtual_addr = (uintptr_t) result + page * 0x1000;
          unsigned pte_i = (virtual_addr >> 12) % 512;
  
          uint64_t old_pte = x86_get_pte(virtual_addr);
          TEMP_PAGE.indices[pte_i] = physical_addr | bits;
  
          if (page == 0) {
              uint64_t old_physical_addr = old_pte & PT_ADDR_MASK;
  
              /* FIXME: does this work? Won't work if TLSF joins different buffers. */
              add_pages_to_pool(old_physical_addr, old_physical_addr + 0x1000 * pages);
          }
      }
  
      cr3_write((uint32_t) &pdpt | PT_CR3_BITS); /* flush tlb */
      return result;
  }
  
  void *x86_get_virtual_pages(unsigned pages, uint64_t bits)
  {
      if (bits & PT_XD) {
          bits &= ~PT_XD;
          bits |= pt_xd;
      }
  
      char *result = memalign(0x1000, pages * 0x1000);
      if (!result) {
          return (void *) -1ul;
      }
  
      virtual_pages_set_bits((uintptr_t) result, pages, bits);
      return result;
  }
  
  void x86_release_virtual_pages(uint32_t virtual_start, unsigned pages)
  {
      virtual_pages_set_bits(virtual_start, pages, PT_HEAP_BITS);
      free((void *) virtual_start);
  }