#include "page.h"
#include <string.h>
#include "filesys/file.h"
#include "filesys/filesys.h"
#include "threads/malloc.h"
#include "threads/palloc.h"
#include "threads/vaddr.h"
#include "userprog/pagedir.h"
#include "userprog/process.h"
#include "vm/frame.h"

static unsigned page_hash (const struct hash_elem *e, void *aux UNUSED);
static bool page_less (const struct hash_elem *a_, const struct hash_elem *b_,
                       void *aux UNUSED);
static struct shared_page_entry *shared_page_insert (struct file *file,
                                                     void *upage, void *frame);
static struct shared_file_entry *shared_file_get (struct file *file);
static struct shared_page_entry *shared_page_get (struct file *file,
                                                  void *upage);
static unsigned shared_page_hash (const struct hash_elem *e, void *aux UNUSED);
static bool shared_page_less (const struct hash_elem *a_,
                              const struct hash_elem *b_, void *aux UNUSED);

/* Initialise a thread's supplemental pages table. */
bool
init_pages (struct thread *t)
{
  return hash_init (&t->pages, page_hash, page_less, NULL);
}

/* 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, void *aux UNUSED)
{
  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;
}

/* Allocate and insert a new page entry into the thread's page table. */
struct page_entry *
page_insert (struct file *file, off_t ofs, void *upage, uint32_t read_bytes,
             uint32_t zero_bytes, bool writable, enum page_type type)
{
  /* If page exists, just update it. */
  struct page_entry *existing = page_get (upage);
  if (existing != NULL)
    {
      existing->read_bytes = read_bytes;
      existing->zero_bytes = zero_bytes;
      existing->writable = writable;
      return existing;
    }

  /* Otherwise allocate a new one. */
  struct page_entry *page = malloc(sizeof (struct page_entry));
  if (page == NULL)
    return NULL;
  page->file = file;
  page->offset = ofs;
  page->upage = upage;
  page->read_bytes = read_bytes;
  page->zero_bytes = zero_bytes;
  page->writable = writable;
  page->shared = false;
  page->type = type;
  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 (struct page_entry *page)
{
  /* If the page is read-only, we want to check if it is a shared page already
     loaded into memory. If it is, we can just map the page to the frame. */
  if (!page->writable)
    {
      lock_acquire (&shared_files_lock);
      struct shared_page_entry *shared_page =
        shared_page_get (page->file, page->upage);

      /* Mark page as shared and install the shared frame. */
      if (shared_page != NULL)
        {
          lock_release (&shared_files_lock);
          if (!install_page (page->upage, shared_page->frame, false))
            return false;
          page->shared = true;
          return true;
        }
    }

  /* 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. */
  void *frame = frame_alloc (PAL_USER, page->upage, thread_current ());
  if (frame == NULL)
    PANIC ("Could not allocate a frame to load page into memory.");

  /* Ensure page is not marked as shared while it doesn't exist in the
     shared_files table, to avoid memory leaks. */
  page->shared = false;

  /* Map the page to the frame. */
  if (!install_page (page->upage, frame, page->writable))
    goto fail;

  /* 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)
    goto fail;

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

  /* If the page is read-only, we need to add it to the shared pages table. */
  if (!page->writable)
    {
      struct shared_page_entry *shared_page
          = shared_page_insert (page->file, page->upage, frame);
      if (shared_page == NULL)
        goto fail;
      lock_release (&shared_files_lock);
      page->shared = true;
    }

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

fail:
  if (!page->writable)
    lock_release (&shared_files_lock);
  frame_free (frame);
  return false;
}

/* 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)
{
  struct page_entry *page = hash_entry (e, struct page_entry, elem);

  /* If page is shared then mark it as not present to avoid it being freed. */
  uint32_t *pd = thread_current ()->pagedir;
  if (pd != NULL && page->shared)
    pagedir_clear_page (pd, page->upage);
  free (page);
}

/* Initialise the shared files and table and lock. */
bool
shared_files_init ()
{
  lock_init (&shared_files_lock);
  if (!hash_init (&shared_files, shared_file_hash, shared_file_less, NULL))
    PANIC ("Failed to initialise shared_files table.");
}

/* Hashing function needed for the shared_file table. Returns a hash for an
   entry based on its file pointer. */
unsigned
shared_file_hash (const struct hash_elem *e, void *aux UNUSED)
{
  return file_hash (hash_entry (e, struct shared_file_entry, elem)->file);
}

/* Less function needed for the shared_file table. */
bool
shared_file_less (const struct hash_elem *a_, const struct hash_elem *b_,
                  void *aux UNUSED)
{
  const struct shared_file_entry *a = hash_entry (a_, struct shared_file_entry,
                                                  elem);
  const struct shared_file_entry *b = hash_entry (b_, struct shared_file_entry,
                                                  elem);

  return !file_compare (a->file, b->file);
}

/* Hashing function needed for the shared pages table. Returns a hash for an
   entry based on its user virtual address (upage) pointer. */
static unsigned
shared_page_hash (const struct hash_elem *e, void *aux UNUSED)
{
  return hash_ptr (hash_entry (e, struct shared_page_entry, elem)->upage);
}

/* Less function needed for the shared pages table. */
static bool
shared_page_less (const struct hash_elem *a_, const struct hash_elem *b_,
                  void *aux UNUSED)
{
  const struct shared_page_entry *a = hash_entry (a_, struct shared_page_entry,
                                                  elem);
  const struct shared_page_entry *b = hash_entry (b_, struct shared_page_entry,
                                                  elem);

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

static struct shared_file_entry *
shared_file_get (struct file *file)
{
  struct shared_file_entry fake_shared_file;
  fake_shared_file.file = file;

  struct hash_elem *e = hash_find (&shared_files, &fake_shared_file.elem);

  if (e == NULL)
    return NULL;

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

/* Gets a shared_page_entry from the shared_pages table using the file and upage
   of the page. Returns NULL if no such page_entry exists in the hash map.*/
struct shared_page_entry *
shared_page_get (struct file *file, void *upage)
{
  /* Search first for the file within the shared_pages structure */
  struct shared_file_entry *shared_file = shared_file_get (file);
  if (shared_file == NULL)
    return NULL;

  /* Search for the page within the shared_file's hash table */
  struct shared_page_entry fake_shared_page_entry;
  fake_shared_page_entry.upage = upage;

  struct hash_elem *e = hash_find (&shared_file->pages, &fake_shared_page_entry.elem);

  if (e == NULL)
    return NULL;

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

static struct shared_page_entry *
shared_page_insert (struct file *file, void *upage, void *frame)
{
  struct shared_file_entry *shared_file = shared_file_get (file);

  /* Allocate a new shared_page_entry first for easier error handling. */
  struct shared_page_entry *shared_page
    = malloc (sizeof (struct shared_page_entry));
  if (shared_page == NULL)
    return NULL;

  /* If shared file doesn't exist in table, also create it. */
  if (shared_file == NULL)
    {
      shared_file = malloc (sizeof (struct shared_file_entry));
      if (shared_file == NULL)
        {
          free (shared_page);
          return NULL;
        }

      shared_file->file = file;
      shared_file->ref_count = 0;
      if (!hash_init (&shared_file->pages, shared_page_hash, shared_page_less, NULL))
        {
          free (shared_page);
          free (shared_file);
          return NULL;
        }

      hash_insert (&shared_files, &shared_file->elem);
    }

  shared_page->upage = upage;
  shared_page->frame = frame;
  hash_insert (&shared_file->pages, &shared_page->elem);
  return shared_page;
}

/* Updates the 'owner' thread's page table entry for virtual address 'upage'
   to have a present bit of 0 and stores the specified swap slot value in the
   entry for later retrieval from disk. */
void
page_set_swap (struct thread *owner, void *upage, size_t swap_slot)
{

}

/* 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. Otherwise panics the kernel. */
size_t
page_get_swap (struct thread *owner, void *upage)
{
  return 0;
}