#include "pmm.h"
#include "vmm.h"          // for MEM_PHYS_OFFSET
#include "memory.h"    // memset
#include "stdio.h"

static spinlock_t memory_lock = SPINLOCK_INIT;
static uint8_t *bitmap = NULL;
static uint64_t total_page_count = 0;
static uint64_t last_used_index = 0;
static uint64_t free_pages = 0;

void pmm_init(struct limine_memmap_entry **memmap, size_t memmapentries) {
    uint64_t highest_addr = 0;
    size_t memmap_entries = memmapentries - 1;
    for (size_t i = 0; i < memmap_entries; i++) {
        if (memmap[i]->type != LIMINE_MEMMAP_USABLE &&
            memmap[i]->type != LIMINE_MEMMAP_BOOTLOADER_RECLAIMABLE)
            continue;
        uint64_t top = memmap[i]->base + memmap[i]->length;
        if (top > highest_addr) highest_addr = top;
    }

    total_page_count = highest_addr / PAGE_SIZE;
    size_t bitmap_size = ALIGN_UP(total_page_count / 8, PAGE_SIZE);
    printf("MEM_PHYS_OFFSET: %x\n", MEM_PHYS_OFFSET);
    printf("memmap_entries: %u\n", memmap_entries);
    // Find place for bitmap
    for (size_t i = 0; i < memmap_entries; i++) {
        printf("i: %u\n", i);

        if (memmap[i]->type != LIMINE_MEMMAP_USABLE) continue;
        if (memmap[i]->length >= bitmap_size) {
            printf("test\n");
            bitmap = (uint8_t *)(memmap[i]->base + MEM_PHYS_OFFSET);
            printf("bitmap: %x\n", bitmap);
            printf("bitmap_size: %x\n", bitmap_size);
            memset(bitmap, 0xFF, bitmap_size);
            printf("test3\n");
            printf("memmap base: %x\n", memmap[i]->base);
            printf("memmap length: %x\n", memmap[i]->length);

            memmap[i]->base += bitmap_size;
            memmap[i]->length -= bitmap_size;
            printf("test4\n");
            printf("memmap[7]->type: %u\n", memmap[7]->type);
            printf("memmap[8]->type: %u\n", memmap[8]->type);
            printf("memmap[9]->type: %u\n", memmap[9]->type);
            printf("memmap[10]->type: %u\n", memmap[10]->type); //#define LIMINE_MEMMAP_USABLE                 0
            printf("memmap[11]->type: %u\n", memmap[11]->type); //#define LIMINE_MEMMAP_RESERVED               1
            printf("memmap[12]->type: %u\n", memmap[12]->type); //#define LIMINE_MEMMAP_ACPI_RECLAIMABLE       2
            printf("memmap[13]->type: %u\n", memmap[13]->type); //#define LIMINE_MEMMAP_ACPI_NVS               3
            printf("memmap[14]->type: %u\n", memmap[14]->type); //#define LIMINE_MEMMAP_BAD_MEMORY             4
            printf("memmap[15]->type: %u\n", memmap[15]->type); //#define LIMINE_MEMMAP_BOOTLOADER_RECLAIMABLE 5
            printf("memmap[16]->type: %u\n", memmap[16]->type); //#define LIMINE_MEMMAP_EXECUTABLE_AND_MODULES 6
            printf("memmap[17]->type: %u\n", memmap[17]->type); //#define LIMINE_MEMMAP_FRAMEBUFFER            7
            //printf("memmap[18]->type: %u\n", memmap[18]->type); //#define LIMINE_MEMMAP_RESERVED_MAPPED        8, commented out because it triple faults for some reason
            break;
        }
        
    }
    printf("break");

    // Mark usable pages free
    for (size_t i = 0; i < memmap_entries; i++) {
        if (memmap[i]->type != LIMINE_MEMMAP_USABLE) continue;
        for (uint64_t j = 0; j < memmap[i]->length; j += PAGE_SIZE) {
            size_t bit = (memmap[i]->base + j) / PAGE_SIZE;
            bitmap[bit / 8] &= ~(1u << (bit % 8));
            free_pages++;
        }
    }
    printf("done");
}

static void *inner_alloc(size_t pages, size_t limit) {
    size_t p = 0;
    while (last_used_index < limit) {
        if (!(bitmap[last_used_index / 8] & (1u << (last_used_index % 8)))) {
            if (++p == pages) {
                size_t page = last_used_index - pages + 1;
                for (size_t i = page; i < page + pages; i++)
                    bitmap[i / 8] |= (1u << (i % 8));
                last_used_index = page + pages;
                return (void *)(page * PAGE_SIZE);
            }
        } else {
            p = 0;
        }
        last_used_index++;
    }
    return NULL;
}

void *pmm_alloc(size_t pages) {
    spinlock_acquire_or_wait(&memory_lock);
    size_t last = last_used_index;
    void *ret = inner_alloc(pages, total_page_count);
    if (!ret) {
        last_used_index = 0;
        ret = inner_alloc(pages, last);
    }
    if (ret) free_pages -= pages;
    spinlock_drop(&memory_lock);
    return ret;
}

void *pmm_allocz(size_t pages) {
    void *ret = pmm_alloc(pages);
   
    if (ret) {
        memset((void *)((uintptr_t)ret + MEM_PHYS_OFFSET), 0, pages * PAGE_SIZE); // this is at fault for the page fault
    }
    return ret;
}

void pmm_free(void *addr, size_t pages) {
    spinlock_acquire_or_wait(&memory_lock);
    size_t page = (uintptr_t)addr / PAGE_SIZE;
    for (size_t i = 0; i < pages; i++)
        bitmap[(page + i) / 8] &= ~(1u << ((page + i) % 8));
    free_pages += pages;
    spinlock_drop(&memory_lock);
}