#include <mercury/kernel/utility.hpp>
#include <mercury/kernel/paging.hpp>

//see also ../documentation/memory.txt

extern "C" {
  uint64_t __kernel_p4_paddr;
}

namespace mercury::kernel::paging {

  static constexpr uint64_t kernel_vram_start = 0xffffffffc0000000;
  static constexpr uint64_t kernel_vram_end = 0xffffffffffe00000;
  static constexpr uint64_t kernel_vram_pages =
    (kernel_vram_end - kernel_vram_start) / 4096;
  static constexpr uint64_t syscall_stack_bottom = 0xfffffffffff01000;
  static constexpr uint64_t syscall_stack_top = 0xfffffffffffff000;
  static constexpr uint64_t interrupt_stack_bottom = 0xffffffffffe01000;
  static constexpr uint64_t interrupt_stack_top = 0xffffffffffeff000;

  static constexpr uint64_t pram_pages = 1 << 23;
  static uint64_t pram_usage_bitmap[pram_pages / 64];

  void mark_all_pram_used() {
    utility::mark_bitmap_region_one(pram_usage_bitmap, 0, pram_pages);
  }

  void mark_pram_region_free(uint64_t start_addr, uint64_t end_addr) {
    utility::mark_bitmap_region_zero(
      pram_usage_bitmap, start_addr / 4096,
      utility::min(end_addr / 4096, pram_pages));
  }

  [[gnu::aligned(4096)]] uint64_t kernel_p4[512];
  [[gnu::aligned(4096)]] uint64_t kernel_p3[512];
  [[gnu::aligned(4096)]] uint64_t kernel_p2[512];
  [[gnu::aligned(4096)]] uint64_t kernel_p1s[512 * 512];

  uint64_t kernel_p4e;

  uint64_t encode_pte(
    uint64_t addr, bool user, bool write, bool execute, bool ps
  ) {
    return (addr & 0x0000ffffffffffff) | (execute ? 0 : (1ULL << 63))
         | (ps << 7) | (user << 2) | (write << 1) | 1;
  }

  void init_kernel_page_tables(uint64_t kernel_offset) {
    __kernel_p4_paddr = (uint64_t)kernel_p4 - kernel_offset;
    for (int i = 0; i < 511; ++i)
      kernel_p4[i] = 0;
    kernel_p4e = encode_pte(
      (uint64_t)kernel_p3 - kernel_offset, false, true, true, false);
    kernel_p4[511] = kernel_p4e;
    for (int i = 0; i < 511; ++i)
      kernel_p3[i] = 0;
    kernel_p3[511] = encode_pte(
      (uint64_t)kernel_p2 - kernel_offset, false, true, true, false);
    for (int i = 0; i < 512; ++i)
      kernel_p2[i] = encode_pte(
        (uint64_t)kernel_p1s + 4096 * i - kernel_offset, false, true, true, false);
    for (int i = 0; i < 512 * 512; ++i)
      kernel_p1s[i] = 0;
  }

  void map_kernel_page(
      uint64_t paddr, uint64_t vaddr, bool write, bool execute) {
    uint64_t i = (vaddr - kernel_vram_start) / 4096;
    kernel_p1s[i] = encode_pte(paddr, false, write, execute, false);
  }

  void unmap_kernel_page(uint64_t vaddr) {
    uint64_t i = (vaddr - kernel_vram_start) / 4096;
    kernel_p1s[i] = 0;
    asm volatile (
      "mov %%cr3, %%rax\nmov %%rax, %%cr3" ::: "%rax"
    );
  }

  static uint64_t take_pram_page() {
    for (uint64_t i = 0; i < pram_pages / 64; ++i)
      if (~pram_usage_bitmap[i] != 0)
        for (int j = 0; j < 64; ++j)
          if (~pram_usage_bitmap[i] & (1ULL << j)) {
            pram_usage_bitmap[i] |= (1ULL << j);
            return 4096 * (i * 64 + j);
          }
    //TODO: handle error
    return 0;
  }

  uint64_t take_2mib_pram_page() {
    for (uint64_t i = 0; i < pram_pages / 512; ++i) {
      for (int j = 0; j < 8; ++j)
        if (pram_usage_bitmap[i * 8 + j] != 0)
          goto next_i;
      for (int j = 0; j < 8; ++j)
        pram_usage_bitmap[i * 8 + j] = ~0ULL;
      return 0x200000 * i;
    next_i:
      ;
    }
    //TODO: handle error
    return 0;
  }

  void map_kernel_stacks() {
    for (uint64_t vaddr = syscall_stack_bottom;
         vaddr < syscall_stack_top; vaddr += 4096)
      map_kernel_page(take_pram_page(), vaddr, true, false);
    for (uint64_t vaddr = interrupt_stack_bottom;
         vaddr < interrupt_stack_top; vaddr += 4096)
      map_kernel_page(take_pram_page(), vaddr, true, false);
  }

  uint64_t find_unmapped_vram_region(uint64_t page_count) {
    uint64_t start = 0;
    uint64_t len = 0;
    for (uint64_t i = 0; i < kernel_vram_pages; ++i)
      if (kernel_p1s[i] == 0) {
        ++len;
        if (len == page_count)
          return start * 4096 + kernel_vram_start;
      }
      else {
        start = i + 1;
        len = 0;
      }
    //TODO: handle error
    return 0;
  }

  void *map_new_kernel_pages(uint64_t count) {
    uint64_t vaddr = find_unmapped_vram_region(count);
    for (uint64_t i = 0; i < count; ++i)
      map_kernel_page(take_pram_page(), vaddr + i * 4096, true, false);
    return (void *)vaddr;
  }

  void map_new_kernel_page(uint64_t &vaddr_out, uint64_t &paddr_out) {
    vaddr_out = find_unmapped_vram_region(1);
    paddr_out = take_pram_page();
    map_kernel_page(paddr_out, vaddr_out, true, false);
  }

  uint64_t get_used_vram_page_count() {
    uint64_t count = 0;
    for (uint64_t i = 0; i < kernel_vram_pages; ++i)
      if (kernel_p1s[i] != 0)
        ++count;
    return count;
  }

  uint64_t get_free_pram_page_count() {
    uint64_t used_count = 0;
    for (uint64_t i = 0; i < pram_pages / 64; ++i)
      for (uint64_t j = 0; j < 64; ++j)
        used_count += (pram_usage_bitmap[i] >> j) & 1;
    return pram_pages - used_count;
  }

}