#include "vmm.h"
#include "pmm.h"
#include "memory.h"
#include "stdio.h"


struct pagemap *kernel_pagemap = NULL;

__attribute__((used, section(".limine_requests")))
volatile struct limine_hhdm_request hhdm_request = {
    .id = LIMINE_HHDM_REQUEST_ID, .revision = 0
};

__attribute__((used, section(".limine_requests")))
volatile struct limine_executable_address_request kernel_address_request = {
    .id = LIMINE_EXECUTABLE_ADDRESS_REQUEST_ID, .revision = 0
};

__attribute__((used, section(".limine_requests")))
volatile struct limine_paging_mode_request paging_mode_request = {
    .id = LIMINE_PAGING_MODE_REQUEST_ID, .revision = 0, .mode = LIMINE_PAGING_MODE_X86_64_4LVL
};

static uint64_t *get_next_level(uint64_t *top_level, size_t idx, bool allocate) {
    if (top_level[idx] & 1)
        return (uint64_t *)((top_level[idx] & ~0xFFFULL) + MEM_PHYS_OFFSET);

    if (!allocate) return NULL;

    void *next = pmm_allocz(1);
    if (!next) return NULL;

    top_level[idx] = (uint64_t)next | 0b111;
    return (uint64_t *)((uintptr_t)next + MEM_PHYS_OFFSET);
}

void vmm_init(struct limine_memmap_entry **memmap, size_t memmapentries) {
    if (!hhdm_request.response) {
        printf("PRAISE BE TO LIMINE WE DON'T HAE HHDM RESP");
    }
    size_t memmap_entries = memmapentries - 1;
    kernel_pagemap = kmalloc(sizeof(struct pagemap));
    if (!kernel_pagemap) {
        printf("VMM: OOm allocating kernel_pagemap\n");
        while(1) asm volatile ("hlt");
    }
    spinlock_init(&kernel_pagemap->lock);
    kernel_pagemap->top_level = (uint64_t *)((uintptr_t)pmm_allocz(1) + MEM_PHYS_OFFSET);
    if (!kernel_pagemap->top_level) {
        printf("VMM: OOM allocating PML4\n");
        while (1) asm volatile ("hlt");
    }

    // Identity map lower half + HHDM
    for (uint64_t p = 256; p < 512; p++)
        get_next_level(kernel_pagemap->top_level, p, true);

    // Map all physical memory (HHDM)
    for (uint64_t p = 0; p < 0x8000000000ULL; p += 0x200000)  // 2 MiB pages up to 256 GiB
        vmm_map_page(kernel_pagemap, p + MEM_PHYS_OFFSET, p, PAGE_READ | PAGE_WRITE, Size2MiB);
    


    

    // Map kernel sections with correct permissions
    extern char __text_start_addr[], __text_end_addr[];
    extern char __rodata_start_addr[], __rodata_end_addr[];
    extern char __data_start_addr[], __data_end_addr[];
    extern char __bss_start_addr[], __bss_end_addr[];

    uintptr_t text_start = ALIGN_DOWN((uintptr_t)__text_start_addr, PAGE_SIZE);
    uintptr_t text_end   = ALIGN_UP((uintptr_t)__text_end_addr, PAGE_SIZE);
    uintptr_t rodata_start = ALIGN_DOWN((uintptr_t)__rodata_start_addr, PAGE_SIZE);
    uintptr_t rodata_end   = ALIGN_UP((uintptr_t)__rodata_end_addr, PAGE_SIZE);
    uintptr_t data_start = ALIGN_DOWN((uintptr_t)__data_start_addr, PAGE_SIZE);
    uintptr_t data_end   = ALIGN_UP((uintptr_t)__data_end_addr, PAGE_SIZE);
    uintptr_t bss_start = ALIGN_DOWN((uintptr_t)__bss_start_addr, PAGE_SIZE);
    uintptr_t bss_end   = ALIGN_UP((uintptr_t)__bss_end_addr, PAGE_SIZE);

    uint64_t paddr = kernel_address_request.response->physical_base;
    uint64_t vaddr = kernel_address_request.response->virtual_base;


    // Map everything from kernel load address up to .text (this is the requests section)
    for (uintptr_t va = vaddr; va < text_start; va += PAGE_SIZE) {
        uint64_t pa = va - vaddr + paddr;
        vmm_map_page(kernel_pagemap, va, pa,
                     PAGE_READ | PAGE_WRITE | PAGE_NO_EXECUTE, Size4KiB);
    }
    for (uintptr_t va = text_start; va < text_end; va += PAGE_SIZE) {
        uint64_t pa = va - vaddr + paddr;
        vmm_map_page(kernel_pagemap, va, pa, PAGE_READ, Size4KiB);
    }
    // .rodata : R-- + NX
    for (uintptr_t va = rodata_start; va < rodata_end; va += PAGE_SIZE) {
        uint64_t pa = va - vaddr + paddr;
        vmm_map_page(kernel_pagemap, va, pa, PAGE_READ | PAGE_NO_EXECUTE, Size4KiB);
    }

    // .data : RW- + NX
    for (uintptr_t va = data_start; va < data_end; va += PAGE_SIZE) {
        uint64_t pa = va - vaddr + paddr;
        vmm_map_page(kernel_pagemap, va, pa, PAGE_READ | PAGE_WRITE | PAGE_NO_EXECUTE, Size4KiB);
    }

    // .bss : RW- + NX (stack lives here)
    for (uintptr_t va = bss_start; va < bss_end; va += PAGE_SIZE) {
        uint64_t pa = va - vaddr + paddr;
        vmm_map_page(kernel_pagemap, va, pa, PAGE_READ | PAGE_WRITE | PAGE_NO_EXECUTE, Size4KiB);
    }

    vmm_switch_pagemap(kernel_pagemap);
    printf("VMM: switched to new kernel pagemap\n");
}

void vmm_switch_pagemap(struct pagemap *pagemap) {
    uint64_t phys_cr3 = (uintptr_t)pagemap->top_level - MEM_PHYS_OFFSET;
    printf("VMM: loading CR3 with PML4 @ 0x%lx\n", phys_cr3);
    asm volatile("mov %0, %%cr3" : : "r"((uint64_t)pagemap->top_level - MEM_PHYS_OFFSET) : "memory");
    printf("rah");
}

bool vmm_map_page(struct pagemap *pagemap, uint64_t virt, uint64_t phys,
                  uint64_t flags, enum page_size pg_size) {
    spinlock_acquire_or_wait(&pagemap->lock);

	// Calculate the indices in the various tables using the virtual address
	size_t pml5_entry = (virt & ((uint64_t)0x1FF << 48)) >> 48;
	size_t pml4_entry = (virt & ((uint64_t)0x1FF << 39)) >> 39;
	size_t pml3_entry = (virt & ((uint64_t)0x1FF << 30)) >> 30;
	size_t pml2_entry = (virt & ((uint64_t)0x1FF << 21)) >> 21;
	size_t pml1_entry = (virt & ((uint64_t)0x1FF << 12)) >> 12;

	uint64_t *pml5, *pml4, *pml3, *pml2, *pml1;

	if (paging_mode_request.response->mode == LIMINE_PAGING_MODE_X86_64_5LVL) {
		pml5 = pagemap->top_level;
		goto level5;
	} else {
		pml4 = pagemap->top_level;
		goto level4;
	}

level5:
	pml4 = get_next_level(pml5, pml5_entry, true);
	if (pml4 == NULL) {
		goto die;
	}
level4:
	pml3 = get_next_level(pml4, pml4_entry, true);
	if (pml3 == NULL) {
		goto die;
	}

	pml2 = get_next_level(pml3, pml3_entry, true);
	if (pml2 == NULL) {
		goto die;
	}

	if (pg_size == Size2MiB) {
		pml2[pml2_entry] = phys | flags | (1 << 7);
		spinlock_drop(&pagemap->lock);
		return true;
	}

	pml1 = get_next_level(pml2, pml2_entry, true);
	if (pml1 == NULL) {
	die:
		spinlock_drop(&pagemap->lock);
		return false;
	}

	pml1[pml1_entry] = phys | flags;

    spinlock_drop(&pagemap->lock);
    return true;  // simplified for this response
}

bool vmm_unmap_page(struct pagemap *pagemap, uintptr_t virt, bool locked) {
	if (!locked) {
		spinlock_acquire_or_wait(&pagemap->lock);
	}

	size_t pml5_entry = (virt & ((uint64_t)0x1FF << 48)) >> 48;
	size_t pml4_entry = (virt & ((uint64_t)0x1FF << 39)) >> 39;
	size_t pml3_entry = (virt & ((uint64_t)0x1FF << 30)) >> 30;
	size_t pml2_entry = (virt & ((uint64_t)0x1FF << 21)) >> 21;
	size_t pml1_entry = (virt & ((uint64_t)0x1FF << 12)) >> 12;

	uint64_t *pml5, *pml4, *pml3, *pml2, *pml1;

	if (paging_mode_request.response->mode == LIMINE_PAGING_MODE_X86_64_5LVL) {
		pml5 = pagemap->top_level;
		goto level5;
	} else {
		pml4 = pagemap->top_level;
		goto level4;
	}

level5:
	pml4 = get_next_level(pml5, pml5_entry, false);
	if (pml4 == NULL) {
		goto die;
	}
level4:
	pml3 = get_next_level(pml4, pml4_entry, false);
	if (pml3 == NULL) {
		goto die;
	}

	pml2 = get_next_level(pml3, pml3_entry, false);
	if (pml2 == NULL) {
		goto die;
	}

	pml1 = get_next_level(pml2, pml2_entry, false);
	if (pml1 == NULL) {
		goto die;
	}

	if ((pml1[pml1_entry] & 1) == 0) {
	die:
		if (!locked) {
			spinlock_drop(&pagemap->lock);
		}
		return false;
	}

	pml1[pml1_entry] = 0;

	asm volatile("invlpg (%0)" : : "r"(virt) : "memory");

	if (!locked) {
		spinlock_drop(&pagemap->lock);
	}
	return true;
}

uint64_t *vmm_virt_to_pte(struct pagemap *pagemap, uintptr_t virt_addr,
						  bool allocate) {
	//	spinlock_acquire_or_wait(&pagemap->lock);
	size_t pml5_entry = (virt_addr & ((uint64_t)0x1FF << 48)) >> 48;
	size_t pml4_entry = (virt_addr & ((uint64_t)0x1FF << 39)) >> 39;
	size_t pml3_entry = (virt_addr & ((uint64_t)0x1FF << 30)) >> 30;
	size_t pml2_entry = (virt_addr & ((uint64_t)0x1FF << 21)) >> 21;
	size_t pml1_entry = (virt_addr & ((uint64_t)0x1FF << 12)) >> 12;

	uint64_t *pml5, *pml4, *pml3, *pml2, *pml1;

	if (paging_mode_request.response->mode == LIMINE_PAGING_MODE_X86_64_5LVL) {
		pml5 = pagemap->top_level;
		goto level5;
	} else {
		pml4 = pagemap->top_level;
		goto level4;
	}

level5:
	pml4 = get_next_level(pml5, pml5_entry, allocate);
	if (pml4 == NULL) {
		goto die;
	}
level4:
	pml3 = get_next_level(pml4, pml4_entry, allocate);
	if (pml3 == NULL) {
		goto die;
	}

	pml2 = get_next_level(pml3, pml3_entry, allocate);
	if (pml2 == NULL) {
		goto die;
	}

	pml1 = get_next_level(pml2, pml2_entry, allocate);
	if (pml1 == NULL) {
	die:
		spinlock_drop(&pagemap->lock);
		return NULL;
	}

	spinlock_drop(&pagemap->lock);
	return &pml1[pml1_entry];
}


uint64_t vmm_virt_to_phys(struct pagemap *pagemap, uint64_t virt) {
    spinlock_acquire_or_wait(&pagemap->lock);

    size_t pml4_entry = (virt & ((uint64_t)0x1FF << 39)) >> 39;
    size_t pml3_entry = (virt & ((uint64_t)0x1FF << 30)) >> 30;
    size_t pml2_entry = (virt & ((uint64_t)0x1FF << 21)) >> 21;
    size_t pml1_entry = (virt & ((uint64_t)0x1FF << 12)) >> 12;

    uint64_t *pml4 = pagemap->top_level;

    uint64_t *pml3 = get_next_level(pml4, pml4_entry, false);
    if (!pml3) goto fail;

    uint64_t *pml2 = get_next_level(pml3, pml3_entry, false);
    if (!pml2) goto fail;

    // Check for 2MiB huge page (PS bit = bit 7)
    if (pml2[pml2_entry] & (1 << 7)) {
        uint64_t phys = (pml2[pml2_entry] & ~0x1FFFFFULL) // 2MiB-aligned base
                      + (virt & 0x1FFFFFULL);             // + offset within 2MiB
        spinlock_drop(&pagemap->lock);
        return phys;
    }

    uint64_t *pml1 = get_next_level(pml2, pml2_entry, false);
    if (!pml1 || !(pml1[pml1_entry] & 1)) goto fail;

    uint64_t phys = (pml1[pml1_entry] & 0x000ffffffffff000ULL)
                  + (virt & 0xFFFULL);
    spinlock_drop(&pagemap->lock);
    return phys;

fail:
    spinlock_drop(&pagemap->lock);
    printf("Invalid Phys!\n");
    return INVALID_PHYS;
}