#include "thread/Thread.h"
#include "Log.h"
#include "memory/MemoryManager.h"
#include <luna/Alignment.h>
#include <luna/CString.h>
#include <luna/Check.h>

bool is_in_kernel(Registers* regs)
{
    return regs->cs == 8;
}

void Thread::set_ip(u64 ip)
{
    regs.rip = ip;
}

u64 Thread::ip()
{
    return regs.rip;
}

void Thread::set_sp(u64 sp)
{
    regs.rsp = sp;
}

u64 Thread::sp()
{
    return regs.rsp;
}

void Thread::set_return(u64 ret)
{
    regs.rax = ret;
}

u64 Thread::return_register()
{
    return regs.rax;
}

void Thread::init_regs_kernel()
{
    memset(&regs, 0, sizeof(Registers));
    regs.cs = 0x08;
    regs.ss = 0x10;
    regs.rflags = 1 << 9; // IF (Interrupt enable flag)
}

void Thread::init_regs_user()
{
    memset(&regs, 0, sizeof(Registers));
    regs.cs = 0x18 | 3;
    regs.ss = 0x20 | 3;
    regs.rflags = 1 << 9; // IF (Interrupt enable flag)
}

void Thread::set_arguments(u64 arg1, u64 arg2, u64 arg3, u64 arg4)
{
    regs.rdi = arg1;
    regs.rsi = arg2;
    regs.rdx = arg3;
    regs.rcx = arg4;
}

void switch_context(Thread* old_thread, Thread* new_thread, Registers* regs)
{
    if (!old_thread->is_idle()) memcpy(&old_thread->regs, regs, sizeof(Registers));

    memcpy(regs, &new_thread->regs, sizeof(Registers));
}

// FIXME: Move this function to a common location (also used in ThreadImage)
Result<u64> Thread::push_mem_on_stack(const u8* mem, usize size)
{
    if ((regs.rsp - size) < stack.bottom()) return err(E2BIG);

    if (!MemoryManager::validate_user_write((void*)(regs.rsp - size), size)) return err(EFAULT);

    regs.rsp -= size;

    memcpy((void*)regs.rsp, mem, size);

    return regs.rsp;
}

Result<u64> Thread::pop_mem_from_stack(u8* mem, usize size)
{
    if ((regs.rsp + size) > stack.top()) return err(E2BIG);

    if (!MemoryManager::validate_user_read((void*)regs.rsp, size)) return err(EFAULT);

    memcpy(mem, (void*)regs.rsp, size);

    regs.rsp += size;

    return regs.rsp;
}

bool Thread::deliver_signal(int signo, Registers* current_regs)
{
    auto handler = signal_handlers[signo - 1];
    check(handler.sa_handler != SIG_IGN);
    check(handler.sa_handler != SIG_DFL);

    memcpy(&regs, current_regs, sizeof(regs));

    kinfoln("signal: delivering signal %d for thread %ld, ip=%p, sp=%p, handler=%p, sigreturn=%p", signo, id,
            (void*)regs.rip, (void*)regs.rsp, (void*)handler.sa_handler, (void*)handler.__sa_sigreturn);

    regs.rsp -= 128; // Skip the red zone

    fp_data.save();

    if (push_mem_on_stack((u8*)current_regs, sizeof(*current_regs)).has_error()) return false;
    if (push_mem_on_stack((u8*)&signal_mask, sizeof(signal_mask)).has_error()) return false;
    if (push_mem_on_stack((u8*)fp_data.data(), fp_data.size()).has_error()) return false;

    u64 rsp = regs.rsp;

    regs.rsp = align_down<16>(regs.rsp);
    if (push_mem_on_stack((u8*)&rsp, sizeof(u64)).has_error()) return false;
    if (push_mem_on_stack((u8*)&handler.__sa_sigreturn, sizeof(void*)).has_error()) return false;

    signal_mask = handler.sa_mask;
    if ((handler.sa_flags & SA_NODEFER) == 0) signal_mask |= (1 << (signo - 1));

    rsp = regs.rsp;

    init_regs_user();
    regs.rsp = rsp;
    regs.rip = (u64)handler.sa_handler;
    regs.rdi = signo;

    memcpy(current_regs, &regs, sizeof(regs));

    if (handler.sa_flags & SA_RESETHAND)
    {
        handler.sa_handler = SIG_DFL;
        handler.sa_mask = 0;
        handler.sa_flags = 0;
        signal_handlers[signo - 1] = handler;
    }

    return true;
}

void Thread::sigreturn(Registers* current_regs)
{
    memcpy(&regs, current_regs, sizeof(regs));

    u64 rflags = current_regs->rflags;

    u64 rsp;
    pop_mem_from_stack((u8*)&rsp, sizeof(rsp));
    regs.rsp = rsp;
    pop_mem_from_stack((u8*)fp_data.data(), fp_data.size());
    pop_mem_from_stack((u8*)&signal_mask, sizeof(signal_mask));
    pop_mem_from_stack((u8*)current_regs, sizeof(*current_regs));
    memcpy(&regs, current_regs, sizeof(regs));
    regs.cs = 0x18 | 3;
    regs.ss = 0x20 | 3;
    regs.rflags = (rflags & ~0xdff) | (regs.rflags & 0xdff);

    fp_data.restore();

    kinfoln("sigreturn: restored program state, sp=%p, ip=%p", (void*)regs.rsp, (void*)regs.rip);

    memcpy(current_regs, &regs, sizeof(regs));
}