/*
 * Copyright 2020 Bloomberg Finance LP
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#ifndef INCLUDED_BUILDBOXCASD_EXPIRINGUNORDEREDSET_H
#define INCLUDED_BUILDBOXCASD_EXPIRINGUNORDEREDSET_H

#include <buildboxcommon_logging.h>
#include <buildboxcommon_protos.h>

#include <chrono>
#include <condition_variable>
#include <deque>
#include <mutex>
#include <thread>
#include <unordered_map>

namespace buildboxcasd {

template <typename T> class ExpiringUnorderedSet {
  public:
    /*
     * Implements a thread-safe set that stores keys for a fixed number of
     * seconds after the time of their insertion.
     *
     * For performance reasons, the keys `k` that were inserted to the set but
     * has since expired return `false` for `contains(k)` but they are kept in
     * the underlying storage until `resize()` is called.
     *
     * That allows to make lookups efficient by keeping the complexity of
     * expiring keys out of the critical operations needed for implementing a
     * cache (`insert()` and `contains()`).
     */
    explicit ExpiringUnorderedSet(const int keySecondsToLive)
        : d_keySecondsToLive(keySecondsToLive)
    {
    }

    ~ExpiringUnorderedSet() = default;

    /*
     * Returns whether `key` is present in the set. (That is, `key` was
     * inserted less than `keyTimeToLive` seconds ago.)
     */
    bool contains(const T &key)
    {
        const std::lock_guard<std::mutex> lock(d_mutex);

        const auto it = d_data.find(key);
        return (it != d_data.cend()) ? !entryHasExpired(it->second) : false;
    }

    /*
     * Returns the number of elements currently stored in the underlying
     * storage, which can be larger than the actual number of keys `k` for
     * which `contains(k) == true`.
     *
     * That happens because expired keys are reported missing by `contains()`
     * but are not cleared from the storage until `cleanup()` is invoked.
     */
    size_t capacity()
    {
        const std::lock_guard<std::mutex> lock(d_mutex);

        assert(d_data.size() == d_queue.size());
        return d_data.size();
    };

    /*
     * Insert the given key to the set. Returns whether the set changed after
     * this operation (that is, `true` if the key was not already in the set
     * and `false` if it was.)
     *
     * If the key is already present, it has no effect on the lifetime of the
     * previous occurrence.
     *
     * Runs in constant time.
     */
    bool insert(const T &key)
    {
        const std::lock_guard<std::mutex> lock(d_mutex);

        const Timestamp creationTime = now();

        typename UnorderedMap::iterator it;
        bool entryCreated = false;
        std::tie(it, entryCreated) = d_data.emplace(key, creationTime);

        if (entryCreated) {
            d_queue.push_back(it);
        }

        return entryCreated;
    }

    /*
     * Time in seconds during which a key will report `contains(k) == true`
     * after being inserted.
     */
    int keySecondsToLive() const { return d_keySecondsToLive.count(); }

    /*
     * Remove expired entries from the underlying storage, potentially reducing
     * the value of `capacity()`. (Runs in linear time.)
     */
    void resize()
    {
        const std::lock_guard<std::mutex> lock(d_mutex);

        BUILDBOX_LOG_DEBUG("Removing expired entries. ("
                           << "lookup map size=" << d_data.size()
                           << ", queue size=" << d_queue.size() << ")");

        // Iterating the queue from older elements to newer ones.
        // We'll delete entries in `d_data` and `d_queue` until finding the
        // first element that has not expired yet.
        typename Queue::iterator queueIt = d_queue.begin();
        while (queueIt != d_queue.end()) {
            typename UnorderedMap::iterator dataIt = *queueIt;
            const Timestamp &entryCreationTime = dataIt->second;

            if (!entryHasExpired(entryCreationTime)) {
                break; // No more items to delete.
            }

            // Remove it from the map:
            d_data.erase(dataIt);
            // And get rid of the iterator:
            queueIt = d_queue.erase(queueIt);
        }

        BUILDBOX_LOG_DEBUG("Finished removing expired entries."
                           << "(lookup map size=" << d_data.size()
                           << ", queue size=" << d_queue.size() << ")");
    }

  private:
    std::mutex d_mutex;

    // Timing types:
    typedef std::chrono::system_clock::time_point Timestamp;
    typedef std::chrono::seconds Seconds;

    // Time that a key will be present until its expiry:
    Seconds d_keySecondsToLive;

    // We keep two separate structures:
    // i) An unordered map with entries of the form `(T, Timestamp)` that
    // is employed for lookups:
    typedef std::unordered_map<T, Timestamp> UnorderedMap;
    UnorderedMap d_data;
    // ii) An ordered container for efficient iteration during cleanup:
    typedef std::deque<typename UnorderedMap::iterator> Queue;
    Queue d_queue;

    /*
     * Return the current time.
     */
    static Timestamp now() { return std::chrono::system_clock::now(); }

    /*
     * Given a timestamp, compare it against the present time and return
     * whether the difference is larger than the configured time to live.
     */
    bool entryHasExpired(const Timestamp &creationTime) const
    {
        const auto entryAgeSeconds =
            std::chrono::duration_cast<std::chrono::seconds>(now() -
                                                             creationTime);

        return entryAgeSeconds > d_keySecondsToLive;
    }
};

template <typename T>
class SelfCleaningExpiringUnorderedSet final : public ExpiringUnorderedSet<T> {
  public:
    /*
     * Wraps an `ExpiringUnorderedSet` and spawns a thread that calls its
     * `resize()` method every a specified number of seconds.
     *
     */
    SelfCleaningExpiringUnorderedSet(const int keySecondsToLive,
                                     const int cleanupPeriodSeconds)
        : ExpiringUnorderedSet<T>(keySecondsToLive),
          d_cleanupPeriod(cleanupPeriodSeconds),
          d_cleanupThread(std::bind(
              &SelfCleaningExpiringUnorderedSet::cleanupThread, this))
    {
    }

    ~SelfCleaningExpiringUnorderedSet()
    {
        {
            const std::lock_guard<std::mutex> lock(d_stopRequestedMutex);
            if (d_stopRequested) {
                return;
            }
            d_stopRequested = true;
        }

        // Signaling `cleanupThread()` to wake up and exit:
        d_stopRequestedCondition.notify_one();
        d_cleanupThread.join();
    }

    // Disallowing copies:
    SelfCleaningExpiringUnorderedSet(
        const SelfCleaningExpiringUnorderedSet &) = delete;
    SelfCleaningExpiringUnorderedSet &
    operator=(SelfCleaningExpiringUnorderedSet const &) = delete;

    // Disallowing moves:
    SelfCleaningExpiringUnorderedSet(SelfCleaningExpiringUnorderedSet &&) =
        delete;
    SelfCleaningExpiringUnorderedSet &
    operator=(SelfCleaningExpiringUnorderedSet &&) = delete;

  private:
    const std::chrono::seconds d_cleanupPeriod;

    // A condition variable allows to wake up the thread from sleep when the
    // destructor is called:
    bool d_stopRequested = false;
    std::condition_variable d_stopRequestedCondition;
    std::mutex d_stopRequestedMutex;

    // Periodically call `resize()` on the container, removing keys that have
    // expired.
    std::thread d_cleanupThread;

    void cleanupThread()
    {
        while (true) {
            // Blocking while waiting `d_cleanupPeriodSeconds` to elapse or
            // `d_stopRequested` to be set...
            std::unique_lock<std::mutex> lock(d_stopRequestedMutex);
            const bool stopRequested = d_stopRequestedCondition.wait_for(
                lock, d_cleanupPeriod, [this]() { return d_stopRequested; });

            if (stopRequested) {
                return;
            }

            this->resize();
        }
    }
};

} // namespace buildboxcasd

#endif
