//
// Created by rick on 23-02-21.
//

#include <stdbool.h>
#include <string.h>

#include <myke/libk/libk.h>
#include <myke/mem/pm.h>
#include <myke/mem/vmm.h>

#define NUM_PAGING_INFOS 16

#define BLOCK_USED 1
#define BLOCK_FREE 0

#define bitmap_block_byte(block) ((block) / 8)
#define bitmap_block_bit(block) ((block) % 8)

#define bitmap_set(addr, block) (addr)[bitmap_block_byte(block)] = ((addr)[bitmap_block_byte(block)] | (1 << (bitmap_block_bit(block))))
#define bitmap_unset(addr, block) (addr)[bitmap_block_byte(block)] = ((addr)[bitmap_block_byte(block)] & ~(1 << (bitmap_block_bit(block))))

#define bitmap_get(addr, block) ((addr)[bitmap_block_byte(block)] >> bitmap_block_bit(block))

#define page_to_addr(base_addr, block) ((base_addr) + ((block) * PAGE_SIZE))
#define addr_to_page(base_addr, addr) (((addr) - (base_addr)) / PAGE_SIZE)

typedef struct {
    union {
        uint8_t *page_bitmap;
        void *first_page;
    };
    uint32_t num_pages;
    size_t size;
    struct {
        bool present: 1;
    } flags;
} paging_info;

uint8_t last_paging_info = 0;
paging_info paging_infos[NUM_PAGING_INFOS];

void pm_init(void *start_addr, size_t size) {
    uint32_t num_pages = size / PAGE_SIZE;
    uint32_t bitmap_size = num_pages / 8;
    uint32_t bitmap_pages = bitmap_size / PAGE_SIZE;
    paging_infos[last_paging_info].page_bitmap = start_addr;
    paging_infos[last_paging_info].num_pages = num_pages;
    paging_infos[last_paging_info].size = size;
    paging_infos[last_paging_info].flags.present = true;

    // map into memory
    vmm_assign_page(start_addr, start_addr, bitmap_pages);

    // set state
    memset(start_addr, BLOCK_FREE, bitmap_pages);
    for (uint32_t i = 0; i < bitmap_pages; ++i) {
        bitmap_set((uint8_t *) start_addr, i);
    }
    last_paging_info++;
}

void *pm_get_pages(uint32_t num_pages) {
    for (int i = 0; i < last_paging_info; ++i) {
        paging_info *info = &paging_infos[i];
        if (!info->flags.present) continue;

        for (uint32_t j = 0; j < info->num_pages; ++j) {
            if (bitmap_get(info->page_bitmap, j) == BLOCK_FREE) {
                uint32_t k;
                bool usable = true;
                for (k = j + 1; k < j + num_pages; ++k) {
                    if (bitmap_get(info->page_bitmap, k) == BLOCK_USED) {
                        usable = false;
                        j = k; // skip ahead
                        break;
                    }
                }
                if (!usable) break;
                for (k = j; k < j + num_pages; ++k) {
                    bitmap_set(info->page_bitmap, k);
                }
                void *paddr = page_to_addr(info->first_page, j);
                vmm_assign_page(paddr, paddr, num_pages);
                return paddr;
            }
        }
    }
    k_panics("No page found!");
    return NULL;
}

void pm_free_pages(void *page, uint32_t num_pages) {
    for (int i = 0; i < last_paging_info; ++i) {
        paging_info *info = &paging_infos[i];
        if (!info->flags.present) continue;
        // check if page is in memory area
        if (page < info->first_page || page > (info->first_page + info->size)) continue;

        // free page
        uint32_t first_page = addr_to_page(info->first_page, page);
        for (uint32_t current_page = first_page; current_page < first_page + num_pages; ++current_page) {
            bitmap_unset(info->page_bitmap, current_page);
        }
        vmm_unassign_page(page, num_pages);
    }
}