#include "../include/paging.h"
#include <stdint.h>
#include "../include/multiboot.h"
#include "../include/string.h"

uint32_t kernel_page_directory[1024] __attribute__((aligned(4096)));
uint32_t kernel_page_table[1024] __attribute__((aligned(4096)));
uint32_t user_page_table[1024] __attribute__((aligned(4096)));

uint32_t *kpage_dir;

uint32_t setup_tmp_pgdir(uint32_t magic, uint32_t info)
{
	int n, pd;
	uint32_t pagedir_address, memksize;
	uint32_t *page_table;
	struct multiboot_info *mbi;

	if(magic != MULTIBOOT_BOOTLOADER_MAGIC) {
		/* 4MB of memory assumed */
		memksize = 4096;
	} else {
		mbi = (struct multiboot_info *)(PAGE_OFFSET + info);
		if(!(mbi->flags & MULTIBOOT_INFO_MEMORY)) {
			/* 4MB of memory assumed */
			memksize = 4096;
		} else {
			/* we need to add the first 1MB to memksize */
			memksize = (uint32_t)mbi->mem_upper + 1024;
			if(memksize > ((0xFFFFFFFF - PAGE_OFFSET) / 1024)) {
				memksize = (0xFFFFFFFF - PAGE_OFFSET) / 1024;
			}
		}
	}

    /* address becomes PAGE_OFFSET (0xC0000000) plus the size of memory minus the first page */
    /* TODO: check if I wrote the comment above correctly*/
	pagedir_address = PAGE_OFFSET + (memksize * 1024) - memksize;
	pagedir_address = PAGE_ALIGN(pagedir_address);

    /*
    memory pages:
    0 -> page directory
    1 -> page table
    */
	kpage_dir = (uint32_t *)pagedir_address;
	memset(kpage_dir, 0, PAGE_SIZE);

	pagedir_address += PAGE_SIZE;
	page_table = (uint32_t *)pagedir_address;
	memset(page_table, 0, memksize);

    /* if we have 4MB assumed earlier, the memksize will be 4096 and this loop will run while n < 1024 */
    /* so basically 4MB is 1024 pages */
	for(n = 0; n < memksize / sizeof(uint32_t); n++) {
		page_table[n] = (n << PAGE_SHIFT) | PAGE_PRESENT | PAGE_RW;
		if(!(n % 1024)) {
			pd = n / 1024;
			kpage_dir[pd] = (uint32_t)(pagedir_address + (PAGE_SIZE * pd) + GDT_BASE) | PAGE_PRESENT | PAGE_RW;
			kpage_dir[GET_PGDIR(PAGE_OFFSET) + pd] = (uint32_t)(pagedir_address + (PAGE_SIZE * pd) + GDT_BASE) | PAGE_PRESENT | PAGE_RW;
		}
	}
	return (uint32_t)kpage_dir - PAGE_OFFSET;
}

void enablePaging() {
    asm volatile (
        "mov %%cr0, %%eax\n"
        "or $0x80000000, %%eax\n"
        "mov %%eax, %%cr0\n"
        "jmp 1f\n"
        "1:\n"
        "mov %%cr3, %%eax\n" // Перезагрузка CR3 для сброса TLB
        "mov %%eax, %%cr3\n"
        "nop\n"
        :
        :
        : "eax", "memory"
    );
}

void loadPageDirectory(uint32_t* page_directory) {
    asm volatile (
        "mov %0, %%cr3\n" // Загружаем физический адрес каталога страниц в CR3
        :                 // Нет выходных операндов
        : "r" (page_directory) // Входной операнд — адрес каталога
        :                 // Нет изменяемых регистров
    );
}

void paging_init() {
    int i;

    // Проверяем физические адреса
    uint32_t pd_phys = (uint32_t)kernel_page_directory;
    uint32_t pt_phys = (uint32_t)kernel_page_table;
    uint32_t user_pt_phys = (uint32_t)user_page_table;

    if (pd_phys & 0xFFF || pt_phys & 0xFFF || user_pt_phys & 0xFFF) {
        while (1); // Ошибка выравнивания
    }

    // Заполняем каталог страниц
    for (i = 0; i < 1024; i++) {
        kernel_page_directory[i] = 0x00000002; // not present
    }

    // Identity mapping для 0x00000000–0x003FFFFF
    for (i = 0; i < 1024; i++) {
        kernel_page_table[i] = (i * 0x1000) | 0x3; // supervisor, present
    }
    kernel_page_directory[0] = pt_phys | 0x3;

    // Ring 3: 0x00400000–0x007FFFFF, User-accessible
    for (i = 0; i < 1024; i++) {
        user_page_table[i] = 0; // Изначально не отображено
    }
    kernel_page_directory[1] = user_pt_phys | 0x7; // Present, R/W, User

    // Рекурсивное отображение
    kernel_page_directory[1023] = pd_phys | 0x3;

    // Проверяем перед загрузкой
    if (kernel_page_directory[0] != (pt_phys | 0x3) ||
        kernel_page_directory[1023] != (pd_phys | 0x3)) {
        printf("OSHIBKA ZAPISI BLYA\n");
        while (1); // Ошибка записи
    }
    loadPageDirectory(kernel_page_directory);
    enablePaging();
}

void test_paging() {
    uint32_t* pd = (uint32_t*)0xFFFFF000;
    uint32_t expected = (uint32_t)kernel_page_table | 0x3;
    // Игнорируем дополнительные флаги, проверяем только базовый адрес и present
    if ((pd[0] & ~0xFFF) == (expected & ~0xFFF)) {
        // Успех
    } else {
        printf("RECURSION BROKEN!!\n");
        printf("pd[0] = 0x%X\n");
        printf("(uint32_t)first_page_table | 0x3 = %X", (uint32_t)kernel_page_table | 0x3);
        while (1); // Рекурсия сломана
    }
}