#include <algorithm>
#include <array>
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>

static const size_t STACK_SIZE = 16777216;
static const size_t DEFAULT_THREADS = 4;
struct Runtime;
static Runtime* runtime;

enum class ThreadState {
    Available,
    Ready,
    Running,
};

struct ThreadContext {
    uint64_t rsp, r15, r14, r13, r12, rbx, rbp;
    ThreadContext() : rsp(), r15(), r14(), r13(), r12(), rbx(), rbp() {}
    static __attribute__((noinline)) __attribute__((naked)) void
    switch_to(ThreadContext&, ThreadContext&) {
        __asm__ volatile(
          "movq %rsp, (%rdi)\n\t"
          "movq %r15, 0x8(%rdi)\n\t"
          "movq %r14, 0x10(%rdi)\n\t"
          "movq %r13, 0x18(%rdi)\n\t"
          "movq %r12, 0x20(%rdi)\n\t"
          "movq %rbx, 0x28(%rdi)\n\t"
          "movq %rbp, 0x30(%rdi)\n\t"

          "movq (%rsi), %rsp\n\t"
          "movq 0x8(%rsi), %r15\n\t"
          "movq 0x10(%rsi), %r14\n\t"
          "movq 0x18(%rsi), %r13\n\t"
          "movq 0x20(%rsi), %r12\n\t"
          "movq 0x28(%rsi), %rbx\n\t"
          "movq 0x30(%rsi), %rbp\n\t"

          "ret");
    }
};
static_assert(offsetof(ThreadContext, rsp) == 0, "unexpected offset");
static_assert(offsetof(ThreadContext, r15) == 010, "unexpected offset");
static_assert(offsetof(ThreadContext, r14) == 020, "unexpected offset");
static_assert(offsetof(ThreadContext, r13) == 030, "unexpected offset");
static_assert(offsetof(ThreadContext, r12) == 040, "unexpected offset");
static_assert(offsetof(ThreadContext, rbx) == 050, "unexpected offset");
static_assert(offsetof(ThreadContext, rbp) == 060, "unexpected offset");

struct Thread {
    size_t id;
    void* stack;
    ThreadContext ctx;
    ThreadState state;
    Thread() : id(0), stack(), state(ThreadState::Available) {
        stack = malloc(STACK_SIZE);
        if (!stack) {
            fputs("Could not allocate memory for new thread.\n", stderr);
            exit(1);
        }
    }
};

struct Runtime {
    std::array<Thread, 1 + DEFAULT_THREADS> threads;
    size_t current;
    Runtime() : threads(), current(0) {
        threads.front().state = ThreadState::Running;
        for (size_t i = 1; i <= DEFAULT_THREADS; ++i) { threads[i].id = i; }
    }
    void init() { runtime = this; }
    void run() {
        while (yield())
            ;
        exit(0);
    }
    void ret() {
        if (current) {
            threads[current].state = ThreadState::Available;
            yield();
        }
    }
    bool yield() {
        auto next_ready = std::find_if(
          threads.begin() + current, threads.end(),
          [](Thread const& th) { return th.state == ThreadState::Ready; });
        if (next_ready == threads.end()) {
            next_ready = std::find_if(
              threads.begin(), threads.begin() + current,
              [](Thread const& th) { return th.state == ThreadState::Ready; });
        }
        if (next_ready == threads.begin() + current) { return false; }

        if (threads[current].state != ThreadState::Available)
            threads[current].state = ThreadState::Ready;

        next_ready->state = ThreadState::Running;
        size_t old = current;
        current = next_ready - threads.begin();

        ThreadContext::switch_to(threads[old].ctx, next_ready->ctx);
        return true;
    }
    static void finish() { runtime->ret(); }
    void spawn(void (*func)()) {
        auto avail =
          std::find_if(threads.begin(), threads.end(), [](Thread const& th) {
              return th.state == ThreadState::Available;
          });
        if (avail == threads.end()) {
            fputs("Cannot spawn task. No available threads.\n", stderr);
            exit(2);
        }
        void (*guard)() = Runtime::finish;
        memcpy((unsigned char*) avail->stack + STACK_SIZE - 8, &guard, 8);
        memcpy((unsigned char*) avail->stack + STACK_SIZE - 16, &func, 8);
        avail->ctx.rsp =
          (uint64_t)((unsigned char*) avail->stack + STACK_SIZE - 16);
        avail->state = ThreadState::Ready;
    }
};

static void thread1() {
    puts("thread 1 started");
    for (int i = 0; i < 10; ++i) {
        printf("thread 1 counting %d\n", i);
        runtime->yield();
    }
}

static void thread2() {
    puts("thread 2 started");
    for (int i = 0; i < 10; ++i) {
        printf("thread 2 counting %d\n", i);
        runtime->yield();
    }
}

int main() {
    Runtime().init();
    runtime->spawn(thread1);
    runtime->spawn(thread2);
    runtime->run();
}