// SMP.1 for non-curthr actions; none for curthr
#include "config.h"
#include "globals.h"
#include "mm/slab.h"
#include "util/debug.h"
#include "util/string.h"

/*==========
 * Variables
 *=========*/

/*
 * Global variable maintaining the current thread on the cpu
 */
kthread_t *curthr CORE_SPECIFIC_DATA;

/*
 * Private slab for kthread structs
 */
static slab_allocator_t *kthread_allocator = NULL;

/*=================
 * Helper functions
 *================*/

/*
 * Allocates a new kernel stack. Returns null when not enough memory.
 */
static char *alloc_stack() { return page_alloc_n(DEFAULT_STACK_SIZE_PAGES); }

/*
 * Frees an existing kernel stack.
 */
static void free_stack(char *stack)
{
    page_free_n(stack, DEFAULT_STACK_SIZE_PAGES);
}

/*==========
 * Functions
 *=========*/

/*
 * Initializes the kthread_allocator.
 */
void kthread_init()
{
    KASSERT(__builtin_popcount(DEFAULT_STACK_SIZE_PAGES) == 1 &&
            "stack size should be a power of 2 pages to reduce fragmentation");
    kthread_allocator = slab_allocator_create("kthread", sizeof(kthread_t));
    KASSERT(kthread_allocator);
}

/*
 * Creates and initializes a thread.
 * Returns a new kthread, or NULL on failure.
 *
 * Hints:
 * Use kthread_allocator to allocate a kthread
 * Use alloc_stack() to allocate a kernel stack
 * Use context_setup() to set up the thread's context - 
 *  also use DEFAULT_STACK_SIZE and the process's pagetable (p_pml4)
 * Remember to initialize all the thread's fields
 * Remember to add the thread to proc's threads list
 * Initialize the thread's kt_state to KT_NO_STATE
 * Initialize the thread's kt_recent_core to ~0UL (unsigned -1)
 */
kthread_t *kthread_create(proc_t *proc, kthread_func_t func, long arg1,
                          void *arg2)
{
    kthread_t *new_thread = slab_obj_alloc(kthread_allocator);
    if (new_thread == NULL)
    {
        return NULL;
    }
    new_thread->kt_state = KT_NO_STATE;

    new_thread->kt_kstack = alloc_stack();
    if (new_thread->kt_kstack == NULL)
    {
        slab_obj_free(kthread_allocator, new_thread);
        return NULL;
    }

    new_thread->kt_proc = proc;
    context_setup(&new_thread->kt_ctx, func, arg1, arg2, new_thread->kt_kstack, 
        DEFAULT_STACK_SIZE, new_thread->kt_proc->p_pml4);

    // give default values to rest of struct
    new_thread->kt_retval = NULL;
    new_thread->kt_errno = 0;
    new_thread->kt_cancelled = 0;
    new_thread->kt_wchan = NULL;
    list_link_init(&new_thread->kt_plink);
    list_link_init(&new_thread->kt_qlink);
    list_init(&new_thread->kt_mutexes);
    new_thread->kt_recent_core = 0;

    // put this thread on the process's thread list
    list_insert_tail(&proc->p_threads, &new_thread->kt_plink);

    return new_thread;
}

/*
 * Creates and initializes a thread that is a clone of thr.
 * Returns a new kthread, or null on failure.
 * 
 * P.S. Note that you do not need to implement this function until VM.
 *
 * Hints:
 * The only parts of the context that must be initialized are c_kstack and
 * c_kstacksz. The thread's process should be set outside of this function. Copy
 * over thr's retval, errno, and cancelled; other fields should be freshly
 * initialized. See kthread_create() for more hints.
 */
kthread_t *kthread_clone(kthread_t *thr)
{
    kthread_t *new_thread = slab_obj_alloc(kthread_allocator);
    if (new_thread == NULL)
    {
        return NULL;
    }

    // copy the stack
    char * stk = alloc_stack();
    if (stk == NULL)
    {
        slab_obj_free(kthread_allocator, new_thread);
        return NULL;
    }
    // if not null, set the stack
    new_thread->kt_kstack = (char *)stk;
    new_thread->kt_ctx.c_kstacksz = DEFAULT_STACK_SIZE;
    new_thread->kt_ctx.c_kstack = (uintptr_t)stk;

    // set the retval, errno, cancelled
    new_thread->kt_retval = thr->kt_retval;
    new_thread->kt_errno = thr->kt_errno;
    new_thread->kt_cancelled = thr->kt_cancelled;
    new_thread->kt_recent_core = ~0UL;

    // null fields
    new_thread->kt_wchan = thr->kt_wchan;
    new_thread->kt_proc = NULL;
    new_thread->kt_ctx.c_rbp = 0;
    new_thread->kt_ctx.c_rsp = 0;
    new_thread->kt_ctx.c_rip = 0;
    new_thread->kt_ctx.c_pml4 = NULL;
    new_thread->kt_preemption_count = 0;
    // freshly initialize the rest of the fields
    list_link_init(&new_thread->kt_plink);
    list_link_init(&new_thread->kt_qlink);
    list_init(&new_thread->kt_mutexes);

    return new_thread;
}

/*
 * Free the thread's stack, remove it from its process's list of threads, and
 * free the kthread_t struct itself. Protect access to the kthread using its
 * kt_lock.
 *
 * You cannot destroy curthr.
 */
void kthread_destroy(kthread_t *thr)
{
    KASSERT(thr != curthr);
    KASSERT(thr && thr->kt_kstack);
    if (thr->kt_state != KT_EXITED)
        panic("destroying thread in state %d\n", thr->kt_state);
    free_stack(thr->kt_kstack);
    if (list_link_is_linked(&thr->kt_plink))
        list_remove(&thr->kt_plink);

    slab_obj_free(kthread_allocator, thr);
}

/*
 * Sets the thread's return value and cancels the thread.
 *
 * Note: Check out the use of check_curthr_cancelled() in syscall_handler()
 * to see how a thread eventually notices it is cancelled and handles exiting
 * itself.
 *
 * Hints:
 * This should not be called on curthr.
 * Use sched_cancel() to actually mark the thread as cancelled. This way you
 * can take care of all cancellation cases. 
 */
void kthread_cancel(kthread_t *thr, void *retval)
{
    KASSERT(thr != curthr);

    thr->kt_retval = (void *)retval;
    sched_cancel(thr);
}

/*
 * Wrapper around proc_thread_exiting().
 */
void kthread_exit(void *retval)
{
    proc_thread_exiting(retval);
}