/* 
   Mathieu Stefani, 15 juin 2016
   
   Implementation of the Reactor
*/

#include "reactor.h"

namespace Aio {

struct Reactor::Impl {

    Impl(Reactor* reactor)
        : reactor_(reactor)
    { }

    virtual ~Impl() = default;

    virtual Reactor::Key addHandler(
            const std::shared_ptr<Handler>& handler, bool setKey) = 0;

    virtual std::vector<std::shared_ptr<Handler>> handlers(
            const Reactor::Key& key) const = 0;

    virtual void registerFd(
            const Reactor::Key& key,
            Fd fd,
            Polling::NotifyOn interest,
            Polling::Tag tag,
            Polling::Mode mode = Polling::Mode::Level) = 0;

    virtual void registerFdOneShot(
            const Reactor::Key& key,
            Fd fd,
            Polling::NotifyOn interest,
            Polling::Tag tag,
            Polling::Mode mode = Polling::Mode::Level) = 0;

    virtual void modifyFd(
            const Reactor::Key& key,
            Fd fd,
            Polling::NotifyOn interest,
            Polling::Tag tag,
            Polling::Mode mode = Polling::Mode::Level) = 0;

    virtual void runOnce() = 0;
    virtual void run() = 0;

    virtual void shutdown() = 0;

    Reactor* reactor_;
};

/* Synchronous implementation of the reactor that polls in the context
 * of the same thread
 */
struct SyncImpl : public Reactor::Impl {

    SyncImpl(Reactor* reactor)
        : Reactor::Impl(reactor)
        , handlers_()
        , shutdown_()
        , shutdownFd()
        , poller()
    {
        shutdownFd.bind(poller);
    }

    Reactor::Key addHandler(
            const std::shared_ptr<Handler>& handler, bool setKey = true) {

        handler->registerPoller(poller);

        handler->reactor_ = reactor_;

        auto key = handlers_.add(handler);
        if (setKey) handler->key_ = key;

        return key;
    }

    std::shared_ptr<Handler> handler(const Reactor::Key& key) const {
        return handlers_[key.data()];
    }

    std::vector<std::shared_ptr<Handler>> handlers(const Reactor::Key& key) const {
        std::vector<std::shared_ptr<Handler>> res;

        res.push_back(handler(key));
        return res;
    }

    void registerFd(
            const Reactor::Key& key,
            Fd fd,
            Polling::NotifyOn interest,
            Polling::Tag tag,
            Polling::Mode mode = Polling::Mode::Level) {

        auto pollTag = encodeTag(key, tag);
        poller.addFd(fd, interest, pollTag, mode);
    }

    void registerFdOneShot(
            const Reactor::Key& key,
            Fd fd,
            Polling::NotifyOn interest,
            Polling::Tag tag,
            Polling::Mode mode = Polling::Mode::Level) {

        auto pollTag = encodeTag(key, tag);
        poller.addFdOneShot(fd, interest, pollTag, mode);
    }

    void modifyFd(
            const Reactor::Key& key,
            Fd fd,
            Polling::NotifyOn interest,
            Polling::Tag tag,
            Polling::Mode mode = Polling::Mode::Level) {

        auto pollTag = encodeTag(key, tag);
        poller.rearmFd(fd, interest, pollTag, mode);
    }

    void runOnce() {
        if (handlers_.empty())
            throw std::runtime_error("You need to set at least one handler");

        std::chrono::milliseconds timeout(-1);

        for (;;) {
            std::vector<Polling::Event> events;
            int ready_fds;
            switch (ready_fds = poller.poll(events, 1024, timeout)) {
                case -1: break;
                case 0: break;
                default:
                    if (shutdown_) return;

                    handleFds(std::move(events));

                    timeout = std::chrono::milliseconds(-1);
            }
        }
    }

    void run() {
        handlers_.forEachHandler([](const std::shared_ptr<Handler> handler) {
            handler->context_.tid = std::this_thread::get_id();
        });

        while (!shutdown_)
            runOnce();
    }

    void shutdown() {
        shutdown_.store(true);
        shutdownFd.notify();
    }

    static constexpr size_t MaxHandlers() {
        return HandlerList::MaxHandlers;
    }

private:

    Polling::Tag encodeTag(const Reactor::Key& key, Polling::Tag tag) const {
        uint64_t value = tag.value();
        return HandlerList::encodeTag(key, value);
    }

    std::pair<size_t, uint64_t> decodeTag(const Polling::Tag& tag) const {
        return HandlerList::decodeTag(tag);
    }

    void handleFds(std::vector<Polling::Event> events) const {
        // Fast-path: if we only have one handler, do not bother scanning the fds to find
        // the right handlers
        if (handlers_.size() == 1)
            handlers_[0]->onReady(FdSet(std::move(events)));
        else {
            std::unordered_map<std::shared_ptr<Handler>, std::vector<Polling::Event>> fdHandlers;

            for (auto& event: events) {
                size_t index;
                uint64_t value;

                std::tie(index, value) = decodeTag(event.tag);
                auto handler = handlers_[index];
                auto& evs = fdHandlers[handler];
                evs.push_back(std::move(event));
            }

            for (auto& data: fdHandlers) {
                data.first->onReady(FdSet(std::move(data.second)));
            }
        }
    }

    struct HandlerList {

        // We are using the highest 8 bits of the fd to encode the index of the handler,
        // which gives us a maximum of 2**8 - 1 handler, 255
        static constexpr size_t HandlerBits = 8;
        static constexpr size_t HandlerShift = sizeof(uint64_t) - HandlerBits;
        static constexpr uint64_t DataMask = uint64_t(-1) >> HandlerBits;

        static constexpr size_t MaxHandlers = (1 << HandlerBits) - 1;

        HandlerList()
            :handlers()
            , index_()
        {
            std::fill(std::begin(handlers), std::end(handlers), nullptr);
        }

        HandlerList(const HandlerList& other) = delete;
        HandlerList& operator=(const HandlerList& other) = delete;

        HandlerList(HandlerList&& other) = default;
        HandlerList& operator=(HandlerList&& other) = default;

        HandlerList clone() const {
            HandlerList list;

            for (size_t i = 0; i < index_; ++i) {
                list.handlers[i] = handlers[i]->clone();
            }
            list.index_ = index_;

            return list;
        }

        Reactor::Key add(const std::shared_ptr<Handler>& handler) {
            if (index_ == MaxHandlers)
                throw std::runtime_error("Maximum handlers reached");

            Reactor::Key key(index_);
            handlers[index_++] = handler;

            return key;
        }

        std::shared_ptr<Handler> operator[](size_t index) const {
            return handlers[index];
        }

        std::shared_ptr<Handler> at(size_t index) const {
            if (index >= index_)
                throw std::runtime_error("Attempting to retrieve invalid handler");

            return handlers[index];
        }

        bool empty() const {
            return index_ == 0;
        }

        size_t size() const {
            return index_;
        }

        static Polling::Tag encodeTag(const Reactor::Key& key, uint64_t value) {
            auto index = key.data();
            // The reason why we are using the most significant bits to encode
            // the index of the handler is that in the fast path, we won't need
            // to shift the value to retrieve the fd if there is only one handler as
            // all the bits will already be set to 0.
            auto encodedValue = (index << HandlerShift) | value;
            return Polling::Tag(value);
        }

        static std::pair<size_t, uint64_t> decodeTag(const Polling::Tag& tag) {
            auto value = tag.value();
            size_t index = value >> HandlerShift;
            uint64_t fd = value & DataMask;

            return std::make_pair(index, fd);
        }

        template<typename Func>
        void forEachHandler(Func func) const {
            for (size_t i = 0; i < index_; ++i)
                func(handlers[i]);
        }

    private:
        std::array<std::shared_ptr<Handler>, MaxHandlers> handlers;
        size_t index_;
    };

    HandlerList handlers_;

    std::atomic<bool> shutdown_;
    NotifyFd shutdownFd;

    Polling::Epoll poller;
};

/* Asynchronous implementation of the reactor that spawns a number N of threads
 * and creates a polling fd per thread
 *
 * Implementation detail:
 *
 *  Here is how it works: the implementation simply starts a synchronous variant
 *  of the implementation in its own std::thread. When adding an handler, it will
 *  add a clone() of the handler to every worker (thread), and assign its own key
 *  to the handler. Here is where things start to get interesting. Here is how the
 *  key encoding works for every handler:
 *
 *  [     handler idx      ] [       worker idx         ]
 *  ------------------------ ----------------------------
 *       ^ 32 bits                   ^ 32 bits
 *  -----------------------------------------------------
 *                       ^ 64 bits
 *
 * Since we have up to 64 bits of data for every key, we encode the index of the handler
 * that has been assigned by the SyncImpl in the upper 32 bits, and encode the index of the
 * worker thread in the lowest 32 bits.
 *
 * When registering a fd for a given key, the AsyncImpl then knows which worker to use by
 * looking at the lowest 32 bits of the Key's data. The SyncImpl will then use the highest
 * 32 bits to retrieve the index of the handler.
 */

struct AsyncImpl : public Reactor::Impl {

    static constexpr uint32_t KeyMarker = 0xBADB0B;

    AsyncImpl(Reactor* reactor, size_t threads)
        : Reactor::Impl(reactor) {

        for (size_t i = 0; i < threads; ++i) {
            std::unique_ptr<Worker> wrk(new Worker(reactor));
            workers_.push_back(std::move(wrk));
        }
    }

    Reactor::Key addHandler(
            const std::shared_ptr<Handler>& handler, bool) {

        Reactor::Key keys[SyncImpl::MaxHandlers()];

        for (size_t i = 0; i < workers_.size(); ++i) {
            auto &wrk = workers_[i];

            auto cl = handler->clone();
            auto key = wrk->sync->addHandler(cl, false /* setKey */);
            auto newKey = encodeKey(key, i);
            cl->key_ = newKey;

            keys[i] = key;
        }

        auto data = keys[0].data() << 32 | KeyMarker;

        return Reactor::Key(data);
    }

    std::vector<std::shared_ptr<Handler>> handlers(const Reactor::Key& key) const {

        uint32_t idx;
        uint32_t marker;
        std::tie(idx, marker) = decodeKey(key);
        if (marker != KeyMarker)
            throw std::runtime_error("Invalid key");

        Reactor::Key originalKey(idx);

        std::vector<std::shared_ptr<Handler>> res;
        res.reserve(workers_.size());
        for (auto& wrk: workers_) {
            res.push_back(wrk->sync->handler(originalKey));
        }

        return res;
    }

    void registerFd(
            const Reactor::Key& key,
            Fd fd,
            Polling::NotifyOn interest,
            Polling::Tag tag,
            Polling::Mode mode = Polling::Mode::Level) {
        dispatchCall(key, &SyncImpl::registerFd, fd, interest, tag, mode);
    }

    void registerFdOneShot(
            const Reactor::Key& key,
            Fd fd,
            Polling::NotifyOn interest,
            Polling::Tag tag,
            Polling::Mode mode = Polling::Mode::Level) {
        dispatchCall(key, &SyncImpl::registerFdOneShot, fd, interest, tag, mode);
    }

    void modifyFd(
            const Reactor::Key& key,
            Fd fd,
            Polling::NotifyOn interest,
            Polling::Tag tag,
            Polling::Mode mode = Polling::Mode::Level) {
        dispatchCall(key, &SyncImpl::modifyFd, fd, interest, tag, mode);
    }

    void runOnce() {
    }

    void run() {
        for (auto& wrk: workers_)
            wrk->run();
    }

    void shutdown() {
        for (auto& wrk: workers_)
            wrk->shutdown();
    }

private:
    Reactor::Key encodeKey(const Reactor::Key& originalKey, uint32_t value) const
    {
        auto data = originalKey.data();
        auto newValue = data << 32 | value;
        return Reactor::Key(newValue);
    }

    std::pair<uint32_t, uint32_t> decodeKey(const Reactor::Key& encodedKey) const {
        auto data = encodedKey.data();
        uint32_t hi = data >> 32;
        uint32_t lo = data & 0xFFFFFFFF;
        return std::make_pair(hi, lo);
    }

#define CALL_MEMBER_FN(obj, pmf) (obj->*(pmf))

    template<typename Func, typename... Args>
    void dispatchCall(const Reactor::Key& key, Func func, Args&& ...args) const {
        auto decoded = decodeKey(key);
        auto& wrk = workers_[decoded.second];

        Reactor::Key originalKey(decoded.first);
        CALL_MEMBER_FN(wrk->sync.get(), func)(originalKey, std::forward<Args>(args)...);
    }

#undef CALL_MEMBER_FN

    struct Worker {

        Worker(Reactor* reactor) {
            sync.reset(new SyncImpl(reactor));
        }

        ~Worker() {
            if (thread)
                thread->join();
        }

        void run() {
            thread.reset(new std::thread([=]() {
                sync->run();
            }));
        }

        void shutdown() {
            sync->shutdown();
        }

        std::unique_ptr<std::thread> thread;
        std::unique_ptr<SyncImpl> sync;
    };

    std::vector<std::unique_ptr<Worker>> workers_;
};

Reactor::Key::Key()
    : data_(0)
{ }

Reactor::Key::Key(
        uint64_t data)
    : data_(data)
{ }

Reactor::Reactor() = default;

Reactor::~Reactor() = default;

std::shared_ptr<Reactor>
Reactor::create() {
    return std::make_shared<Reactor>();
}

void
Reactor::init() {
    SyncContext context;
    init(context);
}

void
Reactor::init(const ExecutionContext& context) {
    impl_.reset(context.makeImpl(this));
}

Reactor::Key
Reactor::addHandler(const std::shared_ptr<Handler>& handler) {
    return impl()->addHandler(handler, true);
}

std::vector<std::shared_ptr<Handler>>
Reactor::handlers(const Reactor::Key& key) {
    return impl()->handlers(key);
}

void
Reactor::registerFd(
        const Reactor::Key& key, Fd fd, Polling::NotifyOn interest, Polling::Tag tag,
        Polling::Mode mode)
{
    impl()->registerFd(key, fd, interest, tag, mode);
}

void
Reactor::registerFdOneShot(
        const Reactor::Key& key, Fd fd, Polling::NotifyOn interest, Polling::Tag tag,
        Polling::Mode mode)
{
    impl()->registerFdOneShot(key, fd, interest, tag, mode);
}

void
Reactor::registerFd(
        const Reactor::Key& key, Fd fd, Polling::NotifyOn interest,
        Polling::Mode mode)
{
    impl()->registerFd(key, fd, interest, Polling::Tag(fd), mode);
}

void
Reactor::registerFdOneShot(
        const Reactor::Key& key, Fd fd, Polling::NotifyOn interest,
        Polling::Mode mode)
{
    impl()->registerFdOneShot(key, fd, interest, Polling::Tag(fd), mode);
}

void
Reactor::modifyFd(
        const Reactor::Key& key, Fd fd, Polling::NotifyOn interest, Polling::Tag tag,
        Polling::Mode mode)
{
    impl()->modifyFd(key, fd, interest, tag, mode);
}

void
Reactor::modifyFd(
        const Reactor::Key& key, Fd fd, Polling::NotifyOn interest,
        Polling::Mode mode)
{
    impl()->modifyFd(key, fd, interest, Polling::Tag(fd), mode);
}

void
Reactor::run() {
    impl()->run();
}

void
Reactor::shutdown() {
    impl()->shutdown();
}

void
Reactor::runOnce() {
    impl()->runOnce();
}

Reactor::Impl *
Reactor::impl() const {
    if (!impl_)
        throw std::runtime_error("Invalid object state, you should call init() before");

    return impl_.get();
}

Reactor::Impl*
SyncContext::makeImpl(Reactor* reactor) const {
    return new SyncImpl(reactor);
}

Reactor::Impl*
AsyncContext::makeImpl(Reactor* reactor) const {
    return new AsyncImpl(reactor, threads_);
}

AsyncContext
AsyncContext::singleThreaded() {
    return AsyncContext(1);
}

} // namespace Aio
