#pragma once

#include <cstdint>
#include <iterator> // IWYU pragma: keep
#include <map>
#include <set>

#include "sim/sim.h"
#include "util/assert.h"

namespace sim {

    template<typename Element> struct Queue {
        Scheduler &scheduler;

        Schedulable *reader = nullptr;
        std::set<Schedulable *> writers;

        unsigned int min_latency;

        std::multimap<std::uint64_t, Element> waiting;

        Queue(Scheduler &scheduler, unsigned int min_latency)
            : scheduler(scheduler)
            , min_latency(min_latency)
        { }

        void add_reader(Schedulable *x)
        {
            ASSERT(&x->scheduler == &scheduler, "Reader associated with incorrect Scheduler");
            ASSERT(!reader, "Already have a reader");
            reader = x;
            if (min_latency == 0) {
                for (const auto &w : writers)
                    scheduler.constrain(w, reader);
            }
        }

        void add_writer(Schedulable *x)
        {
            ASSERT(&x->scheduler == &scheduler, "Writer associated with incorrect Scheduler");
            writers.emplace(x);
            if (min_latency == 0 && reader)
                scheduler.constrain(x, reader);
        }

        void write(Element &&x)
        {
            write(std::move(x), min_latency);
        }

        void write(Element &&x, unsigned int latency)
        {
            ASSERT(latency >= min_latency, "Latency too low");
            ASSERT(!scheduler.current_schedulable || writers.count(scheduler.current_schedulable), "Write lacks permission");
            waiting.emplace(scheduler.now + min_latency, std::move(x));
        }

        unsigned int available()
        {
            ASSERT(!scheduler.current_schedulable || reader == scheduler.current_schedulable, "Read lacks permission");
            unsigned int c = 0;
            for (auto x = waiting.begin(); x != waiting.end() && x->first <= scheduler.now; ++x)
                ++c;
            return c;
        }

        Element &peek(unsigned int i=0)
        {
            ASSERT(i < available(), "Peek past available elements");
            auto x = waiting.begin();
            std::advance(x, i);
            return x->second;
        }

        Element read(unsigned int i=0)
        {
            ASSERT(i < available(), "Read past available elements");
            auto x = waiting.begin();
            std::advance(x, i);
            Element e{std::move(x->second)};
            waiting.erase(x);
            return e;
        }
    };

}