#include "../include/syscalls.h"
#include <stddef.h>
#include "../include/stdio.h"
#include "../include/liballoc.h"
#include "../include/paging.h"
#include "../include/string.h"
#include "../include/screen.h"
#include "../include/fat.h"
#include "../include/vfs.h"
#include "../include/isr.h"
#include "../include/task.h"

extern void switchProcess(Process* next);

#define STDIN_FILENO 0
#define STDOUT_FILENO 1
#define STDERR_FILENO 2

#define EBADF 9

#define EXECVE_MAX_ARGUMENT_SIZE 255

const char* flags_to_mode_str(int flags) {
    static char mode[4] = {0};
    int accmode = flags & 3;

    if (accmode == 0) {         // O_RDONLY
        mode[0] = 'r';
        mode[1] = '\0';
    } else if (accmode == 1) {  // O_WRONLY
        if (flags & 0x400) {    // O_APPEND
            mode[0] = 'a';
            mode[1] = '\0';
        } else {
            mode[0] = 'w';
            mode[1] = '\0';
        }
    } else if (accmode == 2) {  // O_RDWR
        if (flags & 0x400) {    // O_APPEND
            mode[0] = 'a';
            mode[1] = '+';
            mode[2] = '\0';
        } else if (flags & 0x200) { // O_TRUNC
            mode[0] = 'w';
            mode[1] = '+';
            mode[2] = '\0';
        } else {
            mode[0] = 'r';
            mode[1] = '+';
            mode[2] = '\0';
        }
    } else {
        mode[0] = 'r';
        mode[1] = '\0';
    }

    return mode;
}

int sys_exit(TrapFrame *tf)
{
    int error_code = tf->ebx;
    debug_log("killing task\n");
    task_kill(current);
    asm("sti");
    while(1){}
    return error_code;
}

extern void trapret(void);

int sys_fork(TrapFrame *tf)
{
    debug_log("FORKFORKFORKFORK\n");
    scheduler_lock();

    uint32_t* new_pd = copy_page_dir(current->pagedir);
    if (!new_pd) {
        debug_log("Failed to copy page directory\n");
        return -1;
    }
    debug_log("new_pd: %X\n", new_pd);

    void *p_physical = alloc_page();
    if (!p_physical) {
        destroy_page_dir(new_pd);
        return -1;
    }
    Process* child = (Process*)((uint32_t)p_physical + 0xC0000000);
    map_page(p_physical, (void*)child, PAGE_PRESENT | PAGE_RW);
    memset(child, 0, sizeof(Process));

    child->pid = ++pid_counter;
    child->state = Ready;
    child->ring = current->ring;
    child->pagedir = new_pd;

    void *kstack_physical = alloc_page();
    if (!kstack_physical) {
        destroy_page_dir(new_pd);
        free_page(p_physical);
        debug_log("Failed to allocate kernel stack\n");
        return -1;
    }
    child->kstack = (char*)((uint32_t)kstack_physical + 0xC0000000);
    map_page(kstack_physical, child->kstack, PAGE_PRESENT | PAGE_RW);
    debug_log("child->kstack: %X\n", child->kstack);

    uint8_t* sp = (uint8_t*)(child->kstack + KSTACKSIZE);

    // 1. Place the TrapFrame on the child's stack.
    sp -= sizeof(TrapFrame);
    child->tf = (TrapFrame*)sp;

    memcpy(&child->tf->gs, &tf->gs, 12 * sizeof(uint32_t));
    uint32_t* parent_cpu_state_ptr = (uint32_t*)&tf->interrupt;
    child->tf->eip          = parent_cpu_state_ptr[0];
    child->tf->cs           = parent_cpu_state_ptr[1];
    child->tf->eflags       = parent_cpu_state_ptr[2];
    child->tf->usermode_esp = parent_cpu_state_ptr[3];
    child->tf->usermode_ss  = parent_cpu_state_ptr[4];

    child->tf->eax = 0;       // Child returns 0
    debug_log("Parent returns PID: %X, Child returns 0\n", child->pid);

    sp -= sizeof(Context);
    child->context = (Context*)sp;
    memset(child->context, 0, sizeof(Context));

    // 5. Set the child's starting instruction pointer to trapret.
    // When the child is scheduled, it will execute trapret, which will
    // restore the registers from the TrapFrame and iret to user mode.
    child->context->eip = (uint32_t)trapret;

    // 6. IMPORTANT: Set the child's kernel stack pointer for the scheduler.
    child->kesp = (uint32_t)sp;

    for (int i = 0; i < MAX_OPEN_FILES; i++) {
        child->file_descriptors[i] = current->file_descriptors[i];
    }
    debug_log("File descriptors copied\n");

    Process* curr = queue;
    while (curr->next) {
        curr = curr->next;
    }
    curr->next = child;
    debug_log("Child added to queue\n");
    
    scheduler_unlock(); // Release the lock

    return child->pid;
}

#include "../../include/keyboard.h"
int sys_read(TrapFrame *tf)
{
    int fd = tf->ebx;
    char *user_buf = (char *)tf->ecx;
    size_t count = tf->edx;

    if (fd == STDIN_FILENO) {
        scheduler_lock();
        while (!stdin_has_line) {//while enter has not been pressed
            current->state = Waiting;
            current->waiting_reason = STDIN;
            schedule();
        }
        scheduler_unlock();

        //find the pos of the first \n
        int line_len = 0;
        while (line_len < stdin_pos && stdin_buffer[line_len] != '\n') line_len++;
        if (line_len < stdin_pos) line_len++;

        //copy the input to the user buffer
        int to_copy = (line_len < count) ? line_len : count;
        memcpy(user_buf, stdin_buffer, to_copy);

        memmove(stdin_buffer, stdin_buffer + to_copy, stdin_pos - to_copy);
        stdin_pos -= to_copy;

        stdin_has_line = false;
        for (int i = 0; i < stdin_pos; i++) {
            if (stdin_buffer[i] == '\n') {
                stdin_has_line = true;
                break;
            }
        }

        return to_copy;
    }
    else if(fd < 3)
    {
        return -EBADF;
    }
    else
    {
        if(current->file_descriptors[fd] == NULL)
        {
            return -EBADF;
        }
        return vfs_read_file_length(current->file_descriptors[fd], (uint8_t*)user_buf, count);
        //return 0;//fat_fread(current->file_descriptors[tf->ebx], (void*)tf->ecx, tf->edx);
    }
}

int sys_write(TrapFrame *tf)
{
    if(tf->ebx < 3)
    {
        //stdin, stdout or stderr
        if(tf->ebx == STDOUT_FILENO)
        {
            //debug_log("\n=================SYS_WRITE output to stdout=====================\n");
            print((char*)tf->ecx, tf->edx);
            //debug_log("\n=================END OF THAT SHIT=====================\n");
            return tf->edx;
        }
        if(tf->ebx == STDERR_FILENO)
        {
            uint8_t prev_color = terminal_getcolor();
            terminal_setcolor(VGA_COLOR_RED);
            print((char*)tf->ecx, tf->edx);
            terminal_setcolor(prev_color);
            return tf->edx;
        }
    }
    else
    {
        if(current->file_descriptors[tf->ebx] == NULL)
        {
            return -EBADF;
        }
        return 0;//fat_fwrite(current->file_descriptors[tf->ebx], (void*)tf->ecx, tf->edx);
    }
}

int sys_open(TrapFrame *tf)
{
    const char* filename = (const char*)tf->ebx;
    int flags = tf->ecx;

    int fd = -1;
    const char* mode_str = flags_to_mode_str(flags);

    //debug_log("\n====filename: %s\n", filename);
    //debug_log("\n====mode_str: %s\n", mode_str);

    for(int i = 3; i < MAX_OPEN_FILES; i++)
    {
        //debug_log("\n====i: %X\n", i);
        if(current->file_descriptors[i] == NULL)
        {
            current->file_descriptors[i] = (l9660_file*)malloc(sizeof(l9660_file));  // Kernel malloc
            l9660_openat(current->file_descriptors[i], root_dir, filename);
            //debug_log("\n====FILE_DESCRIPTOR: %X\n", i);
            int result = 0;//fat_fopen(current->file_descriptors[i], filename, mode_str);
            if (result < 0) {
                free(current->file_descriptors[i]);  // Free on failure
                current->file_descriptors[i] = NULL;
                return result;
            }
            fd = i;
            break;
        }
    }

    return fd;
}

int sys_close(TrapFrame *tf)
{
    int fd = tf->ebx;
    //debug_log("SYS_CLOSE: fd = 0x%X", fd);
    if(current->file_descriptors[fd] != NULL)
    {
        //debug_log("SYS_CLOSE: current->file_descriptors[fd] != NULL");
        int result = 0;//fat_fclose(current->file_descriptors[fd]);
        //debug_log("SYS_CLOSE: fat_fclose result: 0x%X", result);
        free(current->file_descriptors[fd]);
        current->file_descriptors[fd] = NULL;
        return result;
    }
    return -1;
}

#include "../include/exec_from_file.h"

extern void execve_return(TrapFrame *tf);

/**
 * this might be a not really good implementation as it only supports
 * up to EXECVE_MAX_ARGUMENT_SIZE characters in each argument or environment
 * string and for some reason I was unable to use malloc for kernel_argv and
 * kernel_envp
*/
int sys_execve(TrapFrame *tf) {
    const char* filename = (const char*)tf->ebx;

    char **argv = (char**)tf->ecx;
    char **envp = (char**)tf->edx;

    int argc = 0;
    int envc = 0;

    while(argv[argc] != NULL)
    {
        argc++;
    }
    while(envp[envc] != NULL)
    {
        envc++;
    }

    //TODO: check why was there a triple fault when I tried using malloc
    char kernel_argv[argc][EXECVE_MAX_ARGUMENT_SIZE];
    char kernel_envp[envc][EXECVE_MAX_ARGUMENT_SIZE];

    for(int i = 0; i < argc; i++)
    {
        strcpy(kernel_argv[i], argv[i]);
    }

    for(int i = 0; i < envc; i++)
    {
        strcpy(kernel_envp[i], envp[i]);
    }

    uint8_t* file_buffer;
    int err = vfs_read_file_by_path(filename, &file_buffer);
    if (err != 0) {
        debug_log("execve: failed to read file '%s'\n", filename);
        return -1;
    }

    Elf32_Ehdr *elf_header = (Elf32_Ehdr*)file_buffer;
    if (memcmp(elf_header->e_ident, "\x7F" "ELF", 4) != 0) {
        debug_log("execve: '%s' is not a valid ELF file\n", filename);
        free(file_buffer);
        return -1;
    }

    uint32_t* new_page_dir = create_page_dir();
    if (!new_page_dir) {
        free(file_buffer);
        return -1;
    }
    uint32_t new_page_dir_phys = (uint32_t)get_physaddr((void*)new_page_dir);
    uint32_t* old_page_dir_to_free = current->pagedir;

    set_page_dir(new_page_dir_phys);
    current->pagedir = new_page_dir;

    void* kstack_phys = get_physaddr(current->kstack);
    map_page(kstack_phys, current->kstack, PAGE_PRESENT | PAGE_RW);

    uint32_t highest_vaddr = 0;
    for (int i = 0; i < elf_header->e_phnum; i++) {
        Elf32_Phdr *p_header = (Elf32_Phdr *)(file_buffer + elf_header->e_phoff + i * elf_header->e_phentsize);
        if (p_header->p_type == PT_LOAD) {
            uint32_t start_addr = p_header->p_vaddr;
            uint32_t end_addr = start_addr + p_header->p_memsz;

            uint32_t start_page = start_addr & ~0xFFF;
            uint32_t end_page = (end_addr - 1) & ~0xFFF;

            for (uint32_t vaddr = start_page; vaddr <= end_page; vaddr += PAGE_SIZE) {
                void* phys_addr = alloc_page();
                map_page(phys_addr, (void*)vaddr, PAGE_PRESENT | PAGE_RW | PAGE_USER);
            }

            memcpy((void*)p_header->p_vaddr, file_buffer + p_header->p_offset, p_header->p_filesz);
            if (p_header->p_memsz > p_header->p_filesz) {
                memset((void*)(p_header->p_vaddr + p_header->p_filesz), 0, p_header->p_memsz - p_header->p_filesz);
            }
            if (p_header->p_vaddr + p_header->p_memsz > highest_vaddr) {
                highest_vaddr = p_header->p_vaddr + p_header->p_memsz;
            }
        }
    }

    for (uint32_t vaddr = USER_STACK_BOTTOM; vaddr <= USER_STACK_TOP; vaddr += PAGE_SIZE) {
        void* phys_addr = alloc_page();
        map_page(phys_addr, (void*)(vaddr & ~0xFFF), PAGE_PRESENT | PAGE_RW | PAGE_USER);
    }

    uint32_t user_esp = USER_STACK_TOP;
    
    //push environment strings
    uint32_t envp_pointers[envc + 1];
    for (int i = envc - 1; i >= 0; i--) {
        size_t len = strlen(kernel_envp[i]) + 1;
        user_esp -= len;
        memcpy((void*)user_esp, kernel_envp[i], len);
        envp_pointers[i] = user_esp;
    }
    envp_pointers[envc] = 0;//null terminator
    
    //push argument strings
    uint32_t argv_pointers[argc + 1];
    for (int i = argc - 1; i >= 0; i--) {
        size_t len = strlen(kernel_argv[i]) + 1;
        user_esp -= len;
        memcpy((void*)user_esp, kernel_argv[i], len);
        argv_pointers[i] = user_esp;
    }
    argv_pointers[argc] = 0;//null terminator

    //push envp pointers
    user_esp -= (envc + 1) * sizeof(uint32_t);
    memcpy((void*)user_esp, envp_pointers, (envc + 1) * sizeof(uint32_t));
    
    //push argv pointers
    user_esp -= (argc + 1) * sizeof(uint32_t);
    memcpy((void*)user_esp, argv_pointers, (argc + 1) * sizeof(uint32_t));
    
    //push argc
    user_esp -= sizeof(uint32_t);
    *((uint32_t*)user_esp) = argc;

    current->brk = (void*)(DivRoundUp(highest_vaddr, PAGE_SIZE) * PAGE_SIZE);
    free(file_buffer);
    destroy_page_dir(old_page_dir_to_free);

    tf->gs = SEG_UDATA | DPL_USER;
    tf->fs = SEG_UDATA | DPL_USER;
    tf->es = SEG_UDATA | DPL_USER;
    tf->ds = SEG_UDATA | DPL_USER;
    tf->eax = 0; tf->ecx = 0; tf->edx = 0; tf->ebx = 0;
    tf->ebp = 0; tf->esi = 0; tf->edi = 0;
    tf->eip = elf_header->e_entry;
    tf->cs = SEG_UCODE | DPL_USER;
    tf->eflags = FL_IF;
    tf->usermode_esp = user_esp;
    tf->usermode_ss = SEG_UDATA | DPL_USER;

    execve_return(tf);

    return 0;
}

int sys_chdir(TrapFrame *tf)
{
    l9660_dir new_cwd;
    l9660_file placeholder_file;

    int status = follow_path((char*)tf->ebx, &placeholder_file, &new_cwd);

    if(status == L9660_OK)
    {
        *(current->cwd) = new_cwd;
    }

    switch(status)
    {
        case L9660_EIO:
            return -5;//-EIO
            break;
        case L9660_EBADFS:
            return -5;//EIO as well
            break;
        case L9660_ENOENT:
            return -2;
            break;
        case L9660_ENOTDIR:
            return -20;//ENOTDIR
            break;
        default:
            return -5;//EIO
            break;
    }

    //strcat(current->cwd->fat->path, "/");
    //strcat(current->cwd->fat->path, (char*)tf->ebx);

    int result = 0;//fat_opendir(current->cwd, (const char*)tf->ebx);
    if(result == 0)
        return 0;

    return -1;
}

#include "../include/vfs.h"
/**
 * I'm very sorry before myself for this but these directory traversing things are so hard, follow_path in vfs was pretty okay,
 * this is when shit got real though
 */
uint32_t sys_getcwd(TrapFrame *tf)
{
    l9660_dir current_dir = *(current->cwd);

    char path_reversed[256] = {0};
    char component_name[64];

    while (!is_root_dir(&current_dir)) {
        uint32_t child_sector = current_dir.file.first_sector;
        
        // Open the parent directory
        l9660_dir parent_dir;
        l9660_opendirat(&parent_dir, &current_dir, "..");

        // Find our previous directory's name within the parent
        find_name_for_sector(component_name, sizeof(component_name), &parent_dir, child_sector);

        // Prepend the name to our reversed path string (e.g., "dash/" + "katau/" -> "katau/dash/")
        strcat(path_reversed, component_name);
        strcat(path_reversed, "/");

        // Move up one level
        current_dir = parent_dir;
    }

    // Now, reverse the string "dash/katau/" to create the final path "/katau/dash"
    char final_path[256] = "/";
    char* token = strtok(path_reversed, "/");
    while (token != NULL) {
        // A bit of a trick to prepend tokens
        char temp[256];
        strcpy(temp, "/");
        strcat(temp, token);
        strcat(temp, final_path);
        strcpy(final_path, temp);
        
        token = strtok(NULL, "/");
    }
    // Handle the root case where the loop doesn't run
    if (strlen(final_path) > 1) {
        final_path[strlen(final_path) - 1] = '\0'; // Remove trailing slash
    }

    strcpy((char*)tf->ebx, final_path);

    //strcpy((char*)tf->ebx, current->cwd->fat->path);
    
    return strlen(final_path);//tf->ebx;
}

int sys_getpid(TrapFrame *tf)
{
    return current->pid;
}

#define DivRoundUp(number, divisor) ((number + divisor - 1) / divisor)

uint32_t sys_brk(TrapFrame *tf)
{
    if(!tf->ebx)
    {
        //debug_log("current brk is %X\n", (uint32_t)current->brk);
        return (uint32_t)current->brk;
    }
    if(tf->ebx < (uint32_t)current->brk)
    {
        debug_log("PIZDETS!!!! new brk less than current brk\n");
        while(1){}
    }

    uint32_t old_brk_page = DivRoundUp((uint32_t)current->brk, 0x1000);
    uint32_t new_brk_page = DivRoundUp(tf->ebx, 0x1000);

    uint32_t pages_needed = new_brk_page - old_brk_page;
    for(uint32_t i = 0; i < pages_needed; i++)
    {
        void* phys_addr = alloc_page();
        void* virt_addr = (void*)(old_brk_page * 0x1000 + i * 0x1000);
        map_page(phys_addr, virt_addr, 0x7);
        //debug_log("mapping page %X to %X\n", phys_addr, virt_addr);
        memset(virt_addr, 0, 0x1000);
        current->brk = (void*)tf->ebx;
        //debug_log("new brk is %X\n", (uint32_t)current->brk);
    }
    return (uint32_t)current->brk;
}

typedef struct timespec {
    int32_t tv_sec;  // seconds
    int32_t tv_nsec; // nanoseconds
} timespec;

uint32_t sys_nanosleep(TrapFrame *tf)
{
    timespec* duration = (timespec*)tf->ebx;
    size_t ms = duration->tv_sec * 1000 + duration->tv_nsec / 1000000;
    //debug_log("waiting for 0x%X ms\n", ms);

    scheduler_lock();
    current->wake_up_time = timer_ticks + ms;
    current->state = Waiting;
    current->waiting_reason = TIMER;
    schedule();
    scheduler_unlock();
    
    return 0;
}

void handle_syscall(TrapFrame *tf)
{
    /*
    debug_log("EAX: %X ", tf->eax);
    debug_log("EBX: %X ", tf->ebx);
    debug_log("ECX: %s ", tf->ecx);
    debug_log("EDX: %X\n", tf->edx);
    */

    //debug_log("calling 0x%X\n", tf->eax);

    switch(tf->eax)
    {
        case 1://exit
            tf->eax = sys_exit(tf);
            break;

        case 2://fork
            tf->eax = sys_fork(tf);
            break;
    
        case 3://read
            tf->eax = sys_read(tf);
            break;

        case 4://write
            tf->eax = sys_write(tf);
            break;
        case 5://open
            tf->eax = sys_open(tf);
            break;
        case 6://close
            tf->eax = sys_close(tf);
            break;

        case 11://execve
            tf->eax = sys_execve(tf);
            break;
        
        case 12://chdir
            tf->eax = sys_chdir(tf);
            break;
        
        case 20://getpid
            tf->eax = sys_getpid(tf);
            break;
        
        case 183://getcwd
            tf->eax = sys_getcwd(tf);
            break;
        
        case 0x2d://brk
            tf->eax = sys_brk(tf);
            break;

        case 0x36://ioctl
            debug_log("IOCTL WAS CALLED, FUCKFUCK!!!!\n");
            tf->eax = -38;//-ENOSYS
            break;
        
        case 0xa2://nanosleep
            tf->eax = sys_nanosleep(tf);
            break;

        default:
            debug_log("unknown syscall: %X\n", tf->eax);
            tf->eax = -38;//-ENOSYS
            break;
    }
}