#include "vmm.h"
#include "../drivers/serial.h"
#include "string.h"

// Каталог страниц (1024 элемента по 4 байта = 4КБ)
static uint32_t page_directory[1024] __attribute__((aligned(4096)));

// Первая таблица страниц для маппинга первых 4 Мегабайт памяти
static uint32_t first_page_table[1024] __attribute__((aligned(4096)));

extern void vmm_enable_paging(uint32_t* directory);

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

    // 1. Очищаем каталог страниц (ставим не-present)
    // Ставим флаг 0x2 (Read/Write, но Not Present)
    for (int i = 0; i < 1024; i++) {
        page_directory[i] = 0x00000002;
    }

    // 2. Делаем Identity Mapping для первых 4 МБ (0x00000000 - 0x003FFFFF)
    // Пробегаемся по всей первой таблице страниц
    for (uint32_t i = 0; i < 1024; i++) {
        // Физический адрес = i * 4096
        // Флаги: Present (0x1) + Read/Write (0x2)
        first_page_table[i] = (i * 4096) | 3;
    }

    // 3. Засовываем первую таблицу страниц в первый элемент каталога
    // Адрес таблицы | флаги Present + Read/Write
    page_directory[0] = ((uint32_t)first_page_table) | 3;

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

    // 4. Прыгаем в ассемблер, чтобы загрузить CR3 и выставить флаг в CR0
    vmm_enable_paging(page_directory);

    kprint_serial("VMM: HOLY SHIT IT WORKS! Paging is active!\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 бит — индекс таблицы

    // Проверяем, создана ли таблица страниц для этого региона
    // Наше ядро пока создало только одну таблицу для первых 4 МБ (pd_idx == 0)
    if (!(page_directory[pd_idx] & 1)) {
        // Если таблицы нет, по-хорошему надо выделить 4КБ через pmm_alloc_page(),
        // очистить её и засунуть в каталог. Но пока выкинем ошибку для отладки:
        kprint_serial("VMM ERROR: Page table not allocated for this virtual address!\n");
        return;
    }

    // Достаем адрес таблицы страниц из каталога (обнуляем нижние 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");
}