scheduler.hpp

#pragma once

#include <mutex>
#include <chrono>
#include <functional>
#include <thread>
#include <format>

#include <entt/entt.hpp>

#include "services.hpp"

namespace wge::literals {

  constexpr std::chrono::nanoseconds operator""_fps(unsigned long long d) noexcept
  {
    return std::chrono::duration_cast<std::chrono::nanoseconds>(std::chrono::seconds(1) / d);
  }

  constexpr std::chrono::nanoseconds operator""_fps(long double d) noexcept
  {
    return std::chrono::duration_cast<std::chrono::nanoseconds>(std::chrono::seconds(1) / d);
  }

  constexpr std::chrono::nanoseconds operator""_hz(unsigned long long d) noexcept
  {
    return std::chrono::duration_cast<std::chrono::nanoseconds>(std::chrono::seconds(1) / d);
  }

  constexpr std::chrono::nanoseconds operator""_hz(long double d) noexcept
  {
    return std::chrono::duration_cast<std::chrono::nanoseconds>(std::chrono::seconds(1) / d);
  }

} // namespace literals

namespace wge::scheduling {

  using clock = typename std::chrono::high_resolution_clock;
  using time = typename std::chrono::high_resolution_clock;
  using duration = clock::duration;
  using period = clock::period;
  using timestamp = clock::time_point;

  using schedule_id = entt::id_type;

  struct default_context
  {
    timestamp now = clock::now();
    bool is_late = false;
    duration frequency = std::chrono::seconds(1);
    double delta_time = 0.0;
    double elapsed_time = 0.0;

    uint64_t current_frame = 0;
    timestamp current_frame_start = clock::now();

    uint64_t expected_frame = 0;
    timestamp expected_frame_start = clock::now();

    timestamp next_frame_start = clock::now();
  };

  struct cancellable_context : public default_context
  {
    bool cancelled = false;

    void cancel()
    {
      this->cancelled = true;
    }
  };

  struct timer
  {
    struct context : public cancellable_context
    {
    };

    using handler = std::function<void()>;
    using handler_context = std::function<void(context &)>;

    duration interval;
    timestamp schedule;
    handler_context handle;
  };

  struct counter
  {
    struct context : public cancellable_context
    {
      uint64_t value = 0;
    };

    using handler = std::function<void()>;
    using handler_context = std::function<void(context &)>;

    uint32_t step;
    uint32_t count;
    duration interval;
    timestamp schedule;
    handler_context handle;
  };

  struct action
  {
    struct context : public default_context
    {
    };

    using handler = std::function<void()>;
    using handler_context = std::function<void(context &)>;

    timestamp schedule;
    handler_context handle;
  };

  struct frame
  {
    struct context : public cancellable_context
    {
    };

    using handler = std::function<void()>;
    using handler_context = std::function<void(context &)>;

    handler_context handle;
  };

  struct ticker
  {
    struct context : public cancellable_context
    {
    };

    using handler = std::function<void()>;
    using handler_context = std::function<void(context &)>;

    handler_context handle;
  };

  struct task
  {
    struct context
    {
    };

    std::function<void(context &)> handler;

  public:
    task(const std::function<void(context&)>&& handler)
    {
      this->handler = handler;
    }

    void operator()(context &ctx)
    {
    }
  };

  class scheduler
  {
  public:
    using service = typename wge::services::service<scheduler>;

  private:
    uint64_t _frame_index = 0;

    timestamp _now = clock::now();
    timestamp _start_time = timestamp();
    timestamp _next_tick = timestamp();

    duration _delta = std::chrono::seconds(0);

    duration _frequency = std::chrono::milliseconds(0); // Roughly 60fps

    std::recursive_mutex _mutex;
    entt::registry _registry;

  public:
    // TODO(BUG): changing the frequency while running will cause is_late to go crazy.

    void set_fps(double fps)
    {
      if (fps < 0.001) {
        this->_frequency = std::chrono::nanoseconds(0);
        return;
      }
      this->_frequency = std::chrono::duration_cast<std::chrono::nanoseconds>(std::chrono::duration<double>(1.0 / fps));
    }

    void set_frequency(duration frequency)
    {
      this->_frequency = frequency;
    }

    duration get_frequency() const
    {
      return this->_frequency;
    }

    void set_start(timestamp start)
    {
      this->_start_time = start;
    }

    timestamp get_start() const
    {
      return this->_start_time;
    }

    void set_start_now()
    {
      this->_start_time = clock::now();
    }

    void set_start_epoch()
    {
      this->_start_time = timestamp();
    }

    const timestamp get_now() const
    {
      return this->_now;
    }

    void set_now(timestamp now)
    {
      this->_now = now;
    }

    // Returns true if the simulation is running late
    const bool is_late() const
    {
      return this->get_current_frame_index() + 1 < this->get_expected_frame_index();
    }

    // Returns the current frame index
    const uint64_t get_current_frame_index() const
    {
      return this->_frame_index;
    }

    // Returns the expected frame index based on the current wall clock
    const uint64_t get_expected_frame_index() const
    {
      if (this->_frequency.count() == 0) {
        return this->_frame_index + 1;
      }
      return (this->_now - this->_start_time) / this->_frequency;
    }

    // Returns the expected frame's timestamp
    const timestamp get_expected_frame_timestamp() const
    {
      return this->_start_time + (this->get_expected_frame_index() * this->_frequency);
    }

    // Returns the current frame's timestamp
    const timestamp get_current_frame_timestamp() const
    {
      return this->_start_time + (this->_frame_index * this->_frequency);
    }

    // Returns the next frame's timestamp
    const timestamp get_next_frame_timestamp() const
    {
      return this->_start_time + ((this->_frame_index + 1) * this->_frequency);
    }

    const timestamp get_next_tick_time() const
    {
      return this->_next_tick;
    }

    // Gets the elapsed simulation time
    template <typename Duration>
    const Duration get_delta_time() const
    {
      return std::chrono::duration_cast<Duration>(this->_delta);
    }

    // Gets the elapsed simulation time
    const double get_delta_time_seconds() const
    {
      return std::chrono::duration<double>(
               this->get_delta_time<std::chrono::nanoseconds>() / std::chrono::duration<double>(1))
        .count();
    }

    // Gets the elapsed simulation time
    template <typename Duration>
    const Duration get_elapsed_time() const
    {
      return std::chrono::duration_cast<Duration>(this->_now - this->_start_time);
    }

    // Gets the elapsed simulation time
    const double get_elapsed_time_seconds() const
    {
      return std::chrono::duration<double>(
               this->get_elapsed_time<std::chrono::nanoseconds>() / std::chrono::duration<double>(1))
        .count();
    }

    // Schedules an action at the next tick
    auto run(action::handler_context &&handler)
    {
      std::scoped_lock lock(this->_mutex);
      auto entity = this->_registry.create();
      auto &component = this->_registry.emplace<action>(entity);
      component.handle = handler;
      component.schedule = this->_now;
      register_tick(component.schedule);
      return entity;
    }

    // Schedules an action at the next tick
    auto run(action::handler &&handler)
    {
      return this->run([handler = std::move(handler)](auto context) {
        // forward call
        handler();
      });
    }

    // Schedules an action at the next frame
    auto run_next_frame(action::handler_context &&handler)
    {
      std::scoped_lock lock(this->_mutex);
      auto entity = this->_registry.create();
      auto &component = this->_registry.emplace<action>(entity);
      component.handle = handler;
      component.schedule = this->get_next_frame_timestamp();
      register_tick(component.schedule);
      return entity;
    }

    auto run_next_frame(action::handler &&handler)
    {
      return this->run_next_frame([handler = std::move(handler)](auto context) {
        // forward call
        handler();
      });
    }

    // Schedules an action relative to the current timestamp
    template <class Rep, class Period>
    auto run_after(const std::chrono::duration<Rep, Period> &interval, action::handler_context &&handler)
    {
      std::scoped_lock lock(this->_mutex);
      auto entity = this->_registry.create();
      auto &component = this->_registry.emplace<action>(entity);
      component.handle = handler;
      component.schedule = this->get_current_frame_timestamp() + interval;
      register_tick(component.schedule);
      return entity;
    }

    template <class Rep, class Period>
    auto run_after(const std::chrono::duration<Rep, Period> &interval, action::handler &&handler)
    {
      return this->run_after<Rep, Period>(interval, [handler = std::move(handler)](auto context) {
        // forward call
        handler();
      });
    }

    // Schedules an action at a specific time point
    auto run_at(const timestamp &schedule, action::handler_context &&handler)
    {
      std::scoped_lock lock(this->_mutex);
      auto entity = this->_registry.create();
      auto &component = this->_registry.emplace<action>(entity);
      component.handle = handler;
      component.schedule = schedule;
      register_tick(component.schedule);
      return entity;
    }

    auto run_at(const timestamp &schedule, action::handler &&handler)
    {
      return this->run_at(schedule, [handler = std::move(handler)](auto context) {
        // forward call
        handler();
      });
    }

    // Schedules a continous timer
    template <class Rep, class Period>
    auto run_every(const std::chrono::duration<Rep, Period> &interval, timer::handler_context &&handler)
    {
      std::scoped_lock lock(this->_mutex);
      auto entity = this->_registry.create();
      auto &component = this->_registry.emplace<timer>(entity);
      component.interval = std::chrono::duration_cast<duration>(interval);
      component.schedule = this->get_current_frame_timestamp() + component.interval;
      component.handle = handler;
      register_tick(component.schedule);
      return entity;
    }

    // Schedules a continous timer
    template <class Rep, class Period>
    auto run_every(const std::chrono::duration<Rep, Period> &interval, timer::handler &&handler)
    {
      return this->run_every<Rep, Period>(interval, [handler = std::move(handler)](auto context) {
        // forward call
        handler();
      });
    }

    // Schedules a repeatable timer
    template <class Rep, class Period>
    auto run_counter(
      uint32_t count, const std::chrono::duration<Rep, Period> &interval, counter::handler_context &&handler)
    {
      std::scoped_lock lock(this->_mutex);
      auto entity = this->_registry.create();
      auto &component = this->_registry.emplace<counter>(entity);
      component.step = 0;
      component.count = count;
      component.interval = std::chrono::duration_cast<duration>(interval);
      component.schedule = this->_now;
      component.handle = handler;
      register_tick(component.schedule);
      return entity;
    }

    // Schedules a repeatable timer
    template <class Rep, class Period>
    auto run_counter(uint32_t count, const std::chrono::duration<Rep, Period> &interval, counter::handler &&handler)
    {
      return this->run_counter(count, interval, [handler = std::move(handler)](auto context) {
        // forward call
        handler();
      });
    }

    // Schedules a per frame call
    auto run_every_frame(frame::handler_context &&handler)
    {
      std::scoped_lock lock(this->_mutex);
      auto entity = this->_registry.create();
      auto &component = this->_registry.emplace<frame>(entity);
      component.handle = handler;
      return entity;
    }

    // Schedules a per frame call
    auto run_every_frame(frame::handler &&handler)
    {
      return this->run_every_frame([handler = std::move(handler)](auto context) {
        // forward call
        handler();
      });
    }

    // Schedules a repeatable timer
    auto run_every_tick(ticker::handler_context &&handler)
    {
      std::scoped_lock lock(this->_mutex);
      auto entity = this->_registry.create();
      auto &component = this->_registry.emplace<ticker>(entity);
      component.handle = handler;
      return entity;
    }

    // Schedules a repeatable timer
    auto run_every_tick(ticker::handler &&handler)
    {
      return this->run_every_tick([handler = std::move(handler)](auto context) {
        // forward call
        handler();
      });
    }

    timestamp advance_by(const duration& delta)
    {
      this->_delta = delta;
      this->_now += delta;
      return this->tick();
    }

    timestamp advance_to(const timestamp& now)
    {
      this->_delta = now - this->_now;
      this->_now = now;
      return this->tick();
    }

  protected:
    timestamp tick()
    {
      auto next_frame = this->get_next_frame_timestamp();
      this->_next_tick = next_frame;

      this->update_tickers();
      this->update_actions();
      this->update_counters();
      this->update_timers();

      if (next_frame <= this->_now)
      {
        this->update_frames();
        this->_frame_index++;
        next_frame = this->get_next_frame_timestamp();
      }

      if (this->_next_tick > next_frame)
      {
        this->_next_tick = next_frame;
      }

      return this->_next_tick;
    }

  protected:
    void register_tick(timestamp schedule)
    {
      if (schedule < this->_next_tick)
      {
        this->_next_tick = schedule;
      }
    }

    void update_timers()
    {
      std::scoped_lock lock(this->_mutex);

      auto entities = this->_registry.view<timer>().each();
      for (auto [id, component] : entities)
      {
        if (component.schedule + component.interval > this->_now)
        {
          register_tick(component.schedule + component.interval);
          continue;
        }

        auto context = this->make_context<timer::context>();

        component.handle(context);

        if (!context.cancelled)
        {
          component.schedule += component.interval;
          register_tick(component.schedule + component.interval);
          continue;
        }

        this->_registry.destroy(id);
      };
    }

    void update_counters()
    {
      std::scoped_lock lock(this->_mutex);

      auto entities = this->_registry.view<counter>().each();
      for (auto [id, component] : entities)
      {
        if (component.schedule + component.interval > this->_now)
        {
          register_tick(component.schedule + component.interval);
          continue;
        }

        if (component.step < component.count)
        {
          auto context = this->make_context<counter::context>();
          context.value = component.step;

          component.handle(context);
          component.step += 1;

          if (component.step < component.count)
          {
            if (!context.cancelled)
            {
              component.schedule += component.interval;
              register_tick(component.schedule + component.interval);
              continue;
            }
          }
        }

        this->_registry.destroy(id);
      }
    }

    void update_tickers()
    {
      std::scoped_lock lock(this->_mutex);

      auto entities = this->_registry.view<ticker>().each();
      for (auto [id, component] : entities)
      {
        auto context = this->make_context<ticker::context>();
        component.handle(context);
        if (context.cancelled)
        {
          this->_registry.destroy(id);
        }
      }
    }

    void update_actions()
    {
      std::scoped_lock lock(this->_mutex);

      auto entities = this->_registry.view<action>().each();
      for (auto [id, component] : entities)
      {
        if (component.schedule > this->_now)
        {
          register_tick(component.schedule);
          continue;
        }

        auto context = this->make_context<action::context>();
        component.handle(context);
        this->_registry.destroy(id);
      };
    }

    void update_frames()
    {
      std::scoped_lock lock(this->_mutex);

      auto entities = this->_registry.view<frame>().each();
      for (auto [id, component] : entities)
      {
        auto context = this->make_context<frame::context>();
        component.handle(context);
        if (context.cancelled)
        {
          this->_registry.destroy(id);
        }
      };
    }

  private:
    template<typename T>
    T make_context()
    {
      T context;
      context.now = this->get_now();
      context.is_late = this->is_late();
      context.frequency = this->get_frequency();
      context.delta_time = this->get_delta_time_seconds();
      context.elapsed_time = this->get_elapsed_time_seconds();
      context.current_frame = this->get_current_frame_index();
      context.current_frame_start = this->get_current_frame_timestamp();
      context.expected_frame = this->get_expected_frame_index();
      context.expected_frame_start = this->get_expected_frame_timestamp();
      context.next_frame_start = this->get_next_frame_timestamp();
      return context;
    }
  };

} // namespace wge