#include "pmm.h"
#include "string.h"
#include "logger.h"
#include "../drivers/serial.h"

// Указатель на начало нашей битовой карты в памяти
static uint32_t* pmm_bitmap = NULL;
static uint32_t pmm_bitmap_size = 0;
static uint32_t total_pages = 0;

static pmm_stats_t global_stats;

// Вспомогательные инлайны для быстрой работы с битами ядра
static inline void bitmap_set(uint32_t page_index) {
    pmm_bitmap[page_index / 32] |= (1 << (page_index % 32));
}

static inline void bitmap_clear(uint32_t page_index) {
    pmm_bitmap[page_index / 32] &= ~(1 << (page_index % 32));
}

static inline int bitmap_test(uint32_t page_index) {
    return (pmm_bitmap[page_index / 32] & (1 << (page_index % 32))) != 0;
}

void pmm_init(struct multiboot_info* mbi, uint32_t kernel_end) {
    // Рассчитываем менеджер на 128 МБ памяти
    uint32_t mem_size = 128 * 1024 * 1024; 
    total_pages = mem_size / PAGE_SIZE;
    
    // Саму карту пихаем прямо за концом ядра (куда указывает линкер)
    pmm_bitmap = (uint32_t*)kernel_end;
    pmm_bitmap_size = total_pages / 8; // Размер карты в байтах (4096 байт)
    
    uint32_t bitmap_end = kernel_end + pmm_bitmap_size;

    // Шаг 1: Помечаем вообще всё как занятое (1). Безопасность превыше всего.
    memset(pmm_bitmap, 0xFF, pmm_bitmap_size);

    // Шаг 2: Парсим карту памяти от GRUB и открываем доступные регионы
    struct multiboot_mmap_entry* mmap = (struct multiboot_mmap_entry*)mbi->mmap_addr;
    
    while((uint32_t)mmap < mbi->mmap_addr + mbi->mmap_length) {
        // Если GRUB говорит, что этот кусок RAM свободен для юзания
        if (mmap->type == MULTIBOOT_MEMORY_AVAILABLE) {
            uint32_t start_addr = (uint32_t)mmap->addr;
            uint32_t end_addr = start_addr + (uint32_t)mmap->len;

            for (uint32_t addr = start_addr; addr < end_addr; addr += PAGE_SIZE) {
                uint32_t page_idx = addr / PAGE_SIZE;
                if (page_idx < total_pages) {
                    bitmap_clear(page_idx); // Ставим 0 (страница свободна)
                }
            }
        }
        // Шагаем к следующей записи в таблице Multiboot
        mmap = (struct multiboot_mmap_entry*)((uint32_t)mmap + mmap->size + sizeof(mmap->size));
    }

    // Шаг 3: Жестко защищаем саму ОС! 
    // Всё, что лежит от физического адреса 0x0 до bitmap_end (код ядра, стек, сама битовая карта),
    // мы принудительно забиваем единицами, чтобы аллокатор туда никогда не залез.
    for (uint32_t addr = 0; addr < bitmap_end; addr += PAGE_SIZE) {
        bitmap_set(addr / PAGE_SIZE);
    }

    log_info(LOG_NORMAL, "pmm works.");
}

void* pmm_alloc_page(void) {
    // Тупо линейно ищем первый попавшийся бит, равный 0
    for (uint32_t i = 0; i < total_pages; i++) {
        if (!bitmap_test(i)) { 
            bitmap_set(i); // Застолбили страницу (поставили 1)
            return (void*)(i * PAGE_SIZE); // Считаем её реальный физический адрес
        }
    }
    
    log_error(LOG_NORMAL, "PMM: FUCK, WHERE IS ALL OUR FUCKING MEMORY?!");
    return NULL; 
}

void pmm_free_page(void* addr) {
    uint32_t page_idx = ((uint32_t)addr) / PAGE_SIZE;
    if (page_idx < total_pages) {
        bitmap_clear(page_idx); // Сбросили в 0, страница снова чиста
    }
}

void pmm_get_stats(void) {
    kprint_serial("Someone wants memory stats! Running the test!\n");
    uint32_t used = 0;
    uint32_t free = 0;

    for (uint32_t i = 0; i < total_pages; i++) {
        if (bitmap_test(i)) {
            used++;
        } else {
            free++;
        }
    }

    kprint_serial("Okay, i am collected the stats. ");

    // Пишем в ГЛОБАЛЬНУЮ переменную
    global_stats.total_pages = total_pages;
    global_stats.used_pages = used;
    global_stats.free_pages = free;
    global_stats.total_memory = total_pages << 12;
    global_stats.used_memory = used << 12;
    global_stats.free_memory = free << 12;

    kprint_serial("Goodbye!\n");
}

// Добавим функцию-геттер, которая просто вернет указатель на эту глобальную память
pmm_stats_t* pmm_get_stats_ptr(void) {
    return &global_stats;
}