#include "vmm.h"
#include "pmm.h"
#include "serial.h"

// 1. Обязательное выравнивание для CR3 и таблиц страниц
uint32_t page_directory[1024] __attribute__((aligned(4096)));
uint32_t first_page_table[1024] __attribute__((aligned(4096)));

extern void vmm_enable_paging(uint32_t* page_directory);

void vmm_init(void) {
    kprint_serial("VMM: Initializing paging...\n");

    // 2. Очищаем каталог (ставим Not Present)
    for (int i = 0; i < 1024; i++) {
        page_directory[i] = 0x00000002; // Read/Write, Not Present
    }

    // 3. Identity Map для первых 4 МБ
    // i * 4096 дает физический адрес страницы
    for (uint32_t i = 0; i < 1024; i++) {
        first_page_table[i] = (i * 0x1000) | 0x3; // Present + RW
    }

    // 4. Указываем каталогу на таблицу
    // (uint32_t)first_page_table — это уже физический адрес, 
    // если ядро по 1МБ и мы не используем сложные смещения.
    page_directory[0] = ((uint32_t)first_page_table) | 0x3;

    kprint_serial("DEBUG: page_directory is at: 0x");
    kprint_serial_hex((uint32_t)page_directory);
    kprint_serial("\n");

    kprint_serial("DEBUG: page_directory[0] contains: 0x");
    kprint_serial_hex(page_directory[0]);
    kprint_serial("\n");

    kprint_serial("VMM: Paging configured. Enabling CR0...\n");

    // 5. Включаем
    vmm_enable_paging(page_directory);

    kprint_serial("VMM: FUCK, IT WORKS!\n");
}

void vmm_map_page(uint32_t phys, uint32_t virt, uint32_t flags) {
    // Выковыриваем индексы из виртуального адреса
    uint32_t pd_idx = virt >> 22;          // Старшие 10 бит — индекс каталога
    uint32_t pt_idx = (virt >> 12) & 0x3FF; // Средние 10 бит — индекс таблицы

    if (!(page_directory[pd_idx] & 1)) {
        // 1. Выделяем физическую страницу под НОВУЮ таблицу
        uint32_t phys_table = (uint32_t)pmm_alloc_page(); 
        
        // 2. Ставим её в каталог (Present + Write)
        page_directory[pd_idx] = (phys_table & 0xFFFFF000) | 3;
        
        // 3. Очищаем новую таблицу (важно!)
        uint32_t* new_table = (uint32_t*)(page_directory[pd_idx] & 0xFFFFF000);
        for(int i=0; i<1024; i++) new_table[i] = 0;
    }

    // Достаем адрес таблицы страниц из каталога (обнуляем нижние 12 бит флагов)
    uint32_t* table = (uint32_t*)(page_directory[pd_idx] & 0xFFFFF000);

    // Запихиваем физический адрес и флаги в таблицу страниц
    // PAGE_SIZE у нас 4096, поэтому phys должен быть выровнен (нижние 12 бит затрутся флагами)
    table[pt_idx] = (phys & 0xFFFFF000) | flags;

    // ВАЖНО: Процессор кэширует таблицы в TLB (Translation Lookaside Buffer).
    // Если мы изменили маппинг существующей страницы, процессор об этом не узнает.
    // Нужно сбросить кэш для этого адреса инструкцией invlpg:
    __asm__ __volatile__("invlpg (%0)" : : "r"(virt) : "memory");
}

