#ifndef CORE_SIMULATOR_H
#define CORE_SIMULATOR_H

#include <cstdint>
#include <functional>
#include <queue>
#include <vector>
#include <unordered_set>
#include <utility>
#include <tuple>
#include <type_traits>
#include <limits>
#include <memory>

#include "core/time.h"
#include "core/types.h"

namespace dofs {

class Rng;
class Logger;

class Node;
class NetworkNode;
class Link;

constexpr EventId NULL_EVENT = 0;

constexpr InstanceId INVALID_INSTANCE_ID =
        std::numeric_limits<InstanceId>::max();
constexpr LinkId     INVALID_LINK_ID     =
        std::numeric_limits<LinkId>::max();

class Simulator final {
private:
        struct Item {
                Time when;
                EventId id;
                std::function<void()> fn;
        };

        struct Cmp {
                bool operator()(const Item& a, const Item& b) const noexcept;
        };

        std::priority_queue<Item, std::vector<Item>, Cmp> _event_pq;
        std::unordered_set<EventId> _cancelled;
        Time    _now{};
        EventId _next_id{1};
        bool    _locked{false};

        std::unique_ptr<Rng> _rng;

        std::vector<std::unique_ptr<Node>> _nodes;
        std::vector<std::unique_ptr<Link>> _links;

public:
        Simulator() = default;

        static std::pair<InstanceId, Simulator *> create_simulator(InstanceId id);
        static Simulator *get_simulator(InstanceId id) noexcept;

        Time now() const noexcept;

        template <class F, class... Args>
        EventId schedule_at(Time abs_time, F&& f, Args&&... args) {
                if (_locked)
                        return NULL_EVENT;

                if (abs_time < _now)
                        return NULL_EVENT;

                const EventId eid = _next_id++;
                Item it{
                        abs_time,
                        eid,
                        make_callable(std::forward<F>(f), std::forward<Args>(args)...)
                };
                _event_pq.push(std::move(it));
                return eid;
        }

        template <class F, class... Args>
        EventId schedule_after(Time delay, F&& f, Args&&... args) {
                return schedule_at(_now + delay, std::forward<F>(f),
                                   std::forward<Args>(args)...);
        }

        bool cancel(EventId id);
        bool run_next();
        void run_until(Time end_time);

        // ----- Termination helpers -----
        // Prevent any future schedule_* calls from enqueuing events.
        void lock() noexcept;
        bool is_locked() const noexcept {
                return _locked;
        }
        // Lock, then run all events up to now()+grace, finally discard any remaining.
        void flush_after(Time grace) noexcept;

        // ---------- Object management ----------
        Rng* create_rng(std::uint64_t seed);
        Rng*       get_rng() noexcept;
        Rng const* get_rng() const noexcept;

        std::pair<LinkId, Link *> create_link(NetworkNode* a, PortId a_port,
                                              NetworkNode* b, PortId b_port,
                                              Time latency,
                                              double bandwidth_gbps);

        Link*       get_link(LinkId id) noexcept;
        Link const* get_link(LinkId id) const noexcept;

private:
        template <class F, class... Args>
        static auto make_callable(F&& f, Args&&... args) {
                using Fn  = std::decay_t<F>;
                using Tup = std::tuple<std::decay_t<Args>...>;
                return [fn = Fn(std::forward<F>(f)),
                tup = Tup(std::forward<Args>(args)...)]() mutable {
                        std::apply(fn, std::move(tup));
                };
        }

        static std::vector<std::unique_ptr<Simulator>>& _registry();
};

} // namespace dofs

#endif // CORE_SIMULATOR_H
