/*
 * process.c
 *
 * Copyright (C) 2017 Aleksandar Andrejevic <theflash@sdf.lonestar.org>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <process.h>
#include <thread.h>
#include <heap.h>
#include <syscalls.h>
#include <exception.h>
#include <exec/aout.h>
#include <cpu.h>

dword_t load_elf(handle_t file, thread_state_t *initial_state);
dword_t load_aout(handle_t file, thread_state_t *initial_state);

typedef dword_t (*loader_proc)(handle_t, thread_state_t*);

static dword_t pid_alloc_bitmap[MAX_PROCESSES / 32];
static DECLARE_LOCK(pid_bitmap_lock);

static loader_proc loaders[] =
{
    load_elf,
    load_aout,
    NULL
};

process_t *kernel_process;

static dword_t alloc_pid()
{
    int i;
    dword_t pid = (dword_t)-1;

    lock_acquire(&pid_bitmap_lock);

    for (i = 0; i < MAX_PROCESSES; i++)
    {
        if (!test_bit(pid_alloc_bitmap, i))
        {
            pid = i;
            set_bit(pid_alloc_bitmap, pid);
            break;
        }
    }

    lock_release(&pid_bitmap_lock);
    return pid;
}

void destroy_process(process_t *proc)
{
    int i;
    proc->terminating = TRUE;

    delete_address_space(&proc->memory_space);

    lock_acquire(&proc->handle_table_lock);

    for (i = 0; i < proc->handle_table_size / sizeof(proc->handle_table[i]); i++)
    {
        if (proc->handle_table[i].obj)
        {
            close_object_internal(proc->handle_table[i].obj);

            proc->handle_table[i].obj = NULL;
            proc->handle_count--;
        }
    }

    lock_release(&proc->handle_table_lock);

    proc->terminated = TRUE;
    syscall_clock_get_time(&proc->end_time);

    dereference(&proc->header);
}

void process_cleanup(object_t *obj)
{
    process_t *proc = (process_t*)obj;

    free(proc->name);
    dereference(&proc->current_user->header);

    lock_acquire(&pid_bitmap_lock);
    clear_bit(pid_alloc_bitmap, proc->pid);
    lock_release(&pid_bitmap_lock);
}

dword_t process_pre_wait(object_t *obj, void *parameter, wait_condition_t *condition)
{
    process_t *proc = (process_t*)obj;
    UNUSED_PARAMETER(parameter);
    condition->type = WAIT_UNTIL_NOT_EQUAL;
    condition->pointer = &proc->terminated;
    condition->value = 0;
    return ERR_SUCCESS;
}

process_t *get_current_process()
{
    thread_t *thread = get_current_thread();
    if (thread == NULL) return NULL;
    return thread->owner_process;
}

sysret_t syscall_open_process(dword_t pid, handle_t *handle)
{
    dword_t ret;
    process_t *proc = NULL;

    ret = enum_objects_by_type(OBJECT_PROCESS, (object_t**)&proc);
    ASSERT(ret == ERR_SUCCESS || ret == ERR_NOMORE);

    while (ret == ERR_SUCCESS)
    {
        if (proc->pid == pid) break;
        ret = enum_objects_by_type(OBJECT_PROCESS, (object_t**)&proc);
    }

    ASSERT(ret == ERR_SUCCESS || ret == ERR_NOMORE);
    if (ret == ERR_NOMORE) ret = ERR_NOTFOUND;

    if (ret == ERR_SUCCESS)
    {
        if (proc->current_user->uid == get_current_uid() || check_privileges(PRIVILEGE_PROCESS_CONTROL))
        {
            ret = open_object(&proc->header, 0, handle);
        }
        else
        {
            ret = ERR_FORBIDDEN;
        }

        dereference(&proc->header);
    }

    return ret;
}

dword_t load_executable(handle_t file, thread_state_t *initial_state)
{
    loader_proc *loader = loaders;

    while (*loader)
    {
        dword_t ret = (*loader++)(file, initial_state);
        if (ret == ERR_INVALID) continue;
        return ret;
    }

    return ERR_INVALID;
}

void init_user_stack(uintptr_t *stack_pointer, process_params_t *parameters)
{
    uintptr_t stack_top = *stack_pointer;

    if (parameters->command_line)
    {
        *stack_pointer -= (strlen(parameters->command_line) + 1 + sizeof(uintptr_t) - 1) & ~(sizeof(uintptr_t) - 1);
        strcpy((char*)*stack_pointer, parameters->command_line);
        parameters->command_line = (char*)*stack_pointer;
    }

    *stack_pointer -= (sizeof(process_params_t) + sizeof(uintptr_t) - 1) & ~(sizeof(uintptr_t) - 1);
    *(process_params_t*)*stack_pointer = *parameters;
    parameters = (process_params_t*)*stack_pointer;

    if ((stack_top - *stack_pointer) < MAX_PARAMETERS * sizeof(uintptr_t))
    {
        *stack_pointer = stack_top - MAX_PARAMETERS * sizeof(uintptr_t);
    }

    init_thread_stack(stack_pointer, parameters);
}

sysret_t syscall_create_process(const char *path, dword_t flags, process_params_t *parameters, handle_t *process_handle, handle_t *thread_handle)
{
    dword_t ret;
    handle_t file = INVALID_HANDLE;
    process_t *current = get_current_process();
    char *safe_path = NULL;
    process_params_t safe_params;
    handle_t safe_process_handle = INVALID_HANDLE;
    handle_t safe_thread_handle = INVALID_HANDLE;
    bool_t object_created = FALSE;

    if (get_previous_mode() == USER_MODE)
    {
        if (!check_usermode(process_handle, sizeof(handle_t))) return ERR_BADPTR;
        if (!check_usermode(thread_handle, sizeof(handle_t))) return ERR_BADPTR;

        if (path)
        {
            safe_path = copy_user_string(path);
            if (safe_path == NULL) return ERR_BADPTR;
        }

        if (parameters)
        {
            if (!check_usermode(parameters, sizeof(process_params_t)))
            {
                free(safe_path);
                return ERR_BADPTR;
            }

            EH_TRY safe_params = *parameters;
            EH_CATCH EH_ESCAPE(return ERR_BADPTR);
            EH_DONE;

            char *cmdline = copy_user_string(safe_params.command_line);
            if (cmdline == NULL)
            {
                free(safe_path);
                return ERR_BADPTR;
            }

            safe_params.command_line = cmdline;
            parameters = &safe_params;
        }
    }
    else
    {
        safe_path = (char*)path;

        if (parameters)
        {
            safe_params = *parameters;
            parameters = &safe_params;
        }
    }

    if (safe_path)
    {
        ret = syscall(SYSCALL_OPEN_FILE, safe_path, &file, FILE_MODE_READ, 0);
        if (ret != ERR_SUCCESS) return ret;
    }

    process_t *proc = (process_t*)malloc(sizeof(process_t));
    if (proc == NULL) return ERR_NOMEMORY;
    memset(proc, 0, sizeof(process_t));

    init_object(&proc->header, NULL, OBJECT_PROCESS);

    proc->pid = alloc_pid();
    if (proc->pid == (dword_t)-1)
    {
        ret = ERR_NOMEMORY;
        goto cleanup;
    }

    proc->name = strdup(path ? path : get_current_process()->name);

    list_init(&proc->threads);
    lock_init(&proc->thread_list_lock);

    reference(&current->current_user->header);
    proc->original_user = proc->current_user = current->current_user;

    if (!(flags & PROCESS_CREATE_NO_INHERIT) && current != kernel_process)
    {
        dword_t i;
        lock_acquire_shared(&current->handle_table_lock);

        proc->handle_table = (handle_info_t*)heap_alloc(&evictable_heap, STARTUP_HANDLE_TABLE_SIZE);
        memset(proc->handle_table, 0, current->handle_table_size);
        proc->handle_table_size = current->handle_table_size;
        proc->handle_count = current->handle_count;
        lock_init(&proc->handle_table_lock);

        for (i = 0; i < current->handle_table_size; i++)
        {
            if (current->handle_table[i].obj)
            {
                reference(current->handle_table[i].obj);
                current->handle_table[i].obj->open_count++;
                proc->handle_table[i] = current->handle_table[i];
            }
        }

        lock_release(&current->handle_table_lock);
    }
    else
    {
        proc->handle_table = (handle_info_t*)heap_alloc(&evictable_heap, STARTUP_HANDLE_TABLE_SIZE);
        memset(proc->handle_table, 0, STARTUP_HANDLE_TABLE_SIZE);
        proc->handle_table_size = STARTUP_HANDLE_TABLE_SIZE;
        proc->handle_count = 0;
        lock_init(&proc->handle_table_lock);
    }

    if (safe_path)
    {
        ret = create_address_space(PROCESS_POOL_ADDRESS, PROCESS_POOL_SIZE, &proc->memory_space);
        if (ret != ERR_SUCCESS) goto cleanup;
    }
    else
    {
        ret = clone_address_space(&current->memory_space, &proc->memory_space);
        if (ret != ERR_SUCCESS) goto cleanup;
    }

    syscall_clock_get_time(&proc->start_time);

    ret = create_object(&proc->header);
    if (ret != ERR_SUCCESS) goto cleanup;

    object_created = TRUE;

    thread_state_t initial_state;
    memset(&initial_state, 0, sizeof(initial_state));

    ret = open_object(&proc->header, 0, &safe_process_handle);
    if (ret != ERR_SUCCESS) goto cleanup;

    if (file != INVALID_HANDLE)
    {
        process_t *old_process = switch_process(proc);
        ret = load_executable(file, &initial_state);
        init_user_stack((uintptr_t*)&initial_state.regs.esp, parameters);
        switch_process(old_process);
        if (ret != ERR_SUCCESS) goto cleanup;
    }
    else
    {
        initial_state.regs = *get_current_thread()->last_context;
        cpu_save_fpu_state(initial_state.fpu_state);
        initial_state.regs.eax = CLONE_MAGIC;
    }

    if (!(flags & PROCESS_CREATE_NO_THREADS))
    {
        thread_t *thread;

        void *kernel_stack = malloc(KERNEL_STACK_SIZE + sizeof(uintptr_t) - 1);
        if (kernel_stack == NULL)
        {
            ret = ERR_NOMEMORY;
            goto cleanup;
        }

        ret = commit_pages(kernel_stack, KERNEL_STACK_SIZE);
        if (ret != ERR_SUCCESS) goto cleanup;

        ret = create_thread_internal(proc,
                                     &initial_state,
                                     (flags & PROCESS_CREATE_FROZEN_THREAD) ? THREAD_CREATE_FROZEN : 0,
                                     THREAD_PRIORITY_MID,
                                     kernel_stack,
                                     &thread);
        if (ret != ERR_SUCCESS)
        {
            free(kernel_stack);
            goto cleanup;
        }

        ret = open_object(&thread->header, 0, &safe_thread_handle);
        if (ret != ERR_SUCCESS)
        {
            dereference(&thread->header);
            goto cleanup;
        }
    }

    EH_TRY
    {
        *process_handle = safe_process_handle;
        *thread_handle = safe_thread_handle;
    }
    EH_CATCH
    {
        ret = ERR_BADPTR;
    }
    EH_DONE;

cleanup:
    if (file != INVALID_HANDLE) syscall(SYSCALL_CLOSE_OBJECT, file);

    if (object_created)
    {
        if (ret != ERR_SUCCESS)
        {
            syscall(SYSCALL_TERMINATE, proc->pid, 1);
            if (safe_process_handle != INVALID_HANDLE) syscall_close_object(safe_process_handle);
            if (safe_thread_handle != INVALID_HANDLE) syscall_close_object(safe_thread_handle);
            dereference(&proc->header);
        }
    }
    else
    {
        ASSERT(ret != ERR_SUCCESS);

        if (proc->pid != (dword_t)-1)
        {
            lock_acquire(&pid_bitmap_lock);
            clear_bit(pid_alloc_bitmap, proc->pid);
            lock_release(&pid_bitmap_lock);
        }

        if (proc->name) free(proc->name);
        free(proc);
    }

    if (get_previous_mode() == USER_MODE)
    {
        if (path != NULL) free(safe_path);
        free(safe_params.command_line);
    }

    return ret;
}

void terminate_process(process_t *proc, dword_t exit_code)
{
    proc->terminating = TRUE;
    proc->exit_code = exit_code;

    lock_acquire_shared(&proc->thread_list_lock);

    list_entry_t *ptr;
    for (ptr = proc->threads.next; ptr != &proc->threads; ptr = ptr->next)
    {
        thread_t *thread = CONTAINER_OF(ptr, thread_t, in_process_list);;
        terminate_thread_internal(thread, exit_code);
    }

    lock_release(&proc->thread_list_lock);
}

sysret_t syscall_terminate(handle_t handle, dword_t exit_code)
{
    process_t *proc;

    if (handle == INVALID_HANDLE)
    {
        proc = get_current_process();
        reference(&proc->header);
    }
    else
    {
        if (!reference_by_handle(handle, OBJECT_PROCESS, (object_t**)&proc)) return ERR_INVALID;
    }

    terminate_process(proc, exit_code);
    dereference(&proc->header);
    return ERR_SUCCESS;
}

sysret_t syscall_query_process(handle_t handle, process_info_t info_type, void *buffer, dword_t size)
{
    dword_t ret = ERR_SUCCESS;
    void *safe_buffer;

    if (get_previous_mode() == USER_MODE)
    {
        if (!check_usermode(buffer, size)) return ERR_BADPTR;

        safe_buffer = malloc(size);
        if (safe_buffer == NULL) return ERR_NOMEMORY;
        memset(safe_buffer, 0, size);
    }
    else
    {
        safe_buffer = buffer;
    }

    process_t *proc;

    if (handle == INVALID_HANDLE)
    {
        proc = get_current_process();
        reference(&proc->header);
    }
    else
    {
        if (!reference_by_handle(handle, OBJECT_PROCESS, (object_t**)&proc))
        {
            if (get_previous_mode() == USER_MODE) free(safe_buffer);
            return ERR_INVALID;
        }
    }

    switch (info_type)
    {
        case PROCESS_PID_INFO:
        {
            if (size >= sizeof(dword_t)) *((dword_t*)safe_buffer) = proc->pid;
            else ret = ERR_SMALLBUF;

            break;
        }

        case PROCESS_NAME_INFO:
        {
            if (size >= strlen(proc->name) + 1) strcpy(safe_buffer, proc->name);
            else ret = ERR_SMALLBUF;

            break;
        }

        case PROCESS_MEMORY_INFO:
        {
            ret = ERR_NOSYSCALL; // TODO
            break;
        }

        case PROCESS_START_TIME_INFO:
        {
            if (size >= sizeof(clock_time_t)) *((clock_time_t*)safe_buffer) = proc->start_time;
            else ret = ERR_SMALLBUF;

            break;
        }

        case PROCESS_EXIT_CODE_INFO:
        {
            if (size >= sizeof(dword_t)) *((dword_t*)safe_buffer) = proc->exit_code;
            else ret = ERR_SMALLBUF;

            break;
        }

        case PROCESS_USER_INFO:
        {
            if (size >= 2 * sizeof(dword_t))
            {
                ((dword_t*)safe_buffer)[0] = proc->original_user->uid;
                ((dword_t*)safe_buffer)[1] = proc->current_user->uid;
            }
            else
            {
                ret = ERR_SMALLBUF;
            }

            break;
        }

        case PROCESS_THREAD_INFO:
        {
            if (size >= sizeof(dword_t))
            {
                list_entry_t *ptr;
                dword_t i = 1, count = 0;

                lock_acquire_shared(&proc->thread_list_lock);
                for (ptr = proc->threads.next; ptr != &proc->threads; ptr = ptr->next) count++;
                *((dword_t*)safe_buffer) = count;

                for (ptr = proc->threads.next; ptr != &proc->threads && i < (size / sizeof(dword_t)); ptr = ptr->next)
                {
                    ((dword_t*)safe_buffer)[i++] = CONTAINER_OF(ptr, thread_t, in_process_list)->tid;
                }

                if (ptr != &proc->threads) ret = ERR_SMALLBUF;
                lock_release(&proc->thread_list_lock);
            }
            else
            {
                ret = ERR_SMALLBUF;
            }

            break;
        }

        case PROCESS_HANDLE_INFO:
        {
            if (size >= sizeof(dword_t))
            {
                dword_t i, j = 1;

                lock_acquire_shared(&proc->handle_table_lock);
                *((dword_t*)safe_buffer) = proc->handle_count;

                for (i = 0; i < proc->handle_table_size && j < (size / sizeof(dword_t)); i++)
                {
                    if (proc->handle_table[i].obj) ((dword_t*)safe_buffer)[j++] = i;
                }

                if (i != proc->handle_table_size) ret = ERR_SMALLBUF;
                lock_release(&proc->handle_table_lock);
            }
            else
            {
                ret = ERR_SMALLBUF;
            }

            break;
        }

        default:
        {
            ret = ERR_INVALID;
        }
    }

    if (get_previous_mode() == USER_MODE)
    {
        EH_TRY memcpy(buffer, safe_buffer, size);
        EH_CATCH ret = ERR_BADPTR;
        EH_DONE;

        free(safe_buffer);
    }

    dereference(&proc->header);
    return ret;
}

sysret_t syscall_enum_processes(dword_t *pid_array, dword_t *count)
{
    dword_t ret = ERR_SUCCESS;
    dword_t safe_count;
    dword_t cnt = 0;

    if (get_previous_mode() == USER_MODE)
    {
        if (!check_usermode(count, sizeof(*count))) return ERR_BADPTR;

        EH_TRY safe_count = *count;
        EH_CATCH EH_ESCAPE(return ERR_BADPTR);
        EH_DONE;

        if (!check_usermode(pid_array, sizeof(*pid_array) * safe_count))
        {
            ret = ERR_BADPTR;
            goto cleanup;
        }
    }
    else
    {
        safe_count = *count;
    }

    process_t *proc = NULL;
    ret = enum_objects_by_type(OBJECT_PROCESS, (object_t**)&proc);

    while (ret == ERR_SUCCESS)
    {
        if (cnt == safe_count)
        {
            ret = ERR_SMALLBUF;
            break;
        }

        EH_TRY
        {
            pid_array[cnt++] = proc->pid;
        }
        EH_CATCH
        {
            ret = ERR_BADPTR;
            EH_ESCAPE(break);
        }
        EH_DONE;

        ret = enum_objects_by_type(OBJECT_PROCESS, (object_t**)&proc);
    }

    if (ret == ERR_NOMORE)
    {
        ret = ERR_SUCCESS;
    }
    else if (ret == ERR_SUCCESS)
    {
        dereference(&proc->header);
        ret = ERR_SMALLBUF;
    }

cleanup:
    EH_TRY *count = cnt;
    EH_CATCH ret = ERR_BADPTR;
    EH_DONE;

    return ret;
}

process_t *switch_process(process_t *new_process)
{
    thread_t *thread = get_current_thread();
    process_t *old_process = get_current_process();

    while (!__sync_bool_compare_and_swap(&thread->owner_process, old_process, new_process))
    {
        old_process = get_current_process();
    }

    set_page_directory(new_process->memory_space.page_directory);
    return old_process;
}

void process_init(void)
{
    memset(pid_alloc_bitmap, 0, sizeof(pid_alloc_bitmap));

    kernel_process = (process_t*)malloc(sizeof(process_t));
    ASSERT(kernel_process != NULL);

    memset(kernel_process, 0, sizeof(process_t));
    init_object(&kernel_process->header, NULL, OBJECT_PROCESS);

    ASSERT(create_object(&kernel_process->header) == ERR_SUCCESS);

    kernel_process->pid = 0;
    set_bit(pid_alloc_bitmap, 0);
    kernel_process->name = strdup("system");

    memset(&kernel_process->memory_space, 0, sizeof(memory_address_space_t));
    kernel_process->memory_space.page_directory = get_page_directory();

    kernel_process->exit_code = 0;
    kernel_process->terminating = FALSE;
    kernel_process->terminated = FALSE;

    syscall_clock_get_time(&kernel_process->start_time);

    list_init(&kernel_process->threads);
    lock_init(&kernel_process->thread_list_lock);

    kernel_process->handle_table = (handle_info_t*)heap_alloc(&evictable_heap, STARTUP_HANDLE_TABLE_SIZE);
    kernel_process->handle_table_size = STARTUP_HANDLE_TABLE_SIZE;
    kernel_process->handle_count = 0;
    lock_init(&kernel_process->handle_table_lock);
}
