#include <stdint.h>
#include <stddef.h>

#include "lib/memory.h"
#include "lib/ringbuf.h"

#include "memory/mm.h"

#include "process.h"
#include "scheduler.h"
#include "vfs.h"
#include "elf.h"

#include <drivers/vga/vga.h>
#include <drivers/pty/pty.h>

#include <lib/print.h>
#include <lib/string.h>

typedef struct registers {
    uint32_t ds;                                       // push %ds
    uint32_t edi, esi, ebp, esp, ebx, edx, ecx, eax;   // pushal
    uint32_t eip, cs, eflags, useresp, ss;             // pushed by cpu
} __attribute__((packed)) registers_t;

static int32_t sys_exit(int status) {
    process_t* task = current_task();
    task->state = STATE_DEAD;
    kprintf("exiting task %d with status %d\n", task->pid, status);
    exit();
    return 0;
}

static int32_t sys_write(int fd, const void* buf, size_t count) {
    if (fd < 0 || fd >= MAX_PROCESS_FDS) return -1;

    file_t* f = current_task()->fd_table[fd];
    if (!f || !f->node) return -1;

    uint32_t rd = f->node->write(f->node, f->offset, count, (uint8_t*)buf);
    f->offset += rd;
    
    return rd; 
}

static int32_t sys_read(int fd, void* buf, size_t count) {
    if (fd < 0 || fd >= MAX_PROCESS_FDS) return -1;

    file_t* f = current_task()->fd_table[fd];
    if (!f || !f->node) return -1;

    uint32_t rd = f->node->read(f->node, f->offset, count, (uint8_t*)buf);
    f->offset += rd;

    return rd; 
}

static int32_t sys_open(const char* filename, int flags) {
    return vfs_open(filename, flags);
}

static int32_t sys_close(int fd) {
    process_t* task = current_task();
    task->fd_table[fd] = NULL;
    return 1;
}

static int32_t sys_brk(uint32_t new_break) {
    process_t* current = current_task();

    if (new_break == 0) {
        return current->heap_end;
    }

    if (new_break < current->heap_end) {
        current->heap_end = new_break;
        return current->heap_end;
    }

    uint32_t page_start = (current->heap_end + 4095) & ~4095;
    uint32_t page_end   = (new_break + 4095) & ~4095;

    for (uint32_t addr = page_start; addr < page_end; addr += 4096) {
        // 1. Allocate a physical frame using your PMM (Physical Memory Manager)
        void* phys = p_alloc_frame();
        
        // 2. Map it into the current page directory using your VMM (Virtual Memory Manager)
        map_page(addr, (uint32_t)phys, PAGE_PRESENT | PAGE_WRITABLE | PAGE_USER);
    }

    current->heap_end = new_break;
    return current->heap_end;
}

static int32_t sys_pty(int* slave_fd) {
    pty_t* pty = (pty_t*)kalloc(sizeof(pty_t));
    ring_buf_init(&pty->master_rb);
    ring_buf_init(&pty->slave_rb);

    vfs_node_t* master =(vfs_node_t*)kalloc(sizeof(vfs_node_t));
    master->flags = VFS_CHARDEVICE;
    master->size = 0;
    master->read = pty_master_read;
    master->write = pty_master_write;
    master->find = NULL;
    master->data = pty;

    vfs_node_t* slave =(vfs_node_t*)kalloc(sizeof(vfs_node_t));
    slave->flags = VFS_CHARDEVICE;
    slave->size = 0;
    slave->read = pty_slave_read;
    slave->write = pty_slave_write;
    slave->find = NULL;
    slave->data = pty;

    file_t* fmaster = (file_t*)kalloc(sizeof(file_t));
    fmaster->node = master;
    fmaster->offset = 0;
    fmaster->flags = 0;

    file_t* fslave = (file_t*)kalloc(sizeof(file_t));
    fslave->node = slave;
    fslave->offset = 0;
    fslave->flags = 0;

    *slave_fd = alloc_fd(fslave);
    return alloc_fd(fmaster);
}

static int32_t sys_fork(registers_t* parent_regs) {
    process_t* parent = current_task();
    process_t* child = kalloc(sizeof(process_t));
    if (!child) {
        return -1; 
    }
    
    child->pid = get_next_pid();
    child->state = STATE_READY;
    child->flags = parent->flags;
    child->heap_end = parent->heap_end;
    child->sleep = 0;
    
    for (int i = 0; i < MAX_PROCESS_FDS; i++) {
        child->fd_table[i] = parent->fd_table[i];
    }

    uint32_t kstack_size = 4096;
    void* kstack_alloc = kalloc(kstack_size);
    child->kstack_top = (uint32_t)kstack_alloc + kstack_size;

    registers_t* child_regs = (registers_t*)(child->kstack_top - sizeof(registers_t));
    memcpy(parent_regs, child_regs, sizeof(registers_t));
    child_regs->eax = 0;
    child->esp = (uint32_t)child_regs;

    uint32_t cr3 = (uint32_t)clone_current_pd();
    child->cr3 = cr3;

    add_task(child);
    return child->pid;
}

static int32_t sys_exec(registers_t* regs, const char* prgm) {
    clear_current_pd();

    uint32_t eip = load_elf_program(prgm);

    uint32_t stack_sz = PAGE_SIZE * 16;
    uint32_t stack_top = 0xBFFF0000;
    uint32_t stack_bottom = stack_top - stack_sz;
    for (uint32_t vaddr = stack_bottom; vaddr < stack_top; vaddr += PAGE_SIZE) {
        void* phys_frame = p_alloc_frame();
        map_page(vaddr, (uint32_t)phys_frame, PAGE_PRESENT | PAGE_WRITABLE | PAGE_USER);
    }

    uint32_t* u_esp = (uint32_t*)stack_top;
    *(--u_esp) = 0;          // envp = NULL
    *(--u_esp) = 0;          // argv = NULL
    *(--u_esp) = 0;          // argc = 0

    regs->eax = 0;
    regs->ecx = 0;
    regs->edx = 0;
    regs->ebp = 0;
    regs->esi = 0;
    regs->edi = 0;
    regs->eip = eip;
    regs->useresp = (uint32_t)u_esp;

    return 0;
}

static int32_t sys_dup2(int src, int dst) {
    process_t* task = current_task();
    task->fd_table[dst] = task->fd_table[src];
    return dst;
}

void syscall(registers_t* regs) {
    int32_t ret;
    switch(regs->eax) {
        case 1:  ret = sys_exit(regs->ebx); break;
        case 3:  ret = sys_read(regs->ebx, (void*)regs->ecx, regs->edx); break;
        case 4:  ret = sys_write(regs->ebx, (const void*)regs->ecx, regs->edx); break;
        case 5:  ret = sys_open((const char*)regs->ebx, regs->ecx); break;
        case 6:  ret = sys_close(regs->ebx); break;
        case 12: ret = sys_brk(regs->ebx); break;
        case 20: ret = sys_pty((int*)regs->ebx); break;
        case 21: ret = sys_fork(regs); break;
        case 22: ret = sys_exec(regs, (const char*)regs->ebx); break;
        case 23: ret = sys_dup2(regs->ebx, regs->ecx); break;
        default: ret = -1;
    }
    regs->eax = ret;
}