#include <utility>

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

namespace sim {

    void Scheduler::add_schedulable(Schedulable *schedulable)
    {
        unsorted_schedulables.emplace(schedulable);
        sort_needed = true;
    }

    void Scheduler::remove_schedulable(Schedulable *schedulable)
    {
        unsorted_schedulables.erase(schedulable);
        std::erase_if(constraints, [schedulable](const auto &item) { return item.first == schedulable || item.second == schedulable; });
        sort_needed = true;
    }

    void Scheduler::constrain(Schedulable *prior, Schedulable *later)
    {
        ASSERT(unsorted_schedulables.count(prior), "Constraint prior is not associated with this Scheduler");
        ASSERT(unsorted_schedulables.count(later), "Constraint later is not associated with this Scheduler");
        constraints.emplace(later, prior);
        sort_needed = true;
    }

    void Scheduler::topo_sort(std::set<Schedulable *> &live, std::set<Schedulable *> &waiting, Schedulable *candidate)
    {
        ASSERT(!live.count(candidate), "Dependency loop");
        for (auto prereq = constraints.find(candidate); prereq != constraints.end() && prereq->first == candidate; ++prereq) {
            if (prereq->second != candidate && waiting.count(prereq->second)) {
                live.emplace(candidate);
                topo_sort(live, waiting, prereq->second);
            }
        }
        sorted_schedulables.emplace_back(candidate);
        waiting.erase(candidate);
    }

    void Scheduler::sort()
    {
        sorted_schedulables.clear();
        sorted_schedulables.reserve(unsorted_schedulables.size());
        auto waiting = unsorted_schedulables;
        while (!waiting.empty()) {
            std::set<Schedulable *> live;
            topo_sort(live, waiting, *waiting.begin());
        }
        ASSERT(sorted_schedulables.size() == unsorted_schedulables.size(), "Did not sort every schedulable");
        sort_needed = false;
    }

    void Scheduler::clock()
    {
        if (sort_needed)
            sort();
        for (const auto &s : sorted_schedulables) {
            current_schedulable = s;
            s->clock();
        }
        current_schedulable = nullptr;
        ++now;
    }

    Schedulable::Schedulable(Scheduler &scheduler)
        : scheduler(scheduler)
    {
        scheduler.add_schedulable(this);
    }

    Schedulable::~Schedulable()
    {
        scheduler.remove_schedulable(this);
    }

    std::uint64_t Schedulable::now()
    {
        return scheduler.now;
    }

}