#include <rt/abort.h>
#include <rt/arch/init.h>
#include <rt/arch/irq.h>
#include <rt/arch/syscall.h>
#include <rt/context.h>
#include <rt/cycle.h>
#include <rt/idle.h>
#include <rt/interrupt.h>
#include <rt/log.h>
#include <rt/panic.h>
#include <rt/start.h>
#include <rt/syscall.h>
#include <rt/task.h>
#include <rt/tick.h>
#include <rt/trap.h>

#include <pthread.h>
#include <signal.h>
#include <sys/mman.h>
#include <sys/time.h>
#include <unistd.h>

#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>

#if RT_LOG_ENABLE
#include <stdarg.h>
#endif

#if defined(__aarch64__)
#include "aarch64/context.h"
#elif defined(__x86_64__)
#include "x86_64/context.h"
#else
#error "unsupported architecture"
#endif

#ifndef MAP_STACK
#define MAP_STACK 0
#endif

// Signals in ascending priority order. Later signals should mask earlier ones.
#define SIGPENDSYSCALL SIGUSR1
#define SIGTICK SIGALRM
#define SIGIRQ SIGUSR2

static sigset_t sigset_all, sigset_pendsyscall;

__attribute__((constructor)) static void init_sigsets(void)
{
    sigemptyset(&sigset_all);
    sigaddset(&sigset_all, SIGPENDSYSCALL);
    sigaddset(&sigset_all, SIGTICK);
    sigaddset(&sigset_all, SIGIRQ);

    sigemptyset(&sigset_pendsyscall);
    sigaddset(&sigset_pendsyscall, SIGPENDSYSCALL);
}

static void block_async_signals(sigset_t *old_sigset)
{
    pthread_sigmask(SIG_BLOCK, &sigset_all, old_sigset);
}

void rt_block_pending_syscalls(void)
{
    pthread_sigmask(SIG_BLOCK, &sigset_pendsyscall, NULL);
}

void rt_unblock_pending_syscalls(void)
{
    pthread_sigmask(SIG_UNBLOCK, &sigset_pendsyscall, NULL);
}

#if RT_LOG_ENABLE
void rt_logf(const char *format, ...)
{
    va_list vlist;
    va_start(vlist, format);
    sigset_t old_sigset;
    block_async_signals(&old_sigset);
    vprintf(format, vlist);
    fflush(stdout);
    pthread_sigmask(SIG_SETMASK, &old_sigset, NULL);
}
#endif

__attribute__((noreturn)) void rt_abort(void)
{
    abort();
}

__attribute__((noreturn)) void rt_panic(const char *msg)
{
    block_async_signals(NULL);
    fprintf(stderr, "%s\n", msg);
    fflush(stderr);
    abort();
}

__attribute__((noreturn)) void rt_task_entry(void);

void *rt_context_init(uintptr_t fn, uintptr_t arg, void *stack,
                      size_t stack_size)
{
    // The provided stack is unused.
    stack_size = (size_t)PTHREAD_STACK_MIN;
    stack = mmap(NULL, stack_size, PROT_READ | PROT_WRITE,
                 MAP_ANONYMOUS | MAP_PRIVATE | MAP_STACK, -1, 0);
    rt_assert(stack != MAP_FAILED, "failed to mmap a task stack");

    void *const sp = (char *)stack + stack_size;
    struct context *ctx = sp;
    --ctx;
#ifdef __aarch64__
    ctx->x19 = fn;
    ctx->x20 = arg;
    ctx->fp = 0;
    ctx->lr = (uint64_t)rt_task_entry;
#elif defined(__x86_64__)
    __asm__("stmxcsr (%0)" : : "r"(&ctx->mxcsr) : "memory");
    __asm__("fstcw (%0)" : : "r"(&ctx->x87cw) : "memory");
    ctx->in_signal = false;
    ctx->rbx = fn;
    ctx->r12 = arg;
    ctx->rbp = 0;
    ctx->return_addr = (uint64_t)rt_task_entry;
#endif
    return ctx;
}

void rt_syscall_pend(void)
{
    pthread_kill(pthread_self(), SIGPENDSYSCALL);
}

void rt_irq_pend(void)
{
    pthread_kill(pthread_self(), SIGIRQ);
}

bool rt_interrupt_is_active(void)
{
    sigset_t sigset;
    pthread_sigmask(SIG_BLOCK, NULL, &sigset);
    return sigismember(&sigset, SIGPENDSYSCALL) == 1;
}

static void tick_handler(int sig)
{
    (void)sig;
    rt_tick_advance();
}

__attribute__((weak, noreturn)) void rt_idle(void)
{
    for (;;)
    {
        pause();
    }
}

__attribute__((weak, noreturn)) void rt_trap(void)
{
    for (;;)
    {
#if defined(__aarch64__)
        __asm__("brk 0" :::);
#elif defined(__x86_64__)
        __asm__("int3" :::);
#endif
    }
}

void rt_task_drop_privilege(void)
{
}

uint32_t rt_cycle(void)
{
#if defined(__aarch64__)
    uint64_t cycles;
    __asm__ __volatile__("mrs %0, cntvct_el0" : "=r"(cycles));
    return (uint32_t)cycles;
#elif defined(__x86_64__)
    uint32_t cycles;
    __asm__ __volatile__("rdtsc" : "=a"(cycles) : : "edx");
    return cycles;
#else
    return 0;
#endif
}

__attribute__((weak)) void rt_irq_handler(void)
{
    rt_panic("IRQ was triggered with the default handler");
}

static void irq_handler(int arg)
{
    (void)arg;
    rt_irq_handler();
}

static void trap_handler(int arg)
{
    /* If we are not debugging the program, just exit cleanly. Examples use
     * rt_trap() to terminate. */
    (void)arg;
    exit(0);
}

void rt_signal_init_tick(const struct timeval *period)
{
    const struct itimerval timer = {
        .it_interval = *period,
        .it_value = *period,
    };
    setitimer(ITIMER_REAL, &timer, NULL);
}

void rt_signal_init(void)
{
    struct sigaction action;
    action.sa_flags = SA_RESTART;

    // Pending syscalls don't mask other signals.
    action.sa_handler = rt_pending_syscall_handler;
    sigemptyset(&action.sa_mask);
    sigaction(SIGPENDSYSCALL, &action, NULL);

    // IRQ and trap mask all rt-related signals.
    action.sa_handler = irq_handler;
    action.sa_mask = sigset_all;
    sigaction(SIGIRQ, &action, NULL);

    action.sa_handler = trap_handler;
    sigaction(SIGTRAP, &action, NULL);

    // The tick handler must mask SIGPENDSYSCALL.
    action.sa_handler = tick_handler;
    action.sa_mask = sigset_pendsyscall;
    sigaction(SIGTICK, &action, NULL);
}

__attribute__((weak)) int main(void)
{
    // Block pending syscalls until we have started running the first task.
    rt_block_pending_syscalls();
    rt_signal_init();
    static const struct timeval milli = {
        .tv_sec = 0,
        .tv_usec = 1000,
    };
    if (getenv("RT_SIGNAL_EXTERNAL_TICK") == NULL)
    {
        rt_signal_init_tick(&milli);
    }
    rt_start();
}