#include <cstdint>
#include <generic/arch.hpp>
#include <generic/bootloader/bootloader.hpp>
#include <generic/hhdm.hpp>
#include <generic/pmm.hpp>
#include <arch/riscv64/features.hpp>

#define PTE_MASK_VALUE (0x003ffffffffffc00)
#define PTE_PRESENT (1 << 0)
#define PTE_RW ((1 << 1) | (1 << 2) | (1 << 3))
#define PTE_USER (1 << 4)
#define PTE_WC (0)
#define PTE_MMIO (0) 
#define PTE_ACCESSED (1 << 6)
#define PTE_DIRTY (1 << 7)
#define LVL_PG_MASK PTE_MASK_VALUE


int level_to_index(std::uintptr_t virt, int level) {
    if (riscv64::get_paging_level() == 5) {
        switch (level) {
            case 0:  return PTE_INDEX(virt, 48);
            case 1:  return PTE_INDEX(virt, 39);
            case 2:  return PTE_INDEX(virt, 30);
            case 3:  return PTE_INDEX(virt, 21);
            case 4:  return PTE_INDEX(virt, 12);
            default: return 0;
        }
    } else if (riscv64::get_paging_level() == 4) {
        switch (level) {
            case 0:  return PTE_INDEX(virt, 39);
            case 1:  return PTE_INDEX(virt, 30);
            case 2:  return PTE_INDEX(virt, 21);
            case 3:  return PTE_INDEX(virt, 12);
            default: return 0;
        }
    } else if (riscv64::get_paging_level() == 3) {
        switch (level) {
            case 0:  return PTE_INDEX(virt, 30);
            case 1:  return PTE_INDEX(virt, 21);
            case 2:  return PTE_INDEX(virt, 12);
            default: return 0;
        }
    }
    return 0;
}

std::uint64_t extract_ppn(std::uint64_t masked_table) {
    return masked_table >> 10;
}

std::uint64_t make_pte(std::uint64_t ppn) {
    return ppn << 10;
}

int64_t* __paging_next_level_noalloc(std::uint64_t* table, std::uint16_t idx) {
    if (!(table[idx] & PTE_PRESENT))
        return (int64_t*) -1;
    return (int64_t*)((extract_ppn(table[idx] & PTE_MASK_VALUE) << 12) + etc::hhdm());
}


std::int64_t __memory_paging_getphys(std::uint64_t* table, std::uint64_t virt, int level) {
    if (!table && (std::int64_t) table == -1)
        return -1;
    int max_level = riscv64::get_paging_level() - 1;
    if (max_level == level) {
        return (table[level_to_index(virt, level)] & PTE_PRESENT) ? (extract_ppn(table[level_to_index(virt, level)] & LVL_PG_MASK) << 12) : -1;
    } else
        return __memory_paging_getphys((std::uint64_t*) __paging_next_level_noalloc(table, level_to_index(virt, level)), virt, level + 1);
    return -1;
}

uint64_t* riscv64_paging_next_level(std::uint64_t* table, std::uint16_t idx) {
    if (!(table[idx] & PTE_PRESENT))
        table[idx] = (make_pte(pmm::freelist::alloc_4k() >> 12)) | PTE_PRESENT;
    return (uint64_t*) ((extract_ppn(table[idx] & PTE_MASK_VALUE) << 12) + etc::hhdm());
}

std::uint64_t convert_flags(std::uint64_t flags) {
    std::uint64_t result = 0;
    if (flags & PAGING_NC)
        result |= PTE_MMIO;
    if (flags & PAGING_PRESENT)
        result |= PTE_PRESENT;
    if (flags & PAGING_RW)
        result |= PTE_RW;
    if (flags & PAGING_USER)
        result |= PTE_USER;
    if (flags & PAGING_WC)
        result |= PTE_WC;
    result |= PTE_ACCESSED | PTE_DIRTY;
    return result;
}


void riscv64_map_page(std::uintptr_t root, std::uint64_t phys, std::uintptr_t virt, std::uint32_t flags) {
    if (riscv64::get_paging_level() == 4) {
        std::uint64_t* l_root = (std::uint64_t*) (root + etc::hhdm());
        uint64_t* l1 = riscv64_paging_next_level(l_root, PTE_INDEX(virt, 39));
        uint64_t* l2 = riscv64_paging_next_level(l1, PTE_INDEX(virt, 30));
        uint64_t* l3 = riscv64_paging_next_level(l2, PTE_INDEX(virt, 21));
        l3[PTE_INDEX(virt, 12)] = (make_pte(phys >> 12)) | flags;
    } else if (riscv64::get_paging_level() == 3) {
        std::uint64_t* l_root = (std::uint64_t*) (root + etc::hhdm());
        uint64_t* l2 = riscv64_paging_next_level(l_root, PTE_INDEX(virt, 30));
        uint64_t* l3 = riscv64_paging_next_level(l2, PTE_INDEX(virt, 21));
        l3[PTE_INDEX(virt, 12)] = (make_pte(phys >> 12)) | flags;
    } else if (riscv64::get_paging_level() == 5) {
        std::uint64_t* l_root = (std::uint64_t*) (root + etc::hhdm());
        uint64_t* l0 = riscv64_paging_next_level(l_root, PTE_INDEX(virt, 48));
        uint64_t* l1 = riscv64_paging_next_level(l0, PTE_INDEX(virt, 39));
        uint64_t* l2 = riscv64_paging_next_level(l1, PTE_INDEX(virt, 30));
        uint64_t* l3 = riscv64_paging_next_level(l2, PTE_INDEX(virt, 21));
        l3[PTE_INDEX(virt, 12)] = (make_pte(phys >> 12)) | flags;
    }
}

namespace arch {
    [[gnu::weak]] void enable_paging(std::uintptr_t root) {
        std::uint64_t mode_root = (root >> 12) | ((std::uint64_t) riscv64::raw_level_paging << 60);
        asm volatile("csrw satp, %0" : : "r"(mode_root) : "memory");
        asm volatile("sfence.vma");
    }

    [[gnu::weak]] void map_page(std::uintptr_t root, std::uint64_t phys, std::uintptr_t virt, int flags) {
        riscv64_map_page(root,phys,virt,convert_flags(flags));
    }

    [[gnu::weak]] std::uint64_t get_phys_from_page(std::uintptr_t root, std::uintptr_t virt) {
        return __memory_paging_getphys((std::uint64_t*)(root + etc::hhdm()),virt,0);
    }

    [[gnu::weak]] void destroy_root(std::uintptr_t root, int level) {
        std::uint64_t* table = (std::uint64_t*) (root + etc::hhdm());
        if (level != riscv64::get_paging_level() - 1) {
            if (level == 0) {
                for (int i = 0; i < 256; i++) {
                    if (table[i] & PTE_PRESENT) {
                        destroy_root(table[i] & PTE_MASK_VALUE, level + 1);
                    }
                }
            } else {
                for (int i = 0; i < 512; i++) {
                    if (table[i] & PTE_PRESENT) {
                        destroy_root(table[i] & PTE_MASK_VALUE, level + 1);
                    }
                }
            }

            if (level != 0)
                pmm::freelist::free(root);
        }
    }

    [[gnu::weak]] void copy_higher_half(std::uintptr_t root, std::uintptr_t src_root) {
        std::uint64_t* virt_rootcr3 = (std::uint64_t*) (root + etc::hhdm());
        std::uint64_t* virt_srccr3 = (std::uint64_t*) (src_root + etc::hhdm());
        for (int i = 255; i < 512; i++) {
            virt_rootcr3[i] = virt_srccr3[i];
        }
    }
};