#include <zinnia/alloc.h>
#include <zinnia/errno.h>
#include <zinnia/list.h>
#include <zinnia/utils.h>
#include <zinnia/vmo.h>

void vmo_ref(struct vmo* obj) {
    refcount_inc(&obj->rc);
}

void vmo_unref(struct vmo* obj) {
    if (refcount_dec(&obj->rc))
        obj->destroy(obj);
}

errno_t vmo_read(struct vmo* obj, uintptr_t offset, void* buf, size_t len, size_t* actual_read) {
    const size_t page_size = arch_mem_page_size();
    size_t progress = 0;
    errno_t status = 0;

    while (progress < len) {
        const size_t misalign = (progress + offset) % page_size;
        const size_t page_index = (progress + offset) / page_size;
        const size_t copy_size = MIN(page_size - misalign, len - progress);

        struct page* p;
        status = obj->get_page(obj, page_index, &p);
        if (status != 0)
            break;

        uintptr_t page_addr = vm_page_addr(p);

        memcpy(buf + progress, HHDM_PTR(page_addr) + misalign, copy_size);
        progress += copy_size;
    }

    if (actual_read)
        *actual_read = progress;
    return status;
}

errno_t vmo_write(struct vmo* obj, uintptr_t offset, const void* buf, size_t len, size_t* actual_written) {
    const size_t page_size = arch_mem_page_size();
    size_t progress = 0;
    errno_t status = 0;

    while (progress < len) {
        const size_t misalign = (progress + offset) % page_size;
        const size_t page_index = (progress + offset) / page_size;
        const size_t copy_size = MIN(page_size - misalign, len - progress);

        struct page* p;
        status = obj->get_page(obj, page_index, &p);
        if (status != 0)
            break;

        uintptr_t page_addr = vm_page_addr(p);

        memcpy(HHDM_PTR(page_addr) + misalign, buf + progress, copy_size);
        progress += copy_size;
    }

    if (actual_written)
        *actual_written = progress;
    return status;
}

errno_t vmo_copy(
    struct vmo* target,
    uintptr_t target_offset,
    struct vmo* src,
    uintptr_t src_offset,
    size_t len,
    size_t* actual_copied
) {
    const size_t page_size = arch_mem_page_size();
    size_t progress = 0;
    errno_t status = 0;

    while (progress < len) {
        const size_t target_misalign = (progress + target_offset) % page_size;
        const size_t src_misalign = (progress + src_offset) % page_size;
        const size_t target_page_index = (progress + target_offset) / page_size;
        const size_t src_page_index = (progress + src_offset) / page_size;

        const size_t copy_size = MIN(MIN(page_size - target_misalign, page_size - src_misalign), len - progress);

        struct page* tpage;
        status = target->get_page(target, target_page_index, &tpage);
        if (status != 0)
            break;

        struct page* spage;
        status = src->get_page(src, src_page_index, &spage);
        if (status != 0)
            break;

        uintptr_t taddr = vm_page_addr(tpage);
        uintptr_t saddr = vm_page_addr(spage);

        memcpy(HHDM_PTR(taddr) + target_misalign, HHDM_PTR(saddr) + src_misalign, copy_size);
        progress += copy_size;
    }

    if (actual_copied)
        *actual_copied = progress;
    return status;
}

static errno_t paged_get_page(struct vmo* vmo, uintptr_t offset_idx, struct page** out) {
    struct paged_vmo* paged = containerof(vmo, struct paged_vmo, object);

    // Try to find the requested page in the cache.
    struct page_list* iter;
    spin_lock(&paged->lock);
    list_foreach(iter, &paged->cache, next) {
        if (iter->offset != offset_idx)
            continue;

        // We found a cached page!
        *out = iter->value;
        spin_unlock(&paged->lock);
        return 0;
    }
    spin_unlock(&paged->lock);

    // Get page from out pager if we don't have one already.
    struct page* new_page;
    errno_t status = paged->source.get_page(&paged->source, offset_idx, &new_page);
    if (status)
        return status;

    struct page_list* entry = mem_alloc(sizeof(struct page_list), 0);
    if (!entry) {
        paged->source.put_page(&paged->source, offset_idx, new_page);
        return ENOMEM;
    }

    entry->offset = offset_idx;
    entry->value = new_page;

    spin_lock(&paged->lock);
    list_foreach(iter, &paged->cache, next) {
        if (iter->offset != offset_idx)
            continue;

        *out = iter->value;
        spin_unlock(&paged->lock);
        mem_free(entry);
        paged->source.put_page(&paged->source, offset_idx, new_page);
        return 0;
    }

    // Add page to the object cache.
    list_insert_head(&paged->cache, entry, next);
    spin_unlock(&paged->lock);

    *out = new_page;
    return 0;
}

static void paged_destroy(struct vmo* vmo) {
    struct paged_vmo* paged = containerof(vmo, struct paged_vmo, object);

    struct page_list* iter = list_first(&paged->cache);
    while (iter) {
        struct page_list* next = list_next(iter, next);
        paged->source.put_page(&paged->source, iter->offset, iter->value);
        mem_free(iter);
        iter = next;
    }

    mem_free(paged);
}

errno_t vmo_new_paged(struct pager_ops pager, struct paged_vmo** out) {
    struct paged_vmo* result = mem_alloc(sizeof(struct paged_vmo), 0);
    if (!result)
        return ENOMEM;

    refcount_init(&result->object.rc);
    result->source = pager;
    result->object.get_page = paged_get_page;
    result->object.destroy = paged_destroy;

    *out = result;
    return 0;
}

static errno_t phys_get_page(struct pager_ops* pager, uintptr_t offset, struct page** out) {
    phys_t addr;
    errno_t status = mem_phys_alloc(1, 0, &addr);
    if (status != 0)
        return status;
    size_t idx = addr / arch_mem_page_size();
    *out = &vm_pfndb[idx];
    return 0;
}

static errno_t phys_put_page(struct pager_ops* pager, uintptr_t offset, struct page* page) {
    return mem_phys_free(vm_page_addr(page), 1);
}

static const struct pager_ops phys_pager = {
    .get_page = phys_get_page,
    .put_page = phys_put_page,
};

errno_t vmo_new_phys(struct paged_vmo** out) {
    return vmo_new_paged(phys_pager, out);
}

void vmo_truncate_paged(struct paged_vmo* paged, size_t length) {
    const size_t page_size = arch_mem_page_size();
    const size_t keep_pages = ROUND_UP(length, page_size);

    spin_lock(&paged->lock);
    struct page_list** pp = &list_first(&paged->cache);
    while (*pp) {
        struct page_list* entry = *pp;
        if (entry->offset >= keep_pages) {
            *pp = list_next(entry, next);
            paged->source.put_page(&paged->source, entry->offset, entry->value);
            mem_free(entry);
        } else {
            pp = &list_next(entry, next);
        }
    }
    spin_unlock(&paged->lock);
}

struct phys_range_vmo {
    struct vmo object;
    uintptr_t start;
    size_t len;
};

static errno_t phys_range_get_page(struct vmo* vmo, uintptr_t offset_idx, struct page** out) {
    struct phys_range_vmo* p = containerof(vmo, struct phys_range_vmo, object);
    if (offset_idx >= p->len / arch_mem_page_size())
        return EINVAL;
    *out = &vm_pfndb[(p->start / arch_mem_page_size()) + offset_idx];
    return 0;
}

static void phys_range_destroy(struct vmo* vmo) {
    mem_free(containerof(vmo, struct phys_range_vmo, object));
}

errno_t vmo_new_phys_range(uintptr_t addr, size_t length, struct vmo** out) {
    struct phys_range_vmo* vmo = mem_alloc(sizeof(struct phys_range_vmo), 0);
    if (!vmo)
        return ENOMEM;

    refcount_init(&vmo->object.rc);
    vmo->object.get_page = phys_range_get_page;
    vmo->object.destroy = phys_range_destroy;
    vmo->start = ALIGN_DOWN(addr, arch_mem_page_size());
    vmo->len = ALIGN_UP(length, arch_mem_page_size());

    *out = &vmo->object;
    return 0;
}