#include "page.h"
#include <stdint.h>
#include <string.h>
#include <stdio.h>
#include "filesys/file.h"
#include "threads/pte.h"
#include "threads/malloc.h"
#include "threads/palloc.h"
#include "devices/swap.h"
#include "userprog/process.h"
#include "userprog/pagedir.h"
#include "vm/frame.h"

#define SWAP_FLAG_BIT 9
#define ADDR_START_BIT 12

/* Hashing function needed for the SPT table. Returns a hash for an entry,
   based on its upage. */
unsigned
page_hash (const struct hash_elem *e, UNUSED void *aux)
{
  struct page_entry *page = hash_entry (e, struct page_entry, elem);
  return hash_ptr (page->upage);
}

/* Comparator function for the SPT table. Compares two entries based on their
   upages. */
bool
page_less (const struct hash_elem *a_, const struct hash_elem *b_,
           void *aux UNUSED)
{
  const struct page_entry *a = hash_entry (a_, struct page_entry, elem);
  const struct page_entry *b = hash_entry (b_, struct page_entry, elem);

  return a->upage < b->upage;
}

static void page_flag_swap (uint32_t *pte, bool set);
static void page_set_swap (struct thread *owner, uint32_t *pte,
                           size_t swap_slot);

// TODO: Deal with NULL malloc returns
/* Swap out 'owner' process's 'upage' stored at 'kpage'. Then, allocate and
   insert a new page entry into the user process thread's SPT representing
   this swapped out page. */
struct page_entry *
page_insert_swapped (void *upage, void *kpage, struct thread *owner)
{
  /* 1. Initialize swapped page entry. */
  struct page_entry *page = page_get (upage);
  if (page == NULL)
    {
      page = malloc (sizeof (struct page_entry));
      if (page == NULL)
        return NULL;
      page->upage = upage;
      lock_init (&page->lock);
      hash_insert (&owner->pages, &page->elem);
    }

  /* Mark page as 'swapped' and flag the page directory as having
     been modified *before* eviction begins to prevent the owner of the
     victim page from accessing/modifying it mid-eviction. */
  /* TODO: We need to stop the process from destroying pagedir mid-eviction,
    as this could render the page table entry invalid. */
  uint32_t *pte = lookup_page (owner->pagedir, upage, false);

  page_flag_swap (pte, true);
  lock_acquire (&page->lock);
  pagedir_clear_page (owner->pagedir, upage);

  size_t swap_slot = swap_out (kpage);
  page_set_swap (owner, pte, swap_slot);

  lock_release (&page->lock);
  return page;
}

/* Allocate and insert a new page entry into the user process thread's
   SPT representing a file page. */
struct page_entry *
page_insert_file (struct file *file, off_t ofs, void *upage,
                  uint32_t read_bytes, uint32_t zero_bytes, bool writable,
                  enum page_type type)
{
  struct page_entry *page = malloc(sizeof (struct page_entry));
  if (page == NULL)
    return NULL;

  page->type = type;
  page->file = file;
  page->offset = ofs;
  page->upage = upage;
  page->read_bytes = read_bytes;
  page->zero_bytes = zero_bytes;
  page->writable = writable;
  lock_init (&page->lock);

  hash_insert (&thread_current ()->pages, &page->elem);
  return page;
}

/* Gets a page_entry from the starting address of the page. Returns NULL if no
   such page_entry exists in the hash map.*/
struct page_entry *
page_get (void *upage)
{
  struct page_entry fake_page_entry;
  fake_page_entry.upage = upage;

  struct hash_elem *e
    = hash_find (&thread_current ()->pages, &fake_page_entry.elem);
  if (e == NULL)
      return NULL;

  return hash_entry (e, struct page_entry, elem);
}

bool
page_load_file (struct page_entry *page, bool writable)
{
  /* Allocate a frame for the page. If a frame allocation fails, then
     frame_alloc should try to evict a page. If it is still NULL, the OS
     panics as this should not happen if eviction is working correctly. */
  struct thread *t = thread_current ();
  void *frame = frame_alloc (PAL_USER, page->upage, t);
  pagedir_set_accessed (t->pagedir, page->upage, true);
  if (frame == NULL)
    PANIC ("Could not allocate a frame to load page into memory.");

  /* Map the page to the frame. */
  if (!install_page (page->upage, frame, writable))
    {
      frame_free (frame);
      return false;
    }

  /* Move the file pointer to the correct location in the file. Then, read the
     data from the file into the frame. Checks that we were able to read the
     expected number of bytes. */
  file_seek (page->file, page->offset);
  if (file_read (page->file, frame, page->read_bytes) != (int) page->read_bytes)
  {
    frame_free (frame);
    return false;
  }

  /* Zero out the remaining bytes in the frame. */
  memset (frame + page->read_bytes, 0, page->zero_bytes);

  /* Mark the page as loaded successfully. */
  return true;
}

/* Function to clean up a page_entry. Given the elem of that page_entry, frees
   the page_entry itself. */
void
page_cleanup (struct hash_elem *e, void *aux UNUSED)
{
  free (hash_entry (e, struct page_entry, elem));
}

/* Flags the provided page table entry as representing a swapped out page. */
void
page_flag_swap (uint32_t *pte, bool set)
  {
    if (set)
      *pte |= (1 << SWAP_FLAG_BIT);
    else
      *pte &= ~(1 << SWAP_FLAG_BIT);
  }

/* Sets the address bits of the page table entry to the provided swap slot
   value. To be used for later retrieval of the swap slot when page faulting. */
static void
page_set_swap (struct thread *owner, uint32_t *pte, size_t swap_slot)
{
  /* Store the provided swap slot in the address bits of the page table
     entry, truncating excess bits. */
  *pte |= (1 << SWAP_FLAG_BIT);
  uint32_t swap_slot_bits = (swap_slot << ADDR_START_BIT) & PTE_ADDR;
  *pte = (*pte & PTE_FLAGS) | swap_slot_bits;

  invalidate_pagedir (owner->pagedir);
}

/* Returns true iff the page with user address 'upage' owned by 'owner'
   is flagged to be in the swap disk via the owner's page table. */
bool
page_in_swap (struct thread *owner, void *upage)
{
   uint32_t *pte = lookup_page (owner->pagedir, upage, false);
   return pte != NULL &&
          (*pte & (1 << SWAP_FLAG_BIT)) != 0;
}

/* Given that the page with user address 'upage' owned by 'owner' is flagged
   to be in the swap disk via the owner's page table, returns its stored
   swap slot and marks the PTE as not being in swap. */
size_t
page_get_swap (struct thread *owner, void *upage)
{
   uint32_t *pte = lookup_page (owner->pagedir, upage, false);

   ASSERT (pte != NULL);
   ASSERT ((*pte & PTE_P) == 0);

   /* Masks the address bits and returns truncated value. */
   page_flag_swap (pte, false);
   return ((*pte & PTE_ADDR) >> ADDR_START_BIT);
}