/*
 * @author  xamidev <xamidev@riseup.net>
 * @brief   Process linked list implementation
 * @license GPL-3.0-only
 */

#include "mem/paging.h"
#include "mem/vmm.h"
#include <stddef.h>
#include <sched/process.h>
#include <mem/kheap.h>
#include <kernel.h>
#include <string/string.h>
#include <arch/gdt.h>
#include <config.h>
#include <io/serial/serial.h>
#include <io/term/flanterm.h>
#include <mem/utils.h>

extern struct flanterm_context* ft_ctx;

struct process* processes_list;
struct process* current_process;

extern uint64_t *kernel_pml4;

size_t next_free_pid = 0;

/*
 * process_init - Initializes process list
 */
void process_init()
{
    processes_list = NULL;
    current_process = NULL;
}

/*
 * process_display_list - Debug function to display processes
 * @processes_list: head of the process linked list
 *
 * This function prints the linked list of processes
 * to the DEBUG output.
 */
void process_display_list(struct process* processes_list)
{
    int process_view_id = 0;
    struct process* tmp = processes_list;
    while (tmp != NULL) {
        DEBUG("{%d: %p} -> ", process_view_id, tmp);
        tmp = tmp->next;
        process_view_id++;
    }
    DEBUG("NULL");
}

/*
 * process_create - Create a process
 * @name:     name of the process
 * @function: beginning of process executable code
 * @arg:      (optional) argument provided to process
 *
 * This function creates a process, gives it all
 * necessary context and a stack, and adds the
 * process to the linked list.
 *
 * Return:
 * <proc> - pointer to created process
 */
struct process* process_create(char* name, void(*function)(void*), void* arg)
{
    CLEAR_INTERRUPTS;
    struct process* proc = (struct process*)kmalloc(sizeof(struct process));
    struct cpu_status* ctx = (struct cpu_status*)kmalloc(sizeof(struct cpu_status));

    // No more memory?
    if (!proc) return NULL;
    if (!ctx) return NULL;

    strncpy(proc->name, name, PROCESS_NAME_MAX);
    proc->pid = next_free_pid++;
    proc->status = READY;

    uint64_t* stack_top = (uint64_t*)kalloc_stack();
    // push return address to the stack so when "ret" hits we jmp to exit instead of idk what
    // stack grows DOWNWARDS!!
    *(--stack_top) = (uint64_t)process_exit;

    proc->context = ctx;
    proc->context->iret_ss = KERNEL_DATA_SEGMENT; // process will live in kernel mode
    proc->context->iret_rsp = (uint64_t)stack_top;
    proc->context->iret_flags = 0x202; //bit 2 and 9 set (Interrupt Flag)
    proc->context->iret_cs = KERNEL_CODE_SEGMENT;
    proc->context->iret_rip = (uint64_t)function; // beginning of executable code
    proc->context->rdi = (uint64_t)arg; // 1st arg is in rdi (as per x64 calling convention)
    proc->context->rbp = 0; 

    // Kernel PML4 as it already maps code/stack (when switching to userland we'll have to change that)
    proc->root_page_table = kernel_pml4;

    proc->kernel_stack = kalloc_stack();

    proc->next = 0;

    process_add(&processes_list, proc);

    SET_INTERRUPTS;
    return proc;
}

/*
 * process_add - Add a process to the end of the linked list
 * @processes_list: pointer to the head of the linked list
 * @process:        process to add at the end of the linked list
 */
void process_add(struct process** processes_list, struct process* process)
{
    if (!process) return;
    process->next = NULL;

    if (*processes_list == NULL) {
        // List is empty
        *processes_list = process;
        return;
    }

    struct process* tmp = *processes_list;
    while (tmp->next != NULL) {
        tmp = tmp->next;
    }
    // We're at last process before NULL
    tmp->next = process;
}

/*
 * process_delete - Delete a process from the linked list
 * @processes_list: pointer to head of linked list
 * @process:        the process to delete from the list
 */
void process_delete(struct process** processes_list, struct process* process)
{
    if (!processes_list || !*processes_list || !process) return;

    if (*processes_list == process) {
        // process to delete is at head
        *processes_list = process->next;
        process->next = NULL;
        kfree(process);
        return;
    }

    struct process* tmp = *processes_list;
    while (tmp->next && tmp->next != process) {
        tmp = tmp->next;
    }

    if (tmp->next == NULL) {
        // Didn't find the process
        return;
    }

    // We're at process before the one we want to delete
    tmp->next = process->next;
    process->next = NULL;
    kfree(process);
}

/*
 * process_get_next - Get the next process (unused)
 * @process: pointer to process
 *
 * Return:
 * <process->next> - process right after the one specified
 */
struct process* process_get_next(struct process* process)
{
    if (!process) return NULL;
    return process->next;
}

/*
 * process_exit - Exit from a process
 *
 * This function is pushed to all process stacks, as a last
 * return address. Once the process is done executing, it 
 * ends up here.
 *
 * Process is marked as DEAD, and then execution loops.
 * Next time the scheduler sees the process, it will
 * automatically delete it from the linked list.
 */
void process_exit()
{
    DEBUG("Exiting from process '%s'", current_process->name);
    CLEAR_INTERRUPTS;
    if (current_process) {
        current_process->status = DEAD;
    }
    SET_INTERRUPTS;
    
    for (;;) {
        asm("hlt");
    }
}

/*
 * process_jump_to_user - Jump to userland
 * @stack_top: Address of the top of the user stack
 * @user_code: Address of the first instruction of user code
 */
void process_jump_to_user(uintptr_t stack_top, uintptr_t user_code)
{
    // 0x20 | 3 = 0x23 (user data segment | 3)
    // 0x18 | 3 = 0x1B (user code segment | 3)
    asm volatile(" \
        push $0x23 \n\
        push %0 \n\
        push $0x202 \n\
        push $0x1B \n\
        push %1 \n\
        iretq \n\
        " :: "r"(stack_top), "r"(user_code));
}

extern struct tss tss;

/*
 * process_create_user_raw - Create a new user process from raw binary
 * @file: pointer to beginning of binary
 * @size: size of the binary
 * @name: name for the new process
 *
 * This function takes an executable loaded in memory
 * and maps its code, a user stack, sets the TSS RSP0
 * for interrupts, and finally jumps to the user code.
 */
void process_create_user_raw(char* file, int size, char* name)
{
    CLEAR_INTERRUPTS;
    struct process* proc = (struct process*)kmalloc(sizeof(struct process));
    struct cpu_status* ctx = (struct cpu_status*)kmalloc(sizeof(struct cpu_status));

    if (!proc || !ctx) panic(NULL, "out of memory while creating user process");

    memset(proc, 0, sizeof(struct process));
    memset(ctx, 0, sizeof(struct cpu_status));

    strncpy(proc->name, name, PROCESS_NAME_MAX);
    proc->pid = next_free_pid++;
    proc->status = READY;
    proc->next = 0;
    proc->context = ctx;
    proc->context->iret_ss = USER_DATA_SEGMENT | 3;
    proc->context->iret_cs = USER_CODE_SEGMENT | 3;
    proc->context->iret_flags = 0x202; // Interrupt Flag set

        /* Set basic entries for the process's File Descriptor Table */
    proc->fdt[0].fd = 0;
    proc->fdt[0].open = true;
    proc->fdt[0].cursor = 0;
    strncpy(proc->fdt[0].filename, "stdin", PROCESS_NAME_MAX - 1);

    proc->fdt[1].fd = 1;
    proc->fdt[1].open = true;
    proc->fdt[1].cursor = 0;
    strncpy(proc->fdt[1].filename, "stdout", PROCESS_NAME_MAX - 1);

    proc->fdt[2].fd = 2;
    proc->fdt[2].open = true;
    proc->fdt[2].cursor = 0;
    strncpy(proc->fdt[2].filename, "stderr", PROCESS_NAME_MAX - 1);

    proc->next_free_fd = 3; // file descriptors are also bump-allocated

    void* exec_addr = (void*)file;
    uint64_t exec_size = size;

    uint64_t* user_pml4 = vmm_create_address_space();
    if (!user_pml4) panic(NULL, "failed to create user address space");
    proc->root_page_table = user_pml4;

    uintptr_t stack_top = vmm_alloc_user_stack(user_pml4);
    uint64_t code = vmm_alloc_user_code(user_pml4, exec_addr, exec_size);
    
    proc->context->iret_rsp = stack_top;
    proc->context->iret_rip = code;
    proc->kernel_stack = kalloc_stack();
    if (!proc->kernel_stack) panic(NULL, "failed to allocate kernel stack");

    // Copy code into user pages; for that we need to temporarily switch to the user pml4
    load_cr3(VIRT_TO_PHYS((uint64_t)user_pml4));
    memcpy((uint64_t*)code, exec_addr, exec_size);
    load_cr3(VIRT_TO_PHYS((uint64_t)kernel_pml4));

    process_add(&processes_list, proc);
    DEBUG("user process '%s' (pid=%u) enqueued for scheduling", name, proc->pid);
    SET_INTERRUPTS;
}