/* ============================================================================
 * Pactor64 — Bitmap-based Physical Memory Manager (x86-64)
 * Parses multiboot2 memory map, 4KB page granularity
 * ============================================================================ */
#include "../include/pactor64.h"

#define PAGE_SIZE    4096
#define MAX_REGIONS  32

/* Saved available regions from multiboot2 memory map */
typedef struct {
    uint64_t base;
    uint64_t length;
} mem_region_t;

/* === Bitmap === */
static uint8_t *pmm_bitmap   = NULL;
static volatile uint64_t total_pages  = 0;
static volatile uint64_t used_pages   = 0;
static uint64_t bitmap_size  = 0;

/* === Initialize PMM from multiboot2 memory map === */
void pmm_init(uint64_t mb2_info) {
    uint32_t mb2_size = *(uint32_t *)mb2_info;
    uint64_t highest_addr = 0;

    /*
     * We must extract ALL info from the multiboot2 structure FIRST,
     * because placing the bitmap after kernel_end may overwrite it.
     */
    mem_region_t regions[MAX_REGIONS];
    int num_regions = 0;

    /* Walk multiboot2 tags */
    uint64_t off = 8;
    while (off + 8 <= mb2_size) {
        uint32_t ttype = *(uint32_t *)(mb2_info + off);
        uint32_t tsize = *(uint32_t *)(mb2_info + off + 4);
        if (ttype == 0 || tsize < 8) break;

        if (ttype == 6) {
            /* Memory map tag */
            uint32_t entry_size = *(uint32_t *)(mb2_info + off + 8);
            uint8_t *entries = (uint8_t *)(mb2_info + off + 16);
            uint32_t entries_bytes = tsize - 16;
            int n = entries_bytes / entry_size;

            for (int i = 0; i < n && num_regions < MAX_REGIONS; i++) {
                uint64_t base = *(uint64_t *)(entries + i * entry_size);
                uint64_t len  = *(uint64_t *)(entries + i * entry_size + 8);
                uint32_t type = *(uint32_t *)(entries + i * entry_size + 16);

                if (type == 1 && len >= PAGE_SIZE) {
                    regions[num_regions].base   = base;
                    regions[num_regions].length = len;
                    num_regions++;

                    uint64_t top = base + len;
                    if (top > highest_addr) highest_addr = top;
                }
            }
        }

        off += tsize;
        off = (off + 7) & ~7ULL;
    }

    if (highest_addr == 0) highest_addr = 64 * 1024 * 1024;
    total_pages = highest_addr / PAGE_SIZE;

    /* Place bitmap right after the kernel (safe: mb2 data already parsed) */
    pmm_bitmap = (uint8_t *)kernel_end;
    bitmap_size = (total_pages + 7) / 8;

    /* Mark everything as used */
    kmemset(pmm_bitmap, 0xFF, bitmap_size);
    used_pages = total_pages;

    /* Free available regions using saved data */
    for (int r = 0; r < num_regions; r++) {
        uint64_t pg_start = (regions[r].base + PAGE_SIZE - 1) / PAGE_SIZE;
        uint64_t pg_end   = (regions[r].base + regions[r].length) / PAGE_SIZE;
        for (uint64_t pg = pg_start; pg < pg_end && pg < total_pages; pg++) {
            pmm_bitmap[pg / 8] &= ~(1 << (pg % 8));
            used_pages--;
        }
    }

    /* Re-mark kernel + bitmap region as used */
    uint64_t reserved_pages = ((uint64_t)pmm_bitmap + bitmap_size + PAGE_SIZE - 1) / PAGE_SIZE;
    for (uint64_t pg = 0; pg < reserved_pages && pg < total_pages; pg++) {
        if (!((pmm_bitmap[pg / 8] >> (pg % 8)) & 1)) {
            pmm_bitmap[pg / 8] |= (1 << (pg % 8));
            used_pages++;
        }
    }

    kprintf("  [OK] PMM: %d total, %d free (%d MB)\n",
            total_pages, total_pages - used_pages,
            (total_pages - used_pages) * 4 / 1024);
}

void *pmm_alloc_page(void) {
    for (uint64_t pg = 0; pg < total_pages; pg++) {
        if (!((pmm_bitmap[pg / 8] >> (pg % 8)) & 1)) {
            pmm_bitmap[pg / 8] |= (1 << (pg % 8));
            used_pages++;
            return (void *)(pg * PAGE_SIZE);
        }
    }
    return NULL;
}

void pmm_free_page(void *page) {
    uint64_t pg = (uint64_t)page / PAGE_SIZE;
    if (pg < total_pages && ((pmm_bitmap[pg / 8] >> (pg % 8)) & 1)) {
        pmm_bitmap[pg / 8] &= ~(1 << (pg % 8));
        used_pages--;
    }
}

uint64_t pmm_get_free_pages(void) {
    return total_pages - used_pages;
}

uint64_t pmm_get_total_pages(void) {
    return total_pages;
}