#include <zinnia/alloc.h>
#include <zinnia/assert.h>
#include <zinnia/errno.h>
#include <zinnia/mman.h>
#include <zinnia/pmap.h>
#include <zinnia/utils.h>
#include <zinnia/vma.h>
#include <stdckdint.h>

// Try to create a new, empty virtual address space.
errno_t vma_new(struct vma** out) {
    debug_assert(out != nullptr, "Missing out");

    errno_t res;
    if (!out)
        return EINVAL;

    struct vma* vmspace = mem_alloc(sizeof(struct vma), 0);
    if (!vmspace)
        return ENOMEM;

    res = pmap_new_user(&vmspace->pmap, 0);
    if (res)
        goto err1;

    tailq_init(&vmspace->mappings);

    *out = vmspace;
    return 0;

err1:
    mem_free(vmspace);
err0:
    return res;
}

// Delete a virtual address space.
void vma_free(struct vma* v) {
    struct vma_mapping* mapping;
    while ((mapping = tailq_first(&v->mappings)) != nullptr) {
        tailq_remove(&v->mappings, mapping, link);
        vmo_unref(mapping->object);
        mem_free(mapping);
    }
}

errno_t vma_unmap(struct vma* v, uintptr_t vaddr, size_t size) {
    panic("TODO");
}

errno_t vma_protect(struct vma* v, uintptr_t vaddr, size_t size, enum prot_flags flags) {
    panic("TODO");
}

static struct vma_mapping* find_mapping(struct vma* v, size_t page) {
    for (struct vma_mapping* m = tailq_first(&v->mappings); m; m = tailq_next(m, link)) {
        if (page >= m->start_page && page < m->end_page)
            return m;
    }
    return nullptr;
}

static bool can_access(enum prot_flags prot, const struct vma_fault_info* info) {
    if (info->caused_by_write)
        return prot & PROT_WRITE;
    if (info->caused_by_fetch)
        return prot & PROT_EXEC;
    return prot & (PROT_READ | PROT_WRITE);
}

static enum pte_flags prot_to_pte(struct vma* v, enum prot_flags prot) {
    enum pte_flags pf = 0;
    if (prot & PROT_READ)
        pf |= PTE_READ;
    if (prot & PROT_WRITE)
        pf |= PTE_WRITE;
    if (prot & PROT_EXEC)
        pf |= PTE_EXEC;
    if (v->pmap.is_user)
        pf |= PTE_USER;
    return pf;
}

bool vma_page_fault(struct vma* v, const struct vma_fault_info* info) {
    const size_t page_size = arch_mem_page_size();
    const size_t faulty_page = info->addr / page_size;

    spin_lock(&v->lock);
    struct vma_mapping* mapping = find_mapping(v, faulty_page);
    if (!mapping || !can_access(mapping->prot, info)) {
        spin_unlock(&v->lock);
        return false;
    }

    struct vmo* object = mapping->object;
    const size_t object_page = (faulty_page - mapping->start_page) + mapping->offset_page;
    const enum prot_flags prot = mapping->prot;
    vmo_ref(object);
    spin_unlock(&v->lock);

    struct page* page;
    if (object->get_page(object, object_page, &page)) {
        vmo_unref(object);
        return false;
    }

    spin_lock(&v->lock);
    mapping = find_mapping(v, faulty_page);
    const bool still_valid = mapping && mapping->object == object &&
        (faulty_page - mapping->start_page) + mapping->offset_page == object_page && can_access(mapping->prot, info);

    if (still_valid) {
        errno_t status =
            pmap_map(&v->pmap, faulty_page * page_size, vm_page_addr(page), prot_to_pte(v, prot), CACHE_ANY);
        assert(status == 0, "Failed to map a demand-loaded page");
    }
    spin_unlock(&v->lock);

    vmo_unref(object);
    return true;
}

errno_t vma_map(
    struct vma* vma,
    struct vmo* vmo,
    uintptr_t addr,
    size_t len,
    enum prot_flags flags,
    uintptr_t vmo_offset
) {
    if (!vma || !vmo)
        return EINVAL;

    uintptr_t end;
    if (ckd_add(&end, addr, len))
        return ENOMEM;

    const size_t page_size = arch_mem_page_size();
    if (addr % page_size != vmo_offset % page_size)
        return EINVAL;

    struct vma_mapping* mapping = mem_alloc(sizeof(struct vma_mapping), 0);
    if (!mapping)
        return ENOMEM;

    vmo_ref(vmo);
    mapping->start_page = addr / page_size;
    mapping->end_page = ALIGN_UP(end, page_size) / page_size;
    mapping->offset_page = vmo_offset / page_size;
    mapping->object = vmo;
    mapping->prot = flags;

    spin_lock(&vma->lock);
    tailq_insert_head(&vma->mappings, mapping, link);
    spin_unlock(&vma->lock);

    return 0;
}