#include <zinnia/alloc.h>
#include <zinnia/atomic.h>
#include <zinnia/errno.h>
#include <zinnia/irq.h>
#include <zinnia/percpu.h>
#include <zinnia/pmap.h>
#include <zinnia/spin.h>
#include <zinnia/vma.h>
#include <arch/cpu.h>
#include <config.h>

// TODO: For now, we assume that we can do generic PTE lookups on every arch that has a MMU.
static_assert(CONFIG_HAVE_MMU_GENERIC);

errno_t pmap_new_kernel(struct pmap* pt, enum alloc_flags flags) {
    flags &= ~ALLOC_NOZERO;

    phys_t addr;
    errno_t status = mem_phys_alloc(1, flags, &addr);
    if (status)
        return status;

    pte_t* root = HHDM_PTR(addr);
    const size_t num_entries = (size_t)1 << arch_mem_level_bits();
    for (size_t i = num_entries / 2; i < num_entries; i++) {
        phys_t dir;
        status = mem_phys_alloc(1, flags, &dir);
        if (status)
            return status;
        root[i] = pte_build(dir, PTE_DIR, CACHE_ANY);
    }

    struct pmap result = {
        .root = addr,
        .lock = (struct spinlock){0},
        .is_user = false,
    };

    *pt = result;
    return 0;
}

errno_t pmap_new_user(struct pmap* pt, enum alloc_flags flags) {
    phys_t user_l1;
    errno_t status = mem_phys_alloc(1, 0, &user_l1);
    if (status)
        return status;

    void* user_l1_ptr = HHDM_PTR(user_l1);
    void* kernel_l1_ptr = HHDM_PTR(kernel_vma.pmap.root);
    memcpy(user_l1_ptr, kernel_l1_ptr, arch_mem_page_size());

    struct pmap result = {
        .root = user_l1,
        .lock = (struct spinlock){0},
        .is_user = true,
    };

    *pt = result;
    return 0;
}

// Gets a reference to the PTE at the given virtual address.
// If `check_only` is set, only checks if the PTE exists,
// and doesn't allocate new levels if they don't already exist.
// If it can't allocate a page if it has to, returns `nullptr`.
static pte_t* get_pte(struct pmap* pt, uintptr_t vaddr, bool is_user, bool check_only) {
    pte_t* current_head = HHDM_PTR(pt->root);
    size_t index = 0;

    for (int8_t level = arch_mem_num_levels() - 1; level >= 0; level--) {
        const size_t addr_mask = (1 << arch_mem_level_bits()) - 1;
        const size_t addr_shift = arch_mem_page_bits() + (arch_mem_level_bits() * level);
        const enum pte_flags level_flags = PTE_DIR | (is_user ? PTE_USER : 0);

        index = (vaddr >> addr_shift) & addr_mask;

        // The last level is used to access the actual PTE, so break the loop then.
        // We still need to update the index beforehand, that's why we can't just end early.
        if (level == 0)
            break;

        pte_t* pte = &current_head[index];

        if (pte_is_present(pte)) {
            // Get the next level.
            *pte |= pte_build(0, level_flags, CACHE_ANY);
            current_head = HHDM_PTR(pte_address(pte));
        } else {
            // If the current level isn't present, we can skip the rest.
            if (check_only)
                return nullptr;

            phys_t addr;
            if (mem_phys_alloc(1, 0, &addr) != 0)
                return nullptr;

            *pte = pte_build(addr, level_flags, CACHE_ANY);
            current_head = HHDM_PTR(addr);
        }
    }

    return &current_head[index];
}

errno_t pmap_map(struct pmap* pt, uintptr_t vaddr, phys_t paddr, enum pte_flags flags, enum cache_mode cache) {
    spin_lock(&pt->lock);

    errno_t status = 0;
    pte_t* pte = get_pte(pt, vaddr, flags & PTE_USER, false);
    if (pte == nullptr) {
        status = ENOMEM;
        goto fail;
    }

    *pte = pte_build(paddr, flags, cache);

fail:
    spin_unlock(&pt->lock);
    return status;
}

errno_t pmap_unmap(struct pmap* pt, uintptr_t vaddr) {
    spin_lock(&pt->lock);

    errno_t status = 0;
    pte_t* pte = get_pte(pt, vaddr, false, true);
    if (pte == nullptr || !pte_is_present(pte)) {
        status = ENOENT;
        goto fail;
    }

    pte_clear(pte);
    pte_flush(vaddr);

fail:
    spin_unlock(&pt->lock);
    return status;
}

errno_t pmap_extract(struct pmap* pt, uintptr_t vaddr, phys_t* out) {
    spin_lock(&pt->lock);

    errno_t status = 0;
    pte_t* pte = get_pte(pt, vaddr, false, true);
    if (pte == nullptr || !pte_is_present(pte))
        status = ENOENT;
    else
        *out = pte_address(pte);

    spin_unlock(&pt->lock);
    return status;
}

errno_t pmap_protect(struct pmap* pt, uintptr_t vaddr, enum pte_flags flags) {
    spin_lock(&pt->lock);

    errno_t status = 0;
    pte_t* pte = get_pte(pt, vaddr, false, true);
    if (pte == nullptr || !pte_is_present(pte)) {
        status = ENOENT;
        goto fail;
    }

    *pte = pte_build(pte_address(pte), flags, CACHE_ANY);
    pte_flush(vaddr);

fail:
    spin_unlock(&pt->lock);
    return status;
}

static void free_table(phys_t table, int8_t level) {
    if (level > 0) {
        pte_t* entries = HHDM_PTR(table);
        const size_t num_entries = (size_t)1 << arch_mem_level_bits();
        for (size_t i = 0; i < num_entries; i++) {
            if (pte_is_present(&entries[i]))
                free_table(pte_address(&entries[i]), level - 1);
        }
    }
    mem_phys_free(table, 1);
}

void pmap_destroy(struct pmap* pt) {
    if (!pt->is_user)
        return;

    pte_t* root = HHDM_PTR(pt->root);
    const size_t num_entries = (size_t)1 << arch_mem_level_bits();
    for (size_t i = 0; i < num_entries / 2; i++) {
        if (pte_is_present(&root[i]))
            free_table(pte_address(&root[i]), arch_mem_num_levels() - 2);
    }
    mem_phys_free(pt->root, 1);
    pt->root = 0;
}

[[__percpu]]
static atomic_size_t shootdown_req;

[[__percpu]]
static atomic_size_t shootdown_ack;

void pmap_shootdown_handle() {
    const size_t req = atomic_load(percpu_ref(shootdown_req), ATOMIC_ACQUIRE);
    if (atomic_load(percpu_ref(shootdown_ack), ATOMIC_RELAXED) == req)
        return;
    pte_flush_all();
    atomic_store(percpu_ref(shootdown_ack), req, ATOMIC_RELEASE);
}

void pmap_shootdown(struct pmap* pt) {
    const size_t count = percpu_count();
    if (count < 2)
        return;

    irq_lock();
    struct percpu* self = percpu_get();
    for (size_t i = 0; i < count; i++) {
        struct percpu* cpu = percpu_cpu(i);
        if (cpu == self || !cpu->online)
            continue;
        atomic_fetch_add(percpu_ref_on(cpu, shootdown_req), 1, ATOMIC_RELEASE);
        pte_shootdown_kick(cpu);
    }

    for (size_t i = 0; i < count; i++) {
        struct percpu* cpu = percpu_cpu(i);
        if (cpu == self || !cpu->online)
            continue;
        const size_t req = atomic_load(percpu_ref_on(cpu, shootdown_req), ATOMIC_RELAXED);
        while ((ptrdiff_t)(req - atomic_load(percpu_ref_on(cpu, shootdown_ack), ATOMIC_ACQUIRE)) > 0) {
            pmap_shootdown_handle();
            arch_cpu_pause();
        }
    }
    irq_unlock();
}