#include <amd64/apic.h>
#include <amd64/intr_defs.h>
#include <aux/compiler.h>
#include <irq/irq.h>
#include <libk/std.h>
#include <libk/string.h>
#include <limine/requests.h>
#include <mm/pmm.h>
#include <sync/spin_lock.h>
#include <sys/debug.h>
#include <sys/mm.h>
#include <sys/smp.h>

#define AMD64_PG_PRESENT (1 << 0)
#define AMD64_PG_RW      (1 << 1)
#define AMD64_PG_USER    (1 << 2)
#define AMD64_PG_HUGE    (1 << 7)

/* Auxilary struct for page directory walking */
struct pg_index {
  uint16_t pml4, pml3, pml2, pml1;
} PACKED;

/* Kernel page directory */
static struct pd kernel_pd;
static spin_lock_t kernel_pd_lock;

void mm_kernel_lock (spin_lock_ctx_t* ctx) { spin_lock (&kernel_pd_lock, ctx); }

void mm_kernel_unlock (spin_lock_ctx_t* ctx) { spin_lock (&kernel_pd_lock, ctx); }

/* Get current value of CR3 register */
static uintptr_t amd64_current_cr3 (void) {
  uintptr_t cr3;
  __asm__ volatile ("movq %%cr3, %0" : "=r"(cr3)::"memory");
  return cr3;
}

/* Load kernel CR3 as current CR3 */
void amd64_load_kernel_cr3 (void) {
  uintptr_t cr3 = amd64_current_cr3 ();

  if (cr3 != kernel_pd.cr3_paddr) {
    __asm__ volatile ("movq %0, %%cr3" ::"r"(kernel_pd.cr3_paddr) : "memory");
  }
}

struct pd* mm_get_kernel_pd (void) { return &kernel_pd; }

/* Extract PML info from virtual address */
static struct pg_index amd64_mm_page_index (uint64_t vaddr) {
  struct pg_index ret;

  ret.pml4 = ((vaddr >> 39) & 0x1FF);
  ret.pml3 = ((vaddr >> 30) & 0x1FF);
  ret.pml2 = ((vaddr >> 21) & 0x1FF);
  ret.pml1 = ((vaddr >> 12) & 0x1FF);

  return ret;
}

/* Walk paging tables and allocate necessary structures along the way */
static uint64_t* amd64_mm_next_table (uint64_t* table, uint64_t entry_idx, bool alloc) {
  uint64_t entry = table[entry_idx];
  physaddr_t paddr;

  struct limine_hhdm_response* hhdm = limine_hhdm_request.response;

  if (entry & AMD64_PG_PRESENT) {
    if (entry & AMD64_PG_HUGE)
      return NULL;

    paddr = entry & ~0xFFFULL;
  } else {
    if (!alloc)
      return NULL;

    paddr = pmm_alloc (1);

    if (paddr == PMM_ALLOC_ERR)
      return NULL;

    memset ((void*)((uintptr_t)hhdm->offset + (uintptr_t)paddr), 0, PAGE_SIZE);
    table[entry_idx] = paddr | AMD64_PG_PRESENT | AMD64_PG_RW | AMD64_PG_USER;
  }

  return (uint64_t*)((uintptr_t)hhdm->offset + (uintptr_t)paddr);
}

static bool amd64_mm_is_table_empty (uint64_t* table) {
  for (size_t i = 0; i < 512; i++) {
    if (table[i] & AMD64_PG_PRESENT)
      return false;
  }
  return true;
}

/* Convert generic memory management subsystem flags into AMD64-specific flags */
static uint64_t amd64_mm_resolve_flags (uint32_t generic) {
  uint64_t flags = 0;

  flags |= ((generic & MM_PG_PRESENT) ? AMD64_PG_PRESENT : 0);
  flags |= ((generic & MM_PG_RW) ? AMD64_PG_RW : 0);
  flags |= ((generic & MM_PG_USER) ? AMD64_PG_USER : 0);

  return flags;
}

/* Reload the current CR3 value ON A LOCAL CPU */
static void amd64_reload_cr3 (void) {
  uint64_t cr3;
  __asm__ volatile ("movq %%cr3, %0; movq %0, %%cr3" : "=r"(cr3)::"memory");
}

/* Map physical address to virtual address with flags. TLB needs to be flushed afterwards. */
void mm_map_page (struct pd* pd, uintptr_t paddr, uintptr_t vaddr, uint32_t flags) {
  struct limine_hhdm_response* hhdm = limine_hhdm_request.response;

  uint64_t amd64_flags = amd64_mm_resolve_flags (flags);

  uint64_t* pml4 = (uint64_t*)(pd->cr3_paddr + (uintptr_t)hhdm->offset);
  struct pg_index pg_index = amd64_mm_page_index (vaddr);

  uint64_t* pml3 = amd64_mm_next_table (pml4, pg_index.pml4, true);
  if (pml3 == NULL)
    return;

  uint64_t* pml2 = amd64_mm_next_table (pml3, pg_index.pml3, true);
  if (pml2 == NULL)
    return;

  uint64_t* pml1 = amd64_mm_next_table (pml2, pg_index.pml2, true);
  if (pml1 == NULL)
    return;

  uint64_t* pte = &pml1[pg_index.pml1];

  *pte = ((paddr & ~0xFFFULL) | (amd64_flags & 0x7ULL));
}

/* Map a page into kernel page directory */
void mm_map_kernel_page (uintptr_t paddr, uintptr_t vaddr, uint32_t flags) {
  mm_map_page (&kernel_pd, paddr, vaddr, flags);
  amd64_reload_cr3 ();
}

/* Unmap a virtual address. TLB needs to be flushed afterwards */
void mm_unmap_page (struct pd* pd, uintptr_t vaddr) {
  struct limine_hhdm_response* hhdm = limine_hhdm_request.response;

  uint64_t* pml4 = (uint64_t*)(pd->cr3_paddr + (uintptr_t)hhdm->offset);
  struct pg_index pg_index = amd64_mm_page_index (vaddr);

  uint64_t* pml3 = amd64_mm_next_table (pml4, pg_index.pml4, false);
  if (pml3 == NULL)
    return;

  uint64_t* pml2 = amd64_mm_next_table (pml3, pg_index.pml3, false);
  if (pml2 == NULL)
    return;

  uint64_t* pml1 = amd64_mm_next_table (pml2, pg_index.pml2, false);
  if (pml1 == NULL)
    return;

  uint64_t* pte = &pml1[pg_index.pml1];

  if ((*pte) & AMD64_PG_PRESENT)
    *pte = 0;

  if (amd64_mm_is_table_empty (pml1)) {
    uintptr_t pml1_phys = pml2[pg_index.pml2] & ~0xFFFULL;
    pmm_free (pml1_phys, 1);
    pml2[pg_index.pml2] = 0;

    if (amd64_mm_is_table_empty (pml2)) {
      uintptr_t pml2_phys = pml3[pg_index.pml3] & ~0xFFFULL;
      pmm_free (pml2_phys, 1);
      pml3[pg_index.pml3] = 0;

      if (amd64_mm_is_table_empty (pml3)) {
        uintptr_t pml3_phys = pml4[pg_index.pml4] & ~0xFFFULL;
        pmm_free (pml3_phys, 1);
        pml4[pg_index.pml4] = 0;
      }
    }
  }
}

/* Unmap a page from kernel page directory */
void mm_unmap_kernel_page (uintptr_t vaddr) {
  mm_unmap_page (&kernel_pd, vaddr);
  amd64_reload_cr3 ();
}

/* Allocate a userspace-ready page directory */
uintptr_t mm_alloc_user_pd_phys (void) {
  struct limine_hhdm_response* hhdm = limine_hhdm_request.response;

  physaddr_t cr3 = pmm_alloc (1);
  if (cr3 == PMM_ALLOC_ERR)
    return 0;

  uint8_t* vu_cr3 = (uint8_t*)((uintptr_t)hhdm->offset + cr3);
  memset ((void*)vu_cr3, 0, PAGE_SIZE / 2);

  uint8_t* vk_cr3 = (uint8_t*)((uintptr_t)hhdm->offset + (uintptr_t)kernel_pd.cr3_paddr);

  memcpy (&vu_cr3[PAGE_SIZE / 2], &vk_cr3[PAGE_SIZE / 2], PAGE_SIZE / 2);

  return cr3;
}

bool mm_validate (struct pd* pd, uintptr_t vaddr) {
  struct limine_hhdm_response* hhdm = limine_hhdm_request.response;
  bool ret = false;

  uint64_t* pml4 = (uint64_t*)(pd->cr3_paddr + (uintptr_t)hhdm->offset);
  struct pg_index pg_index = amd64_mm_page_index (vaddr);

  uint64_t* pml3 = amd64_mm_next_table (pml4, pg_index.pml4, false);
  if (pml3 == NULL)
    goto done;

  uint64_t* pml2 = amd64_mm_next_table (pml3, pg_index.pml3, false);
  if (pml2 == NULL)
    goto done;

  uint64_t* pml1 = amd64_mm_next_table (pml2, pg_index.pml2, false);
  if (pml1 == NULL)
    goto done;

  uint64_t pte = pml1[pg_index.pml1];
  ret = (pte & AMD64_PG_PRESENT) != 0;

done:
  return ret;
}

bool mm_validate_buffer (struct pd* pd, uintptr_t vaddr, size_t size) {
  bool ok = true;

  for (size_t i = 0; i < size; i++) {
    ok = mm_validate (pd, vaddr + i);
    if (!ok)
      goto done;
  }

done:
  return ok;
}

uintptr_t mm_p2v (struct pd* pd, uintptr_t paddr) {
  struct limine_hhdm_response* hhdm = limine_hhdm_request.response;
  uintptr_t ret = 0;

  uint64_t* pml4 = (uint64_t*)(pd->cr3_paddr + (uintptr_t)hhdm->offset);

  for (size_t i4 = 0; i4 < 512; i4++) {
    if (!(pml4[i4] & AMD64_PG_PRESENT))
      continue;

    uint64_t* pml3 = (uint64_t*)((uintptr_t)hhdm->offset + (pml4[i4] & ~0xFFFULL));
    for (size_t i3 = 0; i3 < 512; i3++) {
      if (!(pml3[i3] & AMD64_PG_PRESENT))
        continue;

      uint64_t* pml2 = (uint64_t*)((uintptr_t)hhdm->offset + (pml3[i3] & ~0xFFFULL));
      for (size_t i2 = 0; i2 < 512; i2++) {
        if (!(pml2[i2] & AMD64_PG_PRESENT))
          continue;

        uint64_t* pml1 = (uint64_t*)((uintptr_t)hhdm->offset + (pml2[i2] & ~0xFFFULL));
        for (size_t i1 = 0; i1 < 512; i1++) {
          if ((pml1[i1] & AMD64_PG_PRESENT) && ((pml1[i1] & ~0xFFFULL) == (paddr & ~0xFFFULL))) {
            struct pg_index idx = {i4, i3, i2, i1};
            ret = (((uint64_t)idx.pml4 << 39) | ((uint64_t)idx.pml3 << 30) |
                   ((uint64_t)idx.pml2 << 21) | ((uint64_t)idx.pml1 << 12) | (paddr & 0xFFFULL));
            goto done;
          }
        }
      }
    }
  }

done:
  return ret;
}

uintptr_t mm_v2p (struct pd* pd, uintptr_t vaddr) {
  struct limine_hhdm_response* hhdm = limine_hhdm_request.response;
  uintptr_t ret = 0;

  uint64_t* pml4 = (uint64_t*)(pd->cr3_paddr + (uintptr_t)hhdm->offset);
  struct pg_index pg_index = amd64_mm_page_index (vaddr);

  uint64_t* pml3 = amd64_mm_next_table (pml4, pg_index.pml4, false);
  if (pml3 == NULL)
    goto done;

  uint64_t* pml2 = amd64_mm_next_table (pml3, pg_index.pml3, false);
  if (pml2 == NULL)
    goto done;

  uint64_t* pml1 = amd64_mm_next_table (pml2, pg_index.pml2, false);
  if (pml1 == NULL)
    goto done;

  uint64_t pte = pml1[pg_index.pml1];

  if (!(pte & AMD64_PG_PRESENT))
    goto done;

  ret = ((pte & ~0xFFFULL) | (vaddr & 0xFFFULL));

done:
  return ret;
}

/* Initialize essentials for the AMD64 memory management subsystem */
void mm_init (void) { kernel_pd.cr3_paddr = amd64_current_cr3 (); }