// SPDX-License-Identifier: ISC

#include <zinnia/alloc.h>
#include <zinnia/atomic.h>
#include <zinnia/exec.h>
#include <zinnia/hashmap.h>
#include <zinnia/percpu.h>
#include <zinnia/print.h>
#include <zinnia/process.h>
#include <zinnia/sched.h>
#include <zinnia/utils.h>
#include <zinnia/vma.h>
#include <zinnia/vmo.h>

static struct process kernel_process;
static atomic_bool kernel_process_ready;
static atomic_int next_pid = 1;

static void process_init_common(struct process* proc) {
    refcount_init(&proc->rc);
    spin_init(&proc->tasks_lock);
    tailq_init(&proc->tasks);
    mutex_init(&proc->lock);
    mutex_init(&proc->files_lock);
    hashmap_init(&proc->files);
    vec_init(&proc->identity.groups);
}

struct process* process_kernel() {
    if (!atomic_load(&kernel_process_ready, ATOMIC_ACQUIRE)) {
        process_init_common(&kernel_process);
        kernel_process.pid = 0;
        kernel_process.space = &kernel_vma;
        atomic_store(&kernel_process_ready, true, ATOMIC_RELEASE);
    }
    return &kernel_process;
}

static void close_files(struct process* proc) {
    mutex_lock(&proc->files_lock);
    size_t i = 0;
    while (i < proc->files.capacity) {
        struct fd_entry entry = proc->files.values[i];
        if (proc->files.hashes[i] == 0) {
            i++;
            continue;
        }
        hashmap_remove(&proc->files, proc->files.keys[i], hashmap_hash_default, hashmap_eq_default);
        vfs_file_unref(entry.file);
    }
    mutex_unlock(&proc->files_lock);
}

void process_ref(struct process* proc) {
    refcount_inc(&proc->rc);
}

void process_unref(struct process* proc) {
    if (proc == &kernel_process || !refcount_dec(&proc->rc))
        return;
    close_files(proc);
    vfs_path_drop(proc->root);
    vfs_path_drop(proc->cwd);
    if (proc->space)
        vma_unref(proc->space);
    vec_destroy(&proc->identity.groups);
    hashmap_destroy(&proc->files);
    mem_free(proc);
}

uint8_t process_priority(struct process* proc) {
    if (proc == &kernel_process)
        return SCHED_PRIO_DEFAULT;

    // Normalize the priority from "nice" ranges to our priorities.
    return 16 + (19 - proc->nice) * 15 / 39;
}

void process_attach_task(struct process* proc, struct task* task) {
    spin_lock(&proc->tasks_lock);
    tailq_insert_tail(&proc->tasks, task, process_link);
    proc->thread_count++;
    spin_unlock(&proc->tasks_lock);
}

void process_detach_task(struct process* proc, struct task* task) {
    spin_lock(&proc->tasks_lock);
    tailq_remove(&proc->tasks, task, process_link);
    spin_unlock(&proc->tasks_lock);
}

static void rlimit_init(struct process_rlimit* limit, rlim_t cur, rlim_t max) {
    atomic_init(&limit->cur, cur);
    limit->max = max;
}

static errno_t process_alloc(struct process** out) {
    struct process* proc = mem_alloc(sizeof(struct process), 0);
    if (!proc)
        return ENOMEM;
    process_init_common(proc);
    proc->pid = atomic_fetch_add(&next_pid, 1, ATOMIC_RELAXED);
    proc->root = vfs_get_root();
    proc->cwd = vfs_get_root();
    proc->umask = 022;
    rlimit_init(&proc->limit_nofile, PROCESS_DEFAULT_FDS, PROCESS_MAX_FDS);

    *out = proc;
    return 0;
}

static errno_t load_image(
    struct process* proc,
    struct vmo* object,
    const char* path,
    const char** argv,
    const char** envp,
    struct vma** out_space,
    uintptr_t* out_ip,
    uintptr_t* out_sp
) {
    struct vma* space;
    errno_t status = vma_new(&space);
    if (status)
        return status;

    struct exec_info info = {
        .space = space,
        .file_obj = object,
        .exec_path = path,
        .argv = argv,
        .envp = envp,
    };
    struct exec_format* format = exec_identify(object);
    status = format ? format->load(proc, &info, out_ip, out_sp) : ENOEXEC;
    if (status) {
        vma_unref(space);
        return status;
    }

    *out_space = space;
    return 0;
}

errno_t process_spawn_init(const char* path, const char** argv, const char** envp) {
    struct process* proc;
    errno_t status = process_alloc(&proc);
    if (status)
        return status;

    struct task_start_info* start = mem_alloc(sizeof(struct task_start_info), 0);
    if (!start) {
        process_unref(proc);
        return ENOMEM;
    }

    struct vmo* object;
    status = exec_open(proc, path, &object);
    if (!status) {
        status = load_image(proc, object, path, argv, envp, &proc->space, &start->ip, &start->sp);
        vmo_unref(object);
    }
    if (status) {
        mem_free(start);
        process_unref(proc);
        return status;
    }

    struct vfs_file* console;
    status = vfs_file_open(proc->root, proc->cwd, "/dev/console", O_RDWR, 0, ident_kernel(), &console);
    if (status) {
        log_warn("process: Unable to open /dev/console for init (%i)\n", status);
    } else {
        for (int fd = 0; fd <= 2; fd++)
            process_fd_replace(proc, fd, console, false);
        vfs_file_unref(console);
    }

    struct task* task;
    status = task_create(argv[0], proc, sched_to_user, (uintptr_t)start, &task);
    process_unref(proc);
    if (status) {
        mem_free(start);
        return status;
    }

    sched_add_task(&percpu_get()->sched, task);
    return 0;
}

errno_t process_fd_alloc(struct process* proc, struct vfs_file* file, bool cloexec, int* out_fd) {
    return process_fd_alloc_min(proc, file, 0, cloexec, out_fd);
}

static struct fd_entry* fd_lookup(struct process* proc, int fd) {
    return hashmap_get(&proc->files, fd, hashmap_hash_default, hashmap_eq_default);
}

size_t process_fd_limit(struct process* proc) {
    return MIN((size_t)PROCESS_MAX_FDS, (size_t)atomic_load(&proc->limit_nofile.cur, ATOMIC_RELAXED));
}

errno_t process_fd_alloc_min(struct process* proc, struct vfs_file* file, int min_fd, bool cloexec, int* out_fd) {
    if (min_fd < 0 || (size_t)min_fd >= process_fd_limit(proc))
        return EINVAL;

    mutex_lock(&proc->files_lock);
    int fd = min_fd;
    while ((size_t)fd < process_fd_limit(proc) && fd_lookup(proc, fd))
        fd++;

    if ((size_t)fd >= process_fd_limit(proc)) {
        mutex_unlock(&proc->files_lock);
        return EMFILE;
    }

    const struct fd_entry entry = {.file = file, .cloexec = cloexec};
    errno_t status = hashmap_insert(&proc->files, fd, entry, hashmap_hash_default, hashmap_eq_default);
    if (status) {
        mutex_unlock(&proc->files_lock);
        return status;
    }
    vfs_file_ref(file);
    mutex_unlock(&proc->files_lock);

    *out_fd = fd;
    return 0;
}

errno_t process_fd_replace(struct process* proc, int fd, struct vfs_file* file, bool cloexec) {
    if (fd < 0 || (size_t)fd >= process_fd_limit(proc))
        return EBADF;

    struct vfs_file* displaced = nullptr;
    const struct fd_entry entry = {.file = file, .cloexec = cloexec};
    mutex_lock(&proc->files_lock);
    struct fd_entry* slot = fd_lookup(proc, fd);
    if (slot) {
        displaced = slot->file;
        *slot = entry;
    } else {
        errno_t status = hashmap_insert(&proc->files, fd, entry, hashmap_hash_default, hashmap_eq_default);
        if (status) {
            mutex_unlock(&proc->files_lock);
            return status;
        }
    }
    vfs_file_ref(file);
    mutex_unlock(&proc->files_lock);

    if (displaced)
        vfs_file_unref(displaced);
    return 0;
}

errno_t process_fd_get_cloexec(struct process* proc, int fd, bool* out_cloexec) {
    errno_t status = EBADF;
    mutex_lock(&proc->files_lock);
    struct fd_entry* entry = fd_lookup(proc, fd);
    if (entry) {
        *out_cloexec = entry->cloexec;
        status = 0;
    }
    mutex_unlock(&proc->files_lock);
    return status;
}

errno_t process_fd_set_cloexec(struct process* proc, int fd, bool cloexec) {
    errno_t status = EBADF;
    mutex_lock(&proc->files_lock);
    struct fd_entry* entry = fd_lookup(proc, fd);
    if (entry) {
        entry->cloexec = cloexec;
        status = 0;
    }
    mutex_unlock(&proc->files_lock);
    return status;
}

struct vfs_file* process_fd_get(struct process* proc, int fd) {
    struct vfs_file* file = nullptr;
    mutex_lock(&proc->files_lock);
    struct fd_entry* entry = fd_lookup(proc, fd);
    if (entry) {
        file = entry->file;
        vfs_file_ref(file);
    }
    mutex_unlock(&proc->files_lock);
    return file;
}

errno_t process_fd_close(struct process* proc, int fd) {
    struct vfs_file* file = nullptr;
    mutex_lock(&proc->files_lock);
    struct fd_entry* entry = fd_lookup(proc, fd);
    if (entry) {
        file = entry->file;
        hashmap_remove(&proc->files, fd, hashmap_hash_default, hashmap_eq_default);
    }
    mutex_unlock(&proc->files_lock);

    if (!file)
        return EBADF;
    vfs_file_unref(file);
    return 0;
}