#ifndef fiber_h
#define fiber_h 1

#if !defined(__x86_64__)
#error "This code is only for x86-64 architecture"
#endif // __x86_64__

#include <cstddef>
#include <cstdint>
#include <functional>
#include <memory>
#include <pthread.h>
#include <queue>
#include <vector>

namespace as
{
    //---------------------------------------------------------
    // context: cpu state_type and swap function
    //---------------------------------------------------------

    struct context
    {
        void *rsp, *rbx, *rbp, *r12, *r13, *r14, *r15, *rip;

        void swap(context &other) noexcept
        {
            asm volatile("   lea 0f(%%rip), %%rax \n"
                         "   mov %%rsp, 0(%%rdi)  \n"
                         "   mov %%rbx, 8(%%rdi)  \n"
                         "   mov %%rbp, 16(%%rdi) \n"
                         "   mov %%r12, 24(%%rdi) \n"
                         "   mov %%r13, 32(%%rdi) \n"
                         "   mov %%r14, 40(%%rdi) \n"
                         "   mov %%r15, 48(%%rdi) \n"
                         "   mov %%rax, 56(%%rdi) \n"
                         "   mov  0(%%rsi), %%rsp \n"
                         "   mov  8(%%rsi), %%rbx \n"
                         "   mov 16(%%rsi), %%rbp \n"
                         "   mov 24(%%rsi), %%r12 \n"
                         "   mov 32(%%rsi), %%r13 \n"
                         "   mov 40(%%rsi), %%r14 \n"
                         "   mov 48(%%rsi), %%r15 \n"
                         "   jmp *56(%%rsi)       \n"
                         "0:                      \n"
                         :
                         : "D"(this), "S"(&other)
                         : "memory", "rax");
        }
    };

    //---------------------------------------------------------
    // stack: fiber state_type storage on the heap
    //---------------------------------------------------------

    struct stack : public std::unique_ptr<std::byte[]>
    {
        static const size_t size = 1 << 16; // 64K x 1B = 64KB

        stack() noexcept
            : std::unique_ptr<std::byte[]>(std::make_unique<std::byte[]>(size))
        {
        }

        void *&top() noexcept
        {
            return *reinterpret_cast<void **>(
                (reinterpret_cast<std::uintptr_t>(
                     reinterpret_cast<std::byte *>(get()) + size) &
                 ~static_cast<std::uintptr_t>(0xF)) -
                sizeof(void *));
        }
    };

    //---------------------------------------------------------
    // fiber: user level thread mimicking std::thread interface
    //---------------------------------------------------------

    class fiber
    {
    public:
        class id
        {
        public:
            id(const void *p = nullptr) : ptr(p) {}
            operator const void *() const { return ptr; }

        private:
            const void *ptr = nullptr;
        };

        fiber() noexcept = default;

        template <typename Callable, typename... Args>
        fiber(Callable &&callable, Args &&...args) noexcept;

        fiber(const fiber &) = delete;
        fiber &operator=(const fiber &) = delete;

        fiber(fiber &&other) noexcept
            : state(std::exchange(other.state, nullptr))
        {
        }

        fiber &operator=(fiber &&other) noexcept
        {
            if (this != &other)
            {
                if (joinable())
                    std::terminate();
                state = std::exchange(other.state, nullptr);
            }
            return *this;
        }

        ~fiber() noexcept
        {
            if (joinable())
                std::terminate();
        }

        void detach()
        {
            if (!joinable())
                throw std::logic_error("detach on non-joinable as::fiber!");
            state.reset();
        }

        id get_id() const noexcept { return id(state.get()); }

        void join();

        bool joinable() const noexcept { return static_cast<bool>(state); }

    private:
        friend class fiber_scheduler;

        static void entry();

        struct state_type // real fiber state
        {
            state_type(std::function<void()> &&f)
                : ctx{.rsp = &stk.top(),
                      .rip = reinterpret_cast<void *>(entry)},
                  function(f)
            {
            }

            stack stk;                      // heap-allocated stack
            context ctx;                    // cpu context
            std::function<void()> function; // user function to execute
            bool finished = false;          // execution finished?
            std::exception_ptr exception;   // exception thrown by function
        };

        std::shared_ptr<state_type> state = nullptr;
    };

    //---------------------------------------------------------
    // fiber_scheduler: fiber scheduler
    //---------------------------------------------------------

    class fiber_scheduler
    {
    public:
        using state_type = fiber::state_type;

        // ~fiber_scheduler() noexcept { while (run()); }

        static fiber_scheduler &active()
        {
            static thread_local fiber_scheduler instance;
            return instance;
        }

        void enqueue(std::shared_ptr<state_type> s)
        {
            ready_queue.push(std::move(s));
        }

        void exit()
        {
            if (!running_state)
                return;
            running_state->finished = true;
            running_state->ctx.swap(contexts.back());
            std::terminate();
        }

        bool run()
        {
            if (ready_queue.empty())
                return false;
            auto saved_owner = std::move(running_state);
            running_state = std::move(ready_queue.front());
            ready_queue.pop();
            contexts.emplace_back();
            contexts.back().swap(running_state->ctx);
            contexts.pop_back();
            if (!running_state->finished)
                ready_queue.push(std::move(running_state));
            running_state = std::move(saved_owner);
            return true;
        }

        state_type *running() const { return running_state.get(); }

        void yield()
        {
            if (!running_state)
                return;
            running_state->ctx.swap(contexts.back());
        }

    private:
        std::vector<context> contexts;
        std::queue<std::shared_ptr<state_type>> ready_queue;
        std::shared_ptr<state_type> running_state;
    };

    //---------------------------------------------------------
    // fiber methods implementations
    //---------------------------------------------------------

    template <typename Callable, typename... Args>
    fiber::fiber(Callable &&callable, Args &&...args) noexcept
        : state(std::make_shared<fiber::state_type>(std::bind_front(
              std::forward<Callable>(callable), std::forward<Args>(args)...)))
    {
        fiber_scheduler::active().enqueue(state);
    }

    inline void fiber::entry()
    {
        auto &fiber_scheduler = fiber_scheduler::active();
        if (auto *state_type = fiber_scheduler.running())
        {
            try
            {
                state_type->function();
            }
            catch (...)
            {
                state_type->exception = std::current_exception();
            }
            fiber_scheduler.exit();
        }
    }

    inline void fiber::join()
    {
        if (!joinable())
            throw std::logic_error("join on non-joinable fiber!");
        if (fiber_scheduler::active().running() == state.get())
            throw std::logic_error("join on current fiber would deadlock!");
        while (!state->finished)
            if (!fiber_scheduler::active().run())
                throw std::runtime_error("join with empty fiber_scheduler!");
        if (state->exception)
            std::rethrow_exception(state->exception);
        state.reset();
    }

    //---------------------------------------------------------
    // this_fiber: local fiber operations
    //---------------------------------------------------------

    namespace this_fiber
    {
        using id = fiber::id;

        id get_id()
        {
            if (auto *s = fiber_scheduler::active().running(); s)
                return id(s);
            else
                return id(reinterpret_cast<const void *>(pthread_self()));
        }

        void yield() { fiber_scheduler::active().yield(); }

    } // namespace this_fiber

} // namespace as

#endif // fiber_h