/* Copyright 2019, Adrien Destugues, pulkomandy@pulkomandy.tk
 * Distributed under the terms of the MIT License.
 */


#include <string.h>

#include <arch_cpu.h>
#include <arch_debug.h>
#include <arch/thread.h>
#include <boot/stage2.h>
#include <commpage.h>
#include <kernel.h>
#include <thread.h>
#include <team.h>
#include <vm/vm_types.h>
#include <vm/VMAddressSpace.h>

#include "RISCV64VMTranslationMap.h"


extern "C" void SVecU();


status_t
arch_thread_init(struct kernel_args *args)
{
        // Initialize the static initial arch_thread state (sInitialState).
        // Currently nothing to do, i.e. zero initialized is just fine.

        return B_OK;
}


status_t
arch_team_init_team_struct(Team *team, bool kernel)
{
        // Nothing to do. The structure is empty.
        return B_OK;
}


status_t
arch_thread_init_thread_struct(Thread *thread)
{
        return B_OK;
}


void
arch_thread_init_kthread_stack(Thread* thread, void* _stack, void* _stackTop,
        void (*function)(void*), const void* data)
{
        memset(&thread->arch_info.context, 0, sizeof(arch_context));
        thread->arch_info.context.sp = (addr_t)_stackTop;
        thread->arch_info.context.s[0] = 0; // fp
        thread->arch_info.context.s[1] = (addr_t)function;
        thread->arch_info.context.s[2] = (addr_t)data;
        thread->arch_info.context.ra = (addr_t)arch_thread_entry;

        memset(&thread->arch_info.fpuContext, 0, sizeof(fpu_context));
}


status_t
arch_thread_init_tls(Thread *thread)
{
        thread->user_local_storage =
                thread->user_stack_base + thread->user_stack_size;
        return B_OK;
}


void
arch_thread_context_switch(Thread *from, Thread *to)
{
        /*
        dprintf("arch_thread_context_switch(%p(%s), %p(%s))\n", from, from->name,
                to, to->name);
        */

        auto fromMap = (RISCV64VMTranslationMap*)from->team->address_space->TranslationMap();
        auto toMap = (RISCV64VMTranslationMap*)to->team->address_space->TranslationMap();

        int cpu = to->cpu->cpu_num;
        toMap->ActiveOnCpus().SetBitAtomic(cpu);
        fromMap->ActiveOnCpus().ClearBitAtomic(cpu);

        // TODO: save/restore FPU only if needed
        save_fpu(&from->arch_info.fpuContext);
        restore_fpu(&to->arch_info.fpuContext);

        SetSatp(toMap->Satp());
        FlushTlbAllAsid(0);

        arch_context_switch(&from->arch_info.context, &to->arch_info.context);
}


void
arch_thread_dump_info(void *info)
{
}


status_t
arch_thread_enter_userspace(Thread *thread, addr_t entry, void *arg1,
        void *arg2)
{
        //dprintf("arch_thread_enter_uspace(%" B_PRId32 "(%s))\n", thread->id, thread->name);

        addr_t commpageAdr = (addr_t)thread->team->commpage_address;
        addr_t threadExitAddr;
        ASSERT(user_memcpy(&threadExitAddr,
                &((addr_t*)commpageAdr)[COMMPAGE_ENTRY_RISCV64_THREAD_EXIT],
                sizeof(threadExitAddr)) >= B_OK);
        threadExitAddr += commpageAdr;

        disable_interrupts();

        arch_stack* stackHeader = (arch_stack*)thread->kernel_stack_top - 1;
        stackHeader->thread = thread;

        iframe frame;
        memset(&frame, 0, sizeof(frame));

        SstatusReg status{.val = Sstatus()};
        status.pie = (1 << modeS); // enable interrupts when enter userspace
        status.spp = modeU;

        frame.status = status.val;
        frame.epc = entry;
        frame.a0 = (addr_t)arg1;
        frame.a1 = (addr_t)arg2;
        frame.ra = threadExitAddr;
        frame.sp = thread->user_stack_base + thread->user_stack_size;
        frame.tp = thread->user_local_storage;

        arch_load_user_iframe(stackHeader, &frame);

        // never return
        return B_ERROR;
}


bool
arch_on_signal_stack(Thread *thread)
{
        struct iframe* frame = thread->arch_info.userFrame;
        if (frame == NULL) {
                panic("arch_on_signal_stack(): No user iframe!");
                return false;
        }

        return frame->sp >= thread->signal_stack_base
                && frame->sp < thread->signal_stack_base
                        + thread->signal_stack_size;
}


static uint8*
get_signal_stack(Thread* thread, struct iframe* frame,
        struct sigaction* action, size_t spaceNeeded)
{
        // use the alternate signal stack if we should and can
        if (
                thread->signal_stack_enabled &&
                (action->sa_flags & SA_ONSTACK) != 0 && (
                        frame->sp < thread->signal_stack_base ||
                        frame->sp >= thread->signal_stack_base + thread->signal_stack_size
                )
        ) {
                addr_t stackTop = thread->signal_stack_base
                        + thread->signal_stack_size;
                return (uint8*)ROUNDDOWN(stackTop - spaceNeeded, 16);
        }
        return (uint8*)ROUNDDOWN(frame->sp - spaceNeeded, 16);
}


status_t
arch_setup_signal_frame(Thread *thread, struct sigaction *sa,
        struct signal_frame_data *signalFrameData)
{
        // dprintf("%s(%" B_PRId32 "(%s))\n", __func__, thread->id, thread->name);
        iframe* frame = thread->arch_info.userFrame;

        // fill signal context
        signalFrameData->context.uc_mcontext.x[ 0] = frame->ra;
        signalFrameData->context.uc_mcontext.x[ 1] = frame->sp;
        signalFrameData->context.uc_mcontext.x[ 2] = frame->gp;
        signalFrameData->context.uc_mcontext.x[ 3] = frame->tp;
        signalFrameData->context.uc_mcontext.x[ 4] = frame->t0;
        signalFrameData->context.uc_mcontext.x[ 5] = frame->t1;
        signalFrameData->context.uc_mcontext.x[ 6] = frame->t2;
        signalFrameData->context.uc_mcontext.x[ 7] = frame->fp;
        signalFrameData->context.uc_mcontext.x[ 8] = frame->s1;
        signalFrameData->context.uc_mcontext.x[ 9] = frame->a0;
        signalFrameData->context.uc_mcontext.x[10] = frame->a1;
        signalFrameData->context.uc_mcontext.x[11] = frame->a2;
        signalFrameData->context.uc_mcontext.x[12] = frame->a3;
        signalFrameData->context.uc_mcontext.x[13] = frame->a4;
        signalFrameData->context.uc_mcontext.x[14] = frame->a5;
        signalFrameData->context.uc_mcontext.x[15] = frame->a6;
        signalFrameData->context.uc_mcontext.x[16] = frame->a7;
        signalFrameData->context.uc_mcontext.x[17] = frame->s2;
        signalFrameData->context.uc_mcontext.x[18] = frame->s3;
        signalFrameData->context.uc_mcontext.x[19] = frame->s4;
        signalFrameData->context.uc_mcontext.x[20] = frame->s5;
        signalFrameData->context.uc_mcontext.x[21] = frame->s6;
        signalFrameData->context.uc_mcontext.x[22] = frame->s7;
        signalFrameData->context.uc_mcontext.x[23] = frame->s8;
        signalFrameData->context.uc_mcontext.x[24] = frame->s9;
        signalFrameData->context.uc_mcontext.x[25] = frame->s10;
        signalFrameData->context.uc_mcontext.x[26] = frame->s11;
        signalFrameData->context.uc_mcontext.x[27] = frame->t3;
        signalFrameData->context.uc_mcontext.x[28] = frame->t4;
        signalFrameData->context.uc_mcontext.x[29] = frame->t5;
        signalFrameData->context.uc_mcontext.x[30] = frame->t6;
        signalFrameData->context.uc_mcontext.pc = frame->epc;
        // TODO: don't assume that kernel code don't use FPU
        save_fpu((fpu_context*)&signalFrameData->context.uc_mcontext.f[0]);
        // end of fill signal context

        signal_get_user_stack(frame->sp, &signalFrameData->context.uc_stack);
/*
        dprintf("  thread->signal_stack_enabled: %d\n",
                thread->signal_stack_enabled);
        if (thread->signal_stack_enabled) {
                dprintf("  signal stack: 0x%" B_PRIxADDR " - 0x%" B_PRIxADDR "\n",
                        thread->signal_stack_base,
                        thread->signal_stack_base + thread->signal_stack_size
                );
        }
*/
        signalFrameData->syscall_restart_return_value = thread->arch_info.oldA0;

        uint8* userStack = get_signal_stack(thread, frame, sa,
                sizeof(*signalFrameData));
        // dprintf("  user stack: 0x%" B_PRIxADDR "\n", (addr_t)userStack);
        status_t res = user_memcpy(userStack, signalFrameData,
                sizeof(*signalFrameData));
        if (res < B_OK)
                return res;

        addr_t commpageAdr = (addr_t)thread->team->commpage_address;
        // dprintf("  commpageAdr: 0x%" B_PRIxADDR "\n", commpageAdr);
        addr_t signalHandlerAddr;
        ASSERT(user_memcpy(&signalHandlerAddr,
                &((addr_t*)commpageAdr)[COMMPAGE_ENTRY_RISCV64_SIGNAL_HANDLER],
                sizeof(signalHandlerAddr)) >= B_OK);
        signalHandlerAddr += commpageAdr;

        frame->ra = frame->epc;
        frame->sp = (addr_t)userStack;
        frame->epc = signalHandlerAddr;
        frame->a0 = frame->sp;

        // WriteTrapInfo();

        return B_OK;
}


int64
arch_restore_signal_frame(struct signal_frame_data* signalFrameData)
{
        // dprintf("arch_restore_signal_frame()\n");
        iframe* frame = thread_get_current_thread()->arch_info.userFrame;

        thread_get_current_thread()->arch_info.oldA0
                = signalFrameData->syscall_restart_return_value;

        frame->ra  = signalFrameData->context.uc_mcontext.x[ 0];
        frame->sp  = signalFrameData->context.uc_mcontext.x[ 1];
        frame->gp  = signalFrameData->context.uc_mcontext.x[ 2];
        frame->tp  = signalFrameData->context.uc_mcontext.x[ 3];
        frame->t0  = signalFrameData->context.uc_mcontext.x[ 4];
        frame->t1  = signalFrameData->context.uc_mcontext.x[ 5];
        frame->t2  = signalFrameData->context.uc_mcontext.x[ 6];
        frame->fp  = signalFrameData->context.uc_mcontext.x[ 7];
        frame->s1  = signalFrameData->context.uc_mcontext.x[ 8];
        frame->a0  = signalFrameData->context.uc_mcontext.x[ 9];
        frame->a1  = signalFrameData->context.uc_mcontext.x[10];
        frame->a2  = signalFrameData->context.uc_mcontext.x[11];
        frame->a3  = signalFrameData->context.uc_mcontext.x[12];
        frame->a4  = signalFrameData->context.uc_mcontext.x[13];
        frame->a5  = signalFrameData->context.uc_mcontext.x[14];
        frame->a6  = signalFrameData->context.uc_mcontext.x[15];
        frame->a7  = signalFrameData->context.uc_mcontext.x[16];
        frame->s2  = signalFrameData->context.uc_mcontext.x[17];
        frame->s3  = signalFrameData->context.uc_mcontext.x[18];
        frame->s4  = signalFrameData->context.uc_mcontext.x[19];
        frame->s5  = signalFrameData->context.uc_mcontext.x[20];
        frame->s6  = signalFrameData->context.uc_mcontext.x[21];
        frame->s7  = signalFrameData->context.uc_mcontext.x[22];
        frame->s8  = signalFrameData->context.uc_mcontext.x[23];
        frame->s9  = signalFrameData->context.uc_mcontext.x[24];
        frame->s10 = signalFrameData->context.uc_mcontext.x[25];
        frame->s11 = signalFrameData->context.uc_mcontext.x[26];
        frame->t3  = signalFrameData->context.uc_mcontext.x[27];
        frame->t4  = signalFrameData->context.uc_mcontext.x[28];
        frame->t5  = signalFrameData->context.uc_mcontext.x[29];
        frame->t6  = signalFrameData->context.uc_mcontext.x[30];
        frame->epc = signalFrameData->context.uc_mcontext.pc;
        restore_fpu((fpu_context*)&signalFrameData->context.uc_mcontext.f[0]);

        return frame->a0;
}


/**     Saves everything needed to restore the frame in the child fork in the
 *      arch_fork_arg structure to be passed to arch_restore_fork_frame().
 *      Also makes sure to return the right value.
 */

void
arch_store_fork_frame(struct arch_fork_arg *arg)
{
/*
        dprintf("arch_store_fork_frame()\n");
        dprintf("  arg: %p\n", arg);
        dprintf("  userFrame: %p\n",
                thread_get_current_thread()->arch_info.userFrame);
*/
        memcpy(&arg->frame, thread_get_current_thread()->arch_info.userFrame,
                sizeof(iframe));
        arg->frame.a0 = 0; // fork return value
}


/** Restores the frame from a forked team as specified by the provided
 *      arch_fork_arg structure.
 *      Needs to be called from within the child team, ie. instead of
 *      arch_thread_enter_uspace() as thread "starter".
 *      This function does not return to the caller, but will enter userland
 *      in the child team at the same position where the parent team left of.
 */

void
arch_restore_fork_frame(struct arch_fork_arg *arg)
{
        //dprintf("arch_restore_fork_frame(%p)\n", arg);
        //dprintf("  thread: %" B_PRId32 "(%s))\n", thread_get_current_thread()->id,
        //      thread_get_current_thread()->name);
        //dprintf("  kernel SP: %#" B_PRIxADDR "\n", thread_get_current_thread()->kernel_stack_top);
        //dprintf("  user PC: "); WritePC(arg->frame.epc); dprintf("\n");

        disable_interrupts();

        arch_stack* stackHeader = (arch_stack*)thread_get_current_thread()->kernel_stack_top - 1;
        stackHeader->thread = thread_get_current_thread();
        SstatusReg status{.val = Sstatus()};
        status.pie = (1 << modeS); // enable interrupts when enter userspace
        status.spp = modeU;
        arg->frame.status = status.val;
        arch_load_user_iframe(stackHeader, &arg->frame);
}