#include <stdint.h>
#include <stddef.h>
#include <limits.h>
#include "paging/paging.h"
#include "hal/hal.h"
#include "bootinfo/bootinfo.h"
#include "util/util.h"
#include "kprintf.h"

enum {
  PF_PRESENT      = 1<<0,
  PF_READWRITE    = 1<<1,
  PF_USER         = 1<<2,
  PF_WRITETHROUGH = 1<<3,
  PF_NOCACHE      = 1<<4,
  PF_ACCESS       = 1<<5,
  PF_DIRTY        = 1<<6,
  PF_PAGESIZE     = 1<<7,
  PF_ATTRTAB      = 1<<7,
  PF_GLOBAL       = 1<<8,
  PF_SHARED       = 1<<9,
  PF_CACHE_WC     = PF_ATTRTAB | PF_WRITETHROUGH,
};

#define PGSHIFT_PML4E   39
#define PGSHIFT_PDPTE   30
#define PGSHIFT_PDE     21
#define PGSHIFT_PTE     12
#define PGMASK_ENTRY    0x1ff
#define PGMASK_OFFSET   0x3ff

#define BITS_TO_VIRT_ADDR(pml4_index, pdpte_index, pd_index, pt_index)         \
  (((uint64_t)pml4_index << PGSHIFT_PML4E) |                                   \
   ((uint64_t)pdpte_index << PGSHIFT_PDPTE) |                                  \
   ((uint64_t)pd_index << PGSHIFT_PDE) | ((uint64_t)pt_index << PGSHIFT_PTE))

// Workaround repeated characters
#define AMD64_MM_STRIPSX(a) ((uintptr_t)(a) & 0xFFFFFFFFFFFF)
#define AMD64_MM_ADDRSX(a)                                                     \
  (((uintptr_t)(a) & (1ULL << 47)) ? (0xFFFFFF0000000000 | ((uintptr_t)(a)))   \
                                   : ((uintptr_t)(a)))

// Virtual address' macros
#define PML4E(a) (((a) >> PGSHIFT_PML4E) & PGMASK_ENTRY)
#define PDPTE(a) (((a) >> PGSHIFT_PDPTE) & PGMASK_ENTRY)
#define PDE(a) (((a) >> PGSHIFT_PDE) & PGMASK_ENTRY)
#define PTE(a) (((a) >> PGSHIFT_PTE) & PGMASK_ENTRY)

#define PTE_ADDR_MASK 0x000ffffffffff000
#define PTE_GET_ADDR(VALUE) ((VALUE) & PTE_ADDR_MASK)
#define PTE_GET_FLAGS(VALUE) ((VALUE) & ~PTE_ADDR_MASK)

#define PAGE_MASK(x) ((1 << (x)) - 1)

#define USER_STACK_PAGES 2048
#define PAGE_SIZE 0x1000
#define PAGE_SIZE_LARGE 0x200000
#define PAGE_SIZE_HUGE 0x40000000

#define P_PHYs_ADD(x) ((x) & ~0xfff)

uint64_t *global_page_dir = NULL;

void paging_init(void) {
  uint64_t pdphys = 0;
  asm volatile("movq %%cr3, %0" : "=r"(pdphys));
  if (!pdphys) {
    ERR("paging", "could not get default page dir\n");
    hal_hang();
  }

  uint64_t pdvirt = pdphys + BOOT_INFO.hhdm_off;
  global_page_dir = (uint64_t *)pdvirt;

  LOG("paging", "global_page_dir = %p\n", global_page_dir);
}

size_t _paging_virt2phys(uint64_t *pgdir, size_t virtaddr) {
  if (!pgdir) return 0;

  if (virtaddr >= BOOT_INFO.hhdm_off
   && virtaddr <= (BOOT_INFO.hhdm_off + MAX(BOOT_INFO.memmap_total, UINT32_MAX))) {
    return virtaddr - BOOT_INFO.hhdm_off;
  }

  size_t virtaddr_init = virtaddr;
  virtaddr &= ~0xfff;

  virtaddr = AMD64_MM_STRIPSX(virtaddr);

  uint32_t pml4idx = PML4E(virtaddr);
  uint32_t pdpidx = PDPTE(virtaddr);
  uint32_t pdidx = PDE(virtaddr);
  uint32_t ptidx = PTE(virtaddr);

  if (!(pgdir[pml4idx] & PF_PRESENT)) {
    return 0;
  }

  size_t *pdp = (size_t *)(PTE_GET_ADDR(pgdir[pml4idx]) + BOOT_INFO.hhdm_off);

  if (!(pdp[pdpidx] & PF_PRESENT)) {
    return 0;
  }

  size_t *pd = (size_t *)(PTE_GET_ADDR(pdp[pdpidx]) + BOOT_INFO.hhdm_off);

  if (!(pd[pdidx] & PF_PRESENT)) {
    return 0;
  }

  size_t *pt = (size_t *)(PTE_GET_ADDR(pd[pdidx]) + BOOT_INFO.hhdm_off);

  if (pt[ptidx] & PF_PRESENT) {
    return (size_t)(PTE_GET_ADDR(pt[ptidx]) + ((size_t)virtaddr_init & 0xfff));
  }

  return 0;
}

size_t paging_virt2phys(size_t virtaddr) {
  return _paging_virt2phys(global_page_dir, virtaddr);
}