#include <memory/pmm.h>
#include <common.h>
#include <serial.h>
#include <string_utils.h>

extern EFI_SYSTEM_TABLE *ST;

pmm_t g_pmm;

static inline void bitmap_set(UINTN idx) {
    g_pmm.bitmap[idx / 8] |= (1 << (idx % 8));
}

static inline void bitmap_clear(UINTN idx) {
    g_pmm.bitmap[idx / 8] &= ~(1 << (idx % 8));
}

static inline BOOLEAN bitmap_test(UINTN idx) {
    return (g_pmm.bitmap[idx / 8] >> (idx % 8)) & 1;
}

// Clean stale entries from free list head
static void clean_free_list() {
    while (g_pmm.free_list_head != NULL &&
           bitmap_test((UINTN)g_pmm.free_list_head / PAGE_SIZE)) {
        g_pmm.free_list_head = *(void**)g_pmm.free_list_head;
    }
}

EFI_STATUS pmm_init() {
    UINTN map_size = 0;
    UINTN map_key;
    UINTN desc_size;
    UINT32 desc_version;

    EFI_STATUS status = uefi_call_wrapper(
        ST->BootServices->GetMemoryMap, 5,
        &map_size, NULL, &map_key, &desc_size, &desc_version
    );

    map_size += desc_size * 64;

    EFI_MEMORY_DESCRIPTOR* mem_map = NULL;
    status = uefi_call_wrapper(
        ST->BootServices->AllocatePool, 3,
        EfiLoaderData, map_size, (void**)&mem_map
    );
    if (EFI_ERROR(status)) return status;

    status = uefi_call_wrapper(
        ST->BootServices->GetMemoryMap, 5,
        &map_size, mem_map, &map_key, &desc_size, &desc_version
    );
    if (EFI_ERROR(status)) {
        uefi_call_wrapper(ST->BootServices->FreePool, 1, mem_map);
        return status;
    }

    UINTN entry_count = map_size / desc_size;

    // First pass: count total pages and find max physical address
    UINT64 max_addr = 0;
    UINT64 total_free = 0;
    for (UINTN i = 0; i < entry_count; i++) {
        EFI_MEMORY_DESCRIPTOR* desc = (EFI_MEMORY_DESCRIPTOR*)((UINT8*)mem_map + i * desc_size);
        UINT64 end = desc->PhysicalStart + desc->NumberOfPages * PAGE_SIZE;
        if (end > max_addr) max_addr = end;
        if (desc->Type == EfiConventionalMemory) {
            total_free += desc->NumberOfPages * PAGE_SIZE;
        }
    }

    g_pmm.base_addr = 0;
    g_pmm.max_addr  = max_addr;

    // How many pages does the bitmap cover?
    UINTN total_pages = (UINTN)(max_addr / PAGE_SIZE);
    g_pmm.total_pages = total_pages;

    // Bitmap size in bytes, rounded up to page boundary
    g_pmm.bitmap_size = ((total_pages + 7) / 8);
    UINTN bitmap_pages = (g_pmm.bitmap_size + PAGE_SIZE - 1) / PAGE_SIZE;
    g_pmm.bitmap_size = bitmap_pages * PAGE_SIZE;  // round to full pages

    // Place bitmap at the end of the highest free conventional memory region
    UINT64 bitmap_addr = 0;
    for (UINTN i = 0; i < entry_count; i++) {
        EFI_MEMORY_DESCRIPTOR* desc = (EFI_MEMORY_DESCRIPTOR*)((UINT8*)mem_map + i * desc_size);
        if (desc->Type == EfiConventionalMemory) {
            UINT64 region_bytes = desc->NumberOfPages * PAGE_SIZE;
            if (region_bytes >= g_pmm.bitmap_size) {
                UINT64 candidate = desc->PhysicalStart + region_bytes - g_pmm.bitmap_size;
                if (candidate > bitmap_addr) {
                    bitmap_addr = candidate;
                }
            }
        }
    }

    if (bitmap_addr == 0) {
        serial_write("PMM: ERROR - no space for bitmap!\n");
        uefi_call_wrapper(ST->BootServices->FreePool, 1, mem_map);
        return EFI_OUT_OF_RESOURCES;
    }

    g_pmm.bitmap = (UINT8*)(UINTN)bitmap_addr;

    // Init bitmap: mark ALL pages as used (0xFF)
    for (UINTN i = 0; i < g_pmm.bitmap_size; i++) {
        g_pmm.bitmap[i] = 0xFF;
    }

    // Mark free pages (EfiConventionalMemory) as free in bitmap
    g_pmm.free_pages = 0;
    UINT64 bm_start_page = bitmap_addr / PAGE_SIZE;
    UINT64 bm_end_page   = (bitmap_addr + g_pmm.bitmap_size + PAGE_SIZE - 1) / PAGE_SIZE;

    for (UINTN i = 0; i < entry_count; i++) {
        EFI_MEMORY_DESCRIPTOR* desc = (EFI_MEMORY_DESCRIPTOR*)((UINT8*)mem_map + i * desc_size);
        if (desc->Type != EfiConventionalMemory) continue;

        UINT64 start_page = desc->PhysicalStart / PAGE_SIZE;
        UINT64 end_page   = start_page + desc->NumberOfPages;

        for (UINT64 p = start_page; p < end_page; p++) {
            // Skip bitmap pages
            if (p >= bm_start_page && p < bm_end_page) continue;
            bitmap_clear((UINTN)p);
            g_pmm.free_pages++;
        }
    }

    // Mark bitmap pages as used
    for (UINT64 p = bm_start_page; p < bm_end_page; p++) {
        bitmap_set((UINTN)p);
    }

    // Reserve low memory (first 4 MB) — UEFI firmware may use it during BS calls
    UINT64 low_reserve_pages = 0x400;
    for (UINT64 p = 0; p < low_reserve_pages && p < g_pmm.total_pages; p++) {
        if (!bitmap_test((UINTN)p)) {
            bitmap_set((UINTN)p);
            g_pmm.free_pages--;
        }
    }

    // Build free list by linking free pages
    g_pmm.free_list_head = NULL;
    void* prev = NULL;
    for (UINTN i = 0; i < entry_count; i++) {
        EFI_MEMORY_DESCRIPTOR* desc = (EFI_MEMORY_DESCRIPTOR*)((UINT8*)mem_map + i * desc_size);
        if (desc->Type != EfiConventionalMemory) continue;

        UINT64 start_page = desc->PhysicalStart / PAGE_SIZE;
        UINT64 end_page   = start_page + desc->NumberOfPages;

        for (UINT64 p = start_page; p < end_page; p++) {
            if (p >= bm_start_page && p < bm_end_page) continue;
            if (p < low_reserve_pages) continue;
            void* page = (void*)(UINTN)(p * PAGE_SIZE);
            if (prev) {
                *(void**)prev = page;
            } else {
                g_pmm.free_list_head = page;
            }
            prev = page;
        }
    }
    if (prev) *(void**)prev = NULL;

    uefi_call_wrapper(ST->BootServices->FreePool, 1, mem_map);

    serial_write("PMM: init OK, total ");
    serial_write_hex(total_pages);
    serial_write(" pages (");
    serial_write_hex(total_free / (1024*1024));
    serial_write(" MB free), bitmap ");
    serial_write_hex(bitmap_pages);
    serial_write(" pages @ ");
    serial_write_hex(bitmap_addr);
    serial_write("\n");

    return EFI_SUCCESS;
}

void* pmm_alloc_pages(UINTN n) {
    if (n == 0) return NULL;

    clean_free_list();

    if (n == 1) {
        if (g_pmm.free_list_head == NULL) {
            serial_write("PMM: OOM (no free pages)\n");
            return NULL;
        }
        void* page = g_pmm.free_list_head;
        g_pmm.free_list_head = *(void**)page;
        *(void**)page = NULL;  // clear the next pointer

        UINTN idx = (UINTN)page / PAGE_SIZE;
        bitmap_set(idx);
        g_pmm.free_pages--;

        serial_write("PMM: alloc 1 page @ ");
        serial_write_hex((UINTN)page);
        serial_write("\n");
        return page;
    }

    // n > 1: scan bitmap for n consecutive free pages
    UINTN consecutive = 0;
    UINTN start_idx = 0;
    for (UINTN i = 0; i < g_pmm.total_pages; i++) {
        if (!bitmap_test(i)) {
            if (consecutive == 0) start_idx = i;
            consecutive++;
            if (consecutive == n) {
                // Found n consecutive free pages
                void* base = (void*)(UINTN)(start_idx * PAGE_SIZE);
                for (UINTN j = 0; j < n; j++) {
                    bitmap_set(start_idx + j);
                }
                g_pmm.free_pages -= n;

                serial_write("PMM: alloc ");
                serial_write_hex(n);
                serial_write(" pages @ ");
                serial_write_hex((UINTN)base);
                serial_write("\n");
                return base;
            }
        } else {
            consecutive = 0;
        }
    }

    serial_write("PMM: OOM (no contiguous free pages)\n");
    return NULL;
}

void pmm_free_pages(void* addr, UINTN n) {
    UINTN start_idx = (UINTN)addr / PAGE_SIZE;

    // Add freed pages to free list
    for (UINTN i = 0; i < n; i++) {
        UINTN idx = start_idx + i;
        bitmap_clear(idx);
        g_pmm.free_pages++;

        void* page = (void*)(UINTN)(idx * PAGE_SIZE);
        // Push to front of free list
        *(void**)page = g_pmm.free_list_head;
        g_pmm.free_list_head = page;
    }

    serial_write("PMM: free ");
    serial_write_hex(n);
    serial_write(" pages @ ");
    serial_write_hex((UINTN)addr);
        serial_write("\n");
    }


UINTN pmm_get_free_count() {
    return g_pmm.free_pages;
}

BOOLEAN pmm_is_page_free(void* addr) {
    UINTN idx = (UINTN)addr / PAGE_SIZE;
    if (idx >= g_pmm.total_pages) return FALSE;
    return !bitmap_test(idx);
}