#include "cpen212vm.h"
#include <stdlib.h>
#include <string.h>

#define PAGE_SIZE 4096

typedef struct {
    size_t ppn;
    size_t reserved;
    bool accessed;
    bool user;
    bool executable;
    bool writable;
    bool readable;
    bool present;
    bool valid;
} page_table_entry_t;

typedef struct {
    page_table_entry_t *pte_entry;
    void *phys_mem;
    size_t num_phys_pages;
    FILE *swap;
    size_t num_swap_pages;
    size_t num_processes;
} vm_t;

#define VM_SIZE sizeof(vm_t)
paddr_t allocate_physical_page(vm_t *vm);

// description:
// - initializes a VM system
// arguments:
// - physmem: pointer to an area of at least 4096 * num_phys_pages bytes that models the physical memory
//   - no vm_* functions may access memory outside of [physmem, physmem+4096*num_phys_pages)
//   - all physical addresses are offsets from physmem (i.e., physical address is exactly 0 at physmem)
// - num_phys_pages: total number of 4096-byte physical pages available
//   - it is guaranteed that 4 <= num_phys_pages <= 1048576
//   - physical page 0 starts at physmem
// - swap
//   - if non-null: pointer to a swap file opened in read-write mode with size 4096 * num_swap_pages bytes
//   - if null: no swap space is available for this VM instance
// - num_swap_pages: total number of 4096-byte pages available in the swap file
//   - only relevant if swap is not null
//   - if swap is non-null, it is guaranteed that 2 <= num_swap_pages <= 67108864
// returns:
// - a handle that uniquely identifies this VM instance; this will be passed unchanged to other vm_* functions
//   multiple VM instances may co-exist
void *vm_init(void *phys_mem, size_t num_phys_pages, FILE *swap, size_t num_swap_pages) {
    
    if (num_phys_pages < 4 || num_phys_pages > 1048576) {
        return NULL;
    }
    if (swap != NULL && (num_swap_pages < 2 || num_swap_pages > 67108864)) {
        return NULL;
    }

    if (num_phys_pages * PAGE_SIZE < VM_SIZE) {
        return NULL;
    }

    // Init vm
    static char vm_char[VM_SIZE];
    static vm_t *vm = (vm_t*)vm_char;
    vm->phys_mem = phys_mem;
    vm->num_phys_pages = num_phys_pages;
    vm->swap = swap;
    vm->num_swap_pages = num_swap_pages;
    vm->num_processes = 0;

    return vm;
}

// description:
// - deinitializes a VM system
// - note that any swap files are *not* closed
// arguments:
// - vm: a VM system handle returned from vm_init
void vm_deinit(void *vm_ptr) {
    if (!vm_ptr) {
        return;
    }

    vm_t *vm = (vm_t *)vm_ptr;

    if (vm->swap) {
        fclose(vm->swap);
    }

    free(vm->pte_entry);
    vm->pte_entry = NULL;
    vm->phys_mem = NULL;
    vm->num_phys_pages = 0;
    vm->swap = NULL;
    vm->num_swap_pages = 0;
    free(vm);
}

// description:
// - creates a mapping for a new page in the virtual address space and map it to a physical page
// arguments:
// - vm: a VM system handle returned from vm_init
// - new_process: true iff there is no top-level page table for this process
//   - there may be up to 1000 (inclusive) separate processes with active top-level page tables
// - pt: physical address of the top-level page table of the process (relevant only if new_process is false)
// - addr: the virtual address on a page that is to be mapped (not necessarily the start of the page)
// - user: the page is accessible from user-level processes
// - exec: instructions may be fetched from this page
// - write: data may be written to this page
// - read: data may be read from this page
// returns:
// - the success status of the translation:
//   - VM_OK if the mapping succeeded
//   - VM_OUT_OF_MEM if no free pages remain in the physical memory and any relevant swap
//   - VM_OUT_OF_MEM if new_process is true and there are already 1000 active processes with top-level page tables
//   - VM_DUPLICATE if a mapping for this page already exists
//   - VM_BAD_IO if accessing the swap file failed
// - the physical address of the *top-level* page table for this process (relevant only if status is VM_OK)
vm_result_t vm_map_page(void *vm_ptr, bool new_process, paddr_t pt, vaddr_t addr,
                        bool user, bool exec, bool write, bool read) {
    vm_t *vm = (vm_t*) vm_ptr;
    vm_result_t result;

    // Check if virtual address is within bounds
    if (addr >= 4294967296) { //2^32
        result.status = VM_BAD_ADDR;
        return result;
    }

    // Check if virtual address is aligned
    if (addr % PAGE_SIZE != 0) {
        result.status = VM_BAD_ADDR;
        return result;
    }

    // Check if page is already mapped
    page_table_entry_t *pte = vm->pte_entry + (addr >> 12);
    if (pte->valid) {
        result.status = VM_DUPLICATE;
        return result;
    }

    // Check if enough physical memory
    size_t num_free_phys_pages = vm->num_phys_pages - vm->num_swap_pages;
    if (num_free_phys_pages == 0) {
        result.status = VM_OUT_OF_MEM;
        return result;
    }

    // Allocate a page
    size_t ppn;
    bool* phys_mem = (bool*) vm->phys_mem;
    for (ppn = 0; ppn < vm->num_phys_pages; ppn++) {
        if (!phys_mem[ppn]) {
            phys_mem[ppn] = true;
            break;
        }
    }

    // If this is new process, allocate a new top-level page table (WIP)
    if (new_process) {
        if (vm->num_processes == 1000) {

            result.status = VM_OUT_OF_MEM;
            return result;
        }
        pt = allocate_physical_page(vm);
        if (pt == 0) {
            result.status = VM_OUT_OF_MEM;
            return result;
        }
        vm->num_processes++;
    }

    // Map the page to physical page
    pte->ppn = ppn;
    pte->accessed = false;
    pte->user = user;
    pte->executable = exec;
    pte->writable = write;
    pte->readable = read;
    pte->present = true;
    pte->valid = true;

    result.status = VM_OK;
    result.addr = pt;
    return result;
}

// description:
// - removes the mapping for the page that contains virtual address addr
// - returns any unmapped pages and any page tables with no mappings to the free page pool
// arguments:
// - vm: a VM system handle returned from vm_init
// - pt: physical address of the top-level page table of the process
// - addr: the virtual address on a page that is to be unmapped (not necessarily the start of the page)
// returns:
// - the success status of the translation:
//   - VM_OK if the page was successfully unmapped
//   - VM_BAD_ADDR if this process has no mapping for virtual address addr
//   - VM_BAD_IO if accessing the swap file failed
vm_status_t vm_unmap_page(void *vm_ptr, paddr_t pt, vaddr_t addr) {
    vm_t *vm = (vm_t *) vm_ptr;

    // Check if virtual address is aligned
    if (addr % PAGE_SIZE != 0) {
        return VM_BAD_ADDR;
    }

    // Find page table entry corresponding to virtual address
    page_table_entry_t *pte = vm->pte_entry + (addr >> 12);

    // Check if page is mapped
    if (!pte->valid) {
        return VM_BAD_ADDR;
    }

    // Free corresponding physical page and invalidate page table entry
    size_t ppn = pte->ppn;
    bool* phys_mem = (bool*) vm->phys_mem;
    phys_mem[ppn] = false;
    pte->ppn = 0;
    pte->accessed = false;
    pte->user = false;
    pte->executable = false;
    pte->writable = false;
    pte->readable = false;
    pte->present = false;
    pte->valid = false;

    // Check for errors accessing swap
    if (vm->swap != NULL) {
        if (fwrite(vm->phys_mem + ppn * PAGE_SIZE, PAGE_SIZE, 1, vm->swap) != 1) {
            return VM_BAD_IO;
        }
    }

    return VM_OK;
}

// description:
// - translates a virtual address to a physical address if possible
// arguments:
// - vm: a VM system handle returned from vm_init
// - pt: physical address of the top-level page table of the accessing process
// - addr: the virtual address to translate
// - access: the access being made (instruction fetch, read, or write)
// - user: the access is a user-level access (i.e., not a kernel access)
// returns:
// - the success status of the translation:
//   - VM_OK if translation succeeded
//   - VM_BAD_ADDR if there is no translation for this address
//   - VM_BAD_PERM if permissions are not sufficient for the type / source of access requested
//   - VM_BAD_IO if accessing the swap file failed
// - the resulting physical address (relevant only if status is VM_OK)
vm_result_t vm_translate(void *vm_ptr, paddr_t pt, vaddr_t addr, access_type_t acc, bool user) {
    vm_t *vm = (vm_t *) vm_ptr;
    vm_result_t result = { .status = VM_BAD_ADDR };

    // Calculate page directory index, page table index and offset from virtual address
    size_t pdi = addr >> 22 & 0x3ff;
    size_t pti = (addr >> 12) & 0x3ff;
    size_t offset = addr & 0xfff;

    // Check if page directory entry is present
    page_table_entry_t *pde = &vm->pte_entry[pdi];
    if (!pde->present) {
        result.status = VM_BAD_ADDR;
        return result;
    }

    // Check if page table entry is present
    bool* phys_mem = (bool*) vm->phys_mem;
    page_table_entry_t *pte = (page_table_entry_t*) phys_mem[(pde->ppn << 10) + pti];
    if (!pte->present) {
        result.status = VM_BAD_ADDR;
        return result;
    }

    // Check if access is allowed
    bool is_exec = (acc == VM_EXEC);
    bool is_write = (acc == VM_WRITE);
    bool is_read = (acc == VM_READ);
    if (user) {
        if ((is_exec && !pte->executable) || (is_write && !pte->writable) || (is_read && !pte->readable)) {
            result.status = VM_BAD_PERM;
            return result;
        }
    }

    // Check if page is swapped out or not valid
    if (!pte->valid) {
        // Page is not valid or present, all bits except valid are reserved
        result.status = VM_BAD_ADDR;
        return result;
    } else if (!pte->present) {
        // Page is swapped out, but mapped in virtual address
        int err = fseeko(vm->swap, pte->reserved << 12, SEEK_SET);
        if (err != 0) {
            result.status = VM_BAD_IO;
            return result;
        }
        err = fread((char *)vm->phys_mem + (pte->ppn << 12), PAGE_SIZE, 1, vm->swap);
        if (err <= 0) {
            result.status = VM_BAD_IO;
            return result;
        }
        pte->present = true;
    }

    // Update accessed bit
    if (!pte->accessed) {
        pte->accessed = true;
    }

    // Calculate physical address
    pt = (pte->ppn << 12) | offset;

    // Return physical address and status
    result.status = VM_OK;
    result.addr = pt;
    return result;
}

// description:
// - resets the used bit for all pages used by the specified process in physical memory
// arguments:
// - vm: a VM system handle returned from vm_init
// - pt: physical address of the top-level page table of the process
void vm_reset_accessed(void *vm_ptr, paddr_t pt) {
    vm_t *vm = (vm_t *) vm_ptr;
    
    uintptr_t pt_addr = (uintptr_t) pt;
    page_table_entry_t *pte = (page_table_entry_t *) pt_addr;
    for (size_t i = 0; i < (PAGE_SIZE / sizeof(page_table_entry_t)); i++) {
        // Check if page is valid and present
        if (pte[i].valid && pte[i].present) {
            // Reset accessed bit
            pte[i].accessed = false;
        }
        
        // If page is a next-level page table, recurse
        if (pte[i].valid && !pte[i].present) {
            vm_reset_accessed(vm, pte[i].ppn << 12);
        }
    }
}

paddr_t allocate_physical_page(vm_t *vm) {
    bool* phys_mem = (bool*) vm->phys_mem;
    for (paddr_t ppn = 0; ppn < vm->num_phys_pages; ppn++) {
        if (!phys_mem[ppn]) {
            phys_mem[ppn] = true;
            return ppn * PAGE_SIZE;
        }
    }
    return 0;
}