#include <stdint.h>
#include <string.h>
#include "mm/pmm.h"
#include "mm/vmm.h"
#include "mm/memory.h"
#include "stdio.h"
#include "fs/elf.h"

extern uintptr_t g_hhdm_offset;

#define USER_STACK_TOP 0x00007FFFFFFFE000ULL
#define USER_STACK_PAGES 4
#define USER_STACK_SIZE   (USER_STACK_PAGES * PAGE_SIZE)

#define PTE_PRESENT  (1ULL << 0)
#define PTE_WRITABLE (1ULL << 1)
#define PTE_USER     (1ULL << 2)



static uint64_t user_stack_phys_base = 0;

extern struct pagemap *kernel_pagemap;

static void setup_user_stack()
{
    // Allocate contiguous physical pages for the stack
    user_stack_phys_base = (uint64_t)pmm_alloc(USER_STACK_PAGES);
    if (!user_stack_phys_base) {
        printf("Failed to allocate user stack pages!\n");
        for (;;);
    }

    // Map each page individually into the current pagemap
    // Stack grows downward, so start from the bottom
    uintptr_t virt = USER_STACK_TOP - USER_STACK_SIZE;

    for (int i = 0; i < USER_STACK_PAGES; i++) {
        uint64_t phys = user_stack_phys_base + (i * PAGE_SIZE);

        bool success = vmm_map_page(
            kernel_pagemap,           // Use kernel_pagemap for now (later: process->process_pagemap)
            virt,
            phys,
            PAGE_READ | PAGE_WRITE | PAGE_USER,   // RW + User
            Size4KiB
        );

        if (!success) {
            printf("Failed to map user stack page at 0x%lx\n", virt);
            // TODO: cleanup previous pages if partial failure
            for (;;);
        }

        virt += PAGE_SIZE;
    }

    // Optional: zero the stack (good practice)
    // memset((void*)(USER_STACK_TOP - USER_STACK_SIZE), 0, USER_STACK_SIZE);
}



// usermode.c
__attribute__((naked))
void enter_user_mode(uint64_t rip, uint64_t rsp)
{
    asm volatile(
        "cli\n"
        "mov $0x1B, %ax\n" // Use User Data Selector (0x1B)
        "mov %ax, %ds\n"
        "mov %ax, %es\n"
        "mov %ax, %fs\n"
        "mov %ax, %gs\n"

        "pushq $0x1B\n"     // SS (User Data)
        "pushq %rsi\n"     // RSP
        "pushfq\n"
        // Optional: manually set IF bit in pushed RFLAGS if you want interrupts enabled
        "pop %rax\n"
        "or $0x200, %rax\n"
        "push %rax\n"
        "pushq $0x23\n"     // CS (User Code)
        "pushq %rdi\n"     // RIP
        "iretq\n"
    );
}

void start_userspace(void)
{
    void *entry = NULL;

    if (!ELF_Read("init.elf", &entry)) {
        printf("Failed to load init.elf\n");
        for(;;);
    }

    setup_user_stack();

    printf("Jumping to user entry: 0x%lx with stack top 0x%lx\n",
           (uint64_t)entry, USER_STACK_TOP);

    enter_user_mode((uint64_t)entry, USER_STACK_TOP);
}