#include "memory.h"
#include "debug.h"
#include "../lib/kernel/print.h"
#include "../lib/string.h"
#include "../lib/stdint.h"
#include "../thread/sync.h"
#include "../thread/thread.h"

/************************ 位图地址 ****************************************
 * 0xc009f000 是内核主线程栈顶，0xc009e000 是内核主线程的pcb
 * 一个页框大小的位图可表示128MB内存，位图位置安排在地址0xc009a000
 * 0xc009e000 - 0x4000 = 0xc009a000
 */
#define MEM_BITMAP_BASE 0xc009a000
/***********************************************************************/
/*0xc0000000 是内核从虚拟地址 3G 起，0xc0100000 指跨过低端 1MB 内存，使虚拟地址在逻辑上连续 */
#define K_HEAP_START 0xc0100000

#define PDE_IDX(addr) ((addr & 0xffc00000) >> 22) // 获取虚拟地址的页目录项 pde 的索引(高 10 位)
#define PTE_IDX(addr) ((addr & 0x003ff000) >> 12) // 获取虚拟地址的页表项 pte 的索引(中间 10 位)

// static functions
static void mem_pool_init(uint32_t all_mem);                  // 初始化内存池
static void *vaddr_get(enum pool_flags pf, uint32_t pg_cnt);  // 在 pf 指向的虚拟内存池中申请 pg_cnt 个虚拟页
static void *palloc(struct pool *m_pool);                     // 从内存池 m_pool 中申请 1 个物理页，成功则返回页框的物理地址，失败则返回 NULL
static void page_table_add(void *_vaddr, void *_page_pyaddr); // 在页表中添加虚拟地址 _vaddr 与物理地址 _page_pyaddr 的映射

/* 内存池结构，生成两个实例用于管理内核内存池和用户内存池 */
struct pool
{
    struct bitmap pool_bitmap; // 本内存池用到的位图结构，用于管理物理内存
    uint32_t phy_addr_start;   // 本内存池所管理物理内存的起始地址
    uint32_t pool_size;        // 本内存池字节容量
    struct lock lock;          // 申请内存时互斥
};

struct pool kernel_pool, user_pool; // 生成内核内存池和用户内存池
struct virtual_addr kernel_vaddr;   // 此结构是用来给内核分配虚拟地址

/* 初始化内存池 */
static void mem_pool_init(uint32_t all_mem)
{
    put_str("    mem_pool_init start\n");
    lock_init(&kernel_pool.lock);
    lock_init(&user_pool.lock);
    // 页表大小 = 页目录项数 * 每个页目录项占用的字节数
    // 页目录项数 = 1 页的页目录表 + 第 0 和第 768(内核页目录项)个页目录项指向同一个页表 +
    // 第 769 ~ 1022 个页目录项指向 254 个页表 = 256 个页目录项
    //! 最后 1 个页目录项(第 1023 个 pde)指向页目录表自身
    uint32_t page_table_size = PG_SIZE * 256; // 256 个页目录项 * 4KB = 1MB

    // 1024(0x400) * 1024 =  0x100000  是低端 1MB 内存
    uint32_t used_mem = page_table_size + 0x100000;
    uint32_t free_mem = all_mem - used_mem;
    uint16_t all_free_pages = free_mem / PG_SIZE; // 1 页 = 4KB = 4096 字节

    uint16_t kernel_free_pages = all_free_pages / 2;
    uint16_t user_free_pages = all_free_pages - kernel_free_pages;

    /* 为了简化操作，余数不处理，坏处是这样做会丢失内存，好像是不用做内存的越界检查，因为位图表的内存少于实际物理内存 */
    uint32_t kbm_length = kernel_free_pages / 8; // Kernel BitMap 长度，位图中的一位表示一页，以字节为单位
    uint32_t ubm_length = user_free_pages / 8;   // User BitMap 长度

    uint32_t kp_start = used_mem;                               // Kernel Pool 起始地址,内核内存池的起始地址
    uint32_t up_start = kp_start + kernel_free_pages * PG_SIZE; // User Pool 起始地址,用户内存池的起始地址

    kernel_pool.phy_addr_start = kp_start;
    user_pool.phy_addr_start = up_start;

    kernel_pool.pool_size = kernel_free_pages * PG_SIZE;
    user_pool.pool_size = user_free_pages * PG_SIZE;

    kernel_pool.pool_bitmap.btmp_bytes_len = kbm_length;
    user_pool.pool_bitmap.btmp_bytes_len = ubm_length;

    /********* 内核内存池和用户内存池的位图 ****************
     * 位图是全局的数据，长度不固定
     * 全局或静态的数组需要在编译时知道其长度
     * 而我们需要根据总内存大小算出需要多少字节
     * 所以改为指定一块内存来生成位图
     ************************************************/
    // 内核使用的最高地址是 0xc009f000,这是主线程的栈地址
    // 内核的大小预计为 70KB 左右
    // （32MB = 32 * 1024 * 1024 字节）内存占位图是 1KB
    // 内核内存池位图的位置安排在地址 0xc009a000
    kernel_pool.pool_bitmap.bits = (void *)MEM_BITMAP_BASE;
    // 用户内存池的位图紧跟在内核内存池位图之后
    user_pool.pool_bitmap.bits = (void *)((uintptr_t)MEM_BITMAP_BASE + kbm_length);

    /********************** 输出内存池信息 ******************************/
    put_str("    kernel_pool_bitmap_start: ");
    put_int((uintptr_t)kernel_pool.pool_bitmap.bits);
    put_str(", kernel_pool_phy_addr_start: ");
    put_int(kernel_pool.phy_addr_start);
    put_str("\n");
    put_str("    user_pool_bitmap_start: ");
    put_int((uintptr_t)user_pool.pool_bitmap.bits);
    put_str(", user_pool_phy_addr_start: ");
    put_int(user_pool.phy_addr_start);
    put_str("\n");

    bitmap_init(&kernel_pool.pool_bitmap);
    bitmap_init(&user_pool.pool_bitmap);

    /*********************** 初始化虚拟内存池 ********************************/
    kernel_vaddr.vaddr_bitmap.btmp_bytes_len = kbm_length; // 虚拟地址位图的长度,按实际物理内存大小生成数组
    /*位图的数组指向一块未使用的内存，目前定位在内核内存池和用户内存池之后*/
    kernel_vaddr.vaddr_bitmap.bits = (void *)((uintptr_t)MEM_BITMAP_BASE + kbm_length + ubm_length);
    kernel_vaddr.vaddr_start = K_HEAP_START;
    bitmap_init(&kernel_vaddr.vaddr_bitmap);

    put_str("    mem_pool_init end\n");
}

// 内存管理初始化
void mem_init(void)
{
    put_str("mem_init start\n");
    uint32_t mem_bytes_total = (*(uint32_t *)(0xb00)); // 从地址 0xb00 处取出总内存数
    mem_pool_init(mem_bytes_total);                    // 初始化内存池
    put_str("mem_init done\n");
}

// 从内核内存池中申请 pg_cnt 个页空间，成功则返回其虚拟地址，失败则返回 NULL
static void *vaddr_get(enum pool_flags pf, uint32_t pg_cnt)
{
    uintptr_t vaddr_start = 0;
    int bit_idx_start = -1;
    uint32_t cnt = 0;
    if (PF_KERNEL == pf)
    { // 内核内存池
        bit_idx_start = bitmap_scan(&kernel_vaddr.vaddr_bitmap, pg_cnt);
        if (-1 == bit_idx_start)
        {
            return NULL;
        }
        while (cnt < pg_cnt)
        {
            bitmap_set(&kernel_vaddr.vaddr_bitmap, bit_idx_start + cnt++, 1);
        }
        vaddr_start = kernel_vaddr.vaddr_start + bit_idx_start * PG_SIZE;
    }
    else
    { // 用户内存池 -- 需要使用当前线程来分配
        struct task_struct *cur = running_thread();
        bit_idx_start = bitmap_scan(&cur->userprog_addr.vaddr_bitmap, pg_cnt);
        if (bit_idx_start == -1)
        {
            return NULL;
        }
        while (cnt < pg_cnt)
        {
            bitmap_set(&cur->userprog_addr.vaddr_bitmap, bit_idx_start + cnt++, 1);
        }
        vaddr_start = cur->userprog_addr.vaddr_start + bit_idx_start * PG_SIZE;

        // (0xc0000000 - PG_SIZE) 作为用户 3 级栈已经在 start_process 被分配
        ASSERT(((uint32_t)vaddr_start) < (0xc0000000 - PG_SIZE));
    }
    return (void *)vaddr_start;
}

/****************************************************************************
 * 处理器处理 32 位地址的三个步骤：
 * 1. 分段：通过段描述符找到段基址，然后加上段内偏移地址
 * 2. 分页：通过页目录表和页表找到物理地址
 *    （1）页目录表的索引是虚拟地址的高 10 位，处理 pde 索引，得到页表物理地址
 *    （2）页表的索引是虚拟地址的中间 10 位,处理 pte 索引，得到普通物理页的物理地址
 *    （3）页内偏移地址是虚拟地址的低 12 位,加上普通物理页的物理地址，得到最终物理地址
 * 3. 访问内存
 ****************************************************************************/

// 得到虚拟地址 vaddr 对应的 pte 指针(虚拟地址)
uint32_t *pte_ptr(uint32_t vaddr)
{
    // 先访问到页目录表地址(高 10 位)
    //    最后一个页目录项保存的是页目录表的地址，我们需要让高 10 们指向最后一个页目录项,即第 1023 个页目录项
    //    1023=0x3ff，移动到高 10 位，所以是 0xffc00000,此处是页目录表的物理地址 0x100000
    // 找到页表地址(中间 10 位)
    //    页表物理地址 = 页目录表的物理地址 + vaddr 的页目录项索引
    // 在页表中找到 pte 地址
    //    页表物理地址 + vaddr 的页表项索引 * 4

    uint32_t *pte = (uint32_t *)((uintptr_t)0xffc00000 + ((vaddr & 0xffc00000) >> 10) + PTE_IDX(vaddr) * 4);
    return pte;
}

// 得到虚拟地址 vaddr 对应的 pde 指针(虚拟地址)
uint32_t *pde_ptr(uint32_t vaddr)
{
    // 由于最后一个页目录项中存储的是页目录物理地址，所以 32 位地址中高 20 位为 0xfffff,时表示访问到的是最后一个目录项
    // 即获得了页目录表的物理地址
    // 因此新虚拟地址 new_vaddr = 0xfffff000 + vaddr 的页目录项索引 * 4
    return (uint32_t *)((uintptr_t)0xfffff000 + PDE_IDX(vaddr) * 4);
}

// 在 m_pool 指向的物理内存池中申请1个物理页，成功则返回页框的物理地址，失败则返回NULL
static void *palloc(struct pool *m_pool)
{
    int bit_idx = bitmap_scan(&m_pool->pool_bitmap, 1); // 找一个物理页面
    if (-1 == bit_idx)
    {
        return NULL;
    }
    bitmap_set(&m_pool->pool_bitmap, bit_idx, 1); // 将此位bit_idx置 1
    uint32_t page_phyaddr = ((m_pool->phy_addr_start + bit_idx * PG_SIZE));
    return (void *)(uintptr_t)page_phyaddr;
}

// 页表中添加虚拟地址 _vaddr 与物理地址 _page_phyaddr 的映射
static void page_table_add(void *_vaddr, void *_page_pyaddr)
{
    uint32_t vaddr = (uint32_t)(uintptr_t)_vaddr, page_phyaddr = (uint32_t)(uintptr_t)_page_pyaddr;
    uint32_t *pde = pde_ptr(vaddr);
    uint32_t *pte = pte_ptr(vaddr);

    /*********************************************************************
     * !! 注意：
     * !! 执行 *pte 时，会访问到空的 pte，所以确保 pte 的地址在 pde 创建之后，
     * !! 否则会引发 page_fault。因此 *pde o 0 时，*pte 不会执行
     ********************************************************************/
    // 先在页目录内判断目录项的 p 位，若为 1，则表示该表已存在
    if (*pde & 0x00000001)
    {
        // 页目录项和页表项的第 0 位为 P 位，为 1 表示存在
        ASSERT(!(*pte & 0x00000001));
        if (!(*pte & 0x00000001))
        {
            // 页表项不存在，创建页表项
            *pte = (page_phyaddr | PG_US_U | PG_RW_W | PG_P_1);
        }
        else
        { // todo: 目前应该不会执行到这里，因为上面的 ASSERT 会先执行
            PANIC("pte repeat");
            *pte = (page_phyaddr | PG_US_U | PG_RW_W | PG_P_1);
        }
    }
    else // 页目录项不存在，要先创建页目录再创建页表项
    {
        /* 页表中用到的页框一律从内核空间分配 */
        uint32_t pde_phyaddr = (uint32_t)(uintptr_t)palloc(&kernel_pool);
        *pde = (pde_phyaddr | PG_US_U | PG_RW_W | PG_P_1); // 页目录项赋值页表的物理地址
        /* 分配到的物理页地址 pde_phyaddr 对应的物理内存清 0，
         * 避免里面的陈旧数据变成了页表项，从而让页表混乱
         * 访问到 pde 对应的物理地址，用 pte 取高 20 位。
         * 因为 pte 基于该 pde 对应的物理地址再寻址，把低 12 位置 0 便是该 pde 对应的物理页的起始
         */
        _memset((void *)((uintptr_t)pte & 0xfffff000), 0, PG_SIZE);
        ASSERT(!(*pte & 0x00000001));
        *pte = (page_phyaddr | PG_US_U | PG_RW_W | PG_P_1); // 页表项赋值物理页的地址

        ASSERT(!(*pte & 0x00000001));
        *pte = (page_phyaddr | PG_US_U | PG_RW_W | PG_P_1); // 再次检查并确认
    }
}
// 分配 pg_cnt 个页空间，成功则返回虚拟地址，失败则返回 NULL
void *malloc_page(enum pool_flags pf, uint32_t pg_cnt)
{
    // todo: 内核和用户空间各约 16MB，所以暂时不考虑超过 16MB 的内存申请，保守起见 15MB
    //  pg_cnt< 15 * 1024 / 4KB = 3840 页
    ASSERT(pg_cnt > 0 && pg_cnt < 3840);

    void *vaddr_start = vaddr_get(pf, pg_cnt);
    if (NULL == vaddr_start)
    {
        return NULL;
    }
    uint32_t vaddr = (uint32_t)(uintptr_t)vaddr_start, cnt = pg_cnt;
    struct pool *mem_pool = (PF_KERNEL == pf ? &kernel_pool : &user_pool);

    // 因为虚拟地址是连续的，但物理地址可以不连续,所以逐个做映射
    while (cnt-- > 0)
    {
        void *page_phyaddr = palloc(mem_pool);
        if (NULL == page_phyaddr)
        {
            // 失败时要将曾经已申请的虚拟地址和物理页全部回滚
            return NULL;
        }
        page_table_add((void *)(uintptr_t)vaddr, page_phyaddr);
        vaddr += PG_SIZE; // 下一个虚拟页
    }

    return vaddr_start;
}

// 申请内核内存，成功则返回其虚拟地址，失败则返回 NULL
void *get_kernel_pages(uint32_t pg_cnt)
{
    void *vaddr = malloc_page(PF_KERNEL, pg_cnt);
    if (NULL != vaddr)
    {
        // 若分配的地址不为空，则将页内存清 0 后返回
        _memset(vaddr, 0, pg_cnt * PG_SIZE);
    }
    return vaddr;
}

// 申请用户内存，成功则返回其虚拟地址，失败则返回 NULL
void *get_user_pages(uint32_t pg_cnt)
{
    lock_acquire(&user_pool.lock);
    void *vaddr = malloc_page(PF_USER, pg_cnt);
    if (NULL != vaddr)
    {
        // 若分配的地址不为空，则将页内存清 0 后返回
        _memset(vaddr, 0, pg_cnt * PG_SIZE);
    }
    lock_release(&user_pool.lock);
    return vaddr;
}

// 将地址 vaddr 与 pf 池中的物理地址关联，仅支持一页空间分配
void *get_a_page(enum pool_flags pf, uint32_t vaddr)
{
    struct pool *mem_pool = pf & PF_KERNEL ? &kernel_pool : &user_pool;
    lock_acquire(&mem_pool->lock);

    /* 先将虚拟地址对应的位图置 1 */
    struct task_struct *cur = running_thread();
    int32_t bit_idx = -1;

    if (cur->pgdir != NULL && pf == PF_USER)
    { // 若当前是用户进程申请用户内存，就修改用户进程自己的虚拟地址位图
        bit_idx = (vaddr - cur->userprog_addr.vaddr_start) / PG_SIZE;
        ASSERT(bit_idx > 0);
        bitmap_set(&cur->userprog_addr.vaddr_bitmap, bit_idx, 1);
    }
    else if (cur->pgdir == NULL && pf == PF_KERNEL)
    { // 如果是内核线程申请内存，就修改 kernel_vaddr
        bit_idx = (vaddr - kernel_vaddr.vaddr_start) / PG_SIZE;
        ASSERT(bit_idx > 0);
        bitmap_set(&kernel_vaddr.vaddr_bitmap, bit_idx, 1);
    }
    else
    {
        PANIC("get_a_page: not allow kernel alloc userspace or user alloc kernelspace by get_a_page");
    }

    /* 然后申请物理页 */
    void *page_phyaddr = palloc(mem_pool);
    if (page_phyaddr == NULL)
        return NULL;
    page_table_add((void *)(uintptr_t)vaddr, page_phyaddr);

    lock_release(&mem_pool->lock);
    return (void *)(uintptr_t)vaddr;
}

// 得到虚拟地址映射到的物理地址
uint32_t addr_vaddr2phy(uint32_t vaddr)
{
    uint32_t *pte = pte_ptr(vaddr);
    /* (*pte)是值是页表所在的物理页框地址，去掉其低 12 位的页表项属性+虚拟地址 vaddr 的低 12 位 */
    return ((*pte & 0xfffff000) + (vaddr & 0xfff));
}