#ifndef THIRD_PARTY_MILOTIC_EXTERNAL_CC_TLBMC_SCHEDULER_SCHEDULER_H_

#define THIRD_PARTY_MILOTIC_EXTERNAL_CC_TLBMC_SCHEDULER_SCHEDULER_H_

#include <atomic>

#include <cstddef>

#include <cstdint>

#include <memory>

#include <optional>

#include <thread>  //NOLINT: <thread> is commonly used in BMC

#include "absl/base/thread_annotations.h"

#include "absl/container/flat_hash_map.h"

#include "absl/functional/any_invocable.h"

#include "absl/synchronization/mutex.h"

#include "absl/time/time.h"

#include "boost/asio/io_context.hpp"  //NOLINT: boost is commonly used in BMC

#include "boost/asio/steady_timer.hpp"  //NOLINT: boost is commonly used in BMC

#include <boost/asio/executor_work_guard.hpp>

#include "boost/circular_buffer.hpp"  //NOLINT: boost is commonly used in BMC

#include "time/clock.h"

#include <nlohmann/json.hpp>

namespace milotic_tlbmc {

// Represents the time information of a task.

struct TimeInfo {

  absl::Time last_scheduled_time = absl::InfiniteFuture();

  absl::Time last_run_time = absl::InfiniteFuture();

  boost::circular_buffer<absl::Duration> wait_times =

      boost::circular_buffer<absl::Duration>(1000);

  nlohmann::json ToJson() const;

};

// Represents the mode of a task.

//

// kPeriodic: The task will be scheduled periodically.

// kOnce: The task will be cancelled after it is executed once.

enum class TaskMode : uint8_t {

  kPeriodic = 0,

  kOnce,

};

// Represents a single task to be scheduled and run.

class Task : public std::enable_shared_from_this<Task> {

 public:

  Task(int task_id,

       absl::AnyInvocable<void(absl::AnyInvocable<void()>)> task_fn,

       absl::Duration period, boost::asio::io_context* io,

       absl::Duration task_timeout, TaskMode task_mode = TaskMode::kPeriodic,

       ecclesia::Clock* clock = ecclesia::Clock::RealClock());

  Task(const Task&) = delete;

  Task& operator=(const Task&) = delete;

  void Run() ABSL_LOCKS_EXCLUDED(task_mutex_, time_info_mutex_);

  void ScheduleNextRun() ABSL_LOCKS_EXCLUDED(timer_mutex_);

  absl::Duration GetPeriod() const;

  void SetPeriod(absl::Duration period);

  void Cancel() ABSL_LOCKS_EXCLUDED(timer_mutex_);

  void OnDone() ABSL_LOCKS_EXCLUDED(timer_mutex_, time_info_mutex_);

  bool IsCancelled() const { return cancelled_.load(); }

  int GetExecutionCount() const { return execution_count_.load(); }

  absl::Time GetLastRunTime() const ABSL_LOCKS_EXCLUDED(time_info_mutex_);

  absl::Duration GetAverageWaitTime() const

      ABSL_LOCKS_EXCLUDED(time_info_mutex_);

  absl::Duration GetTaskTimeout() const { return task_timeout_; }

  nlohmann::json ToJson() const;

 private:

  const int task_id_;

  const TaskMode task_mode_;

  // Ecclesia clock is thread safe.

  ecclesia::Clock* clock_ = nullptr;

  std::atomic<absl::Duration> period_ = absl::InfiniteDuration();

  std::atomic<bool> cancelled_ = false;

  std::atomic<int> execution_count_ = 0;

  std::atomic<absl::Duration> last_execution_duration_ = absl::ZeroDuration();

  // Task function to be executed.

  // task_fn_with_ack_: This is used for async tasks i.e tasks with

  // TaskMode kPeriodic that have an ack callback.

  absl::AnyInvocable<void(absl::AnyInvocable<void()>)> task_fn_with_ack_

      ABSL_GUARDED_BY(task_mutex_) = nullptr;

  mutable absl::Mutex task_mutex_;

  // TimeInfo to capture the last, next and average wait time.

  TimeInfo time_info_ ABSL_GUARDED_BY(time_info_mutex_);

  mutable absl::Mutex time_info_mutex_;

  // Timer to schedule the task.

  // Mainly used for async tasks.

  std::unique_ptr<boost::asio::steady_timer> timer_

      ABSL_GUARDED_BY(timer_mutex_) = nullptr;

  absl::Mutex timer_mutex_;

  const absl::Duration task_timeout_;

};

// Schedules tasks with a given periodicity.

//

// The scheduler is responsible for creating a thread for each period and

// scheduling the tasks to run on the corresponding thread based on the period.

// The scheduler also provides a way to reschedule a task with a new period and

// cancel a task.

//

// The scheduler is thread safe.

class TaskScheduler {

 public:

  // Options to configure the scheduler.

  struct Options {

    // Default options for the scheduler.

    //

    // Cleanup tasks period is set to 10 seconds arbitrarily, not shorter to

    // prevent overloading scheduler for cleanup tasks.

    //

    // Default task timeout is set to 60 seconds to allow for long running

    // tasks to complete which today aligns with the default timeout for

    // Redfish requests on client side.

    Options()

        : cleanup_tasks_period(absl::Seconds(10)),

          default_task_timeout(absl::Seconds(60)) {}

    absl::Duration cleanup_tasks_period;

    absl::Duration default_task_timeout;

  };

  explicit TaskScheduler(ecclesia::Clock* clock = ecclesia::Clock::RealClock(),

                         const Options& options = Options());

  TaskScheduler(const TaskScheduler&) = delete;

  TaskScheduler& operator=(const TaskScheduler&) = delete;

  virtual ~TaskScheduler();

  // Schedules the `task` for async invocation with the given `period`.

  //

  // The `task` will be invoked with a `on_done` callback as a parameter.

  // The `on_done` callback will be invoked when the task is done.

  // The `timeout` parameter is the timeout for the task to complete. If not

  // specified, the infinite timeout is used. If a task is not completed within

  // the timeout, the task will be cancelled.

  int RunAndScheduleAsync(

      absl::AnyInvocable<void(absl::AnyInvocable<void()>)> task,

      absl::Duration period,

      std::optional<absl::Duration> timeout = absl::InfiniteDuration());

  // Same as RunAndScheduleAsync but does not run the task immediately before

  // scheduling it.

  int ScheduleAsync(

      absl::AnyInvocable<void(absl::AnyInvocable<void()>)> task,

      absl::Duration period,

      std::optional<absl::Duration> timeout = absl::InfiniteDuration());

  // Schedule a oneshot task that is async.

  int ScheduleOneShotAsync(

      absl::AnyInvocable<void(absl::AnyInvocable<void()>)> task,

      absl::Duration period,

      std::optional<absl::Duration> timeout = std::nullopt);

  void UpdateTaskPeriod(int task_id, absl::Duration period)

      ABSL_LOCKS_EXCLUDED(async_tasks_mutex_);

  void Cancel(int task_id) ABSL_LOCKS_EXCLUDED(async_tasks_mutex_);

  void CancelAll() ABSL_LOCKS_EXCLUDED(async_tasks_mutex_);

  void Stop() ABSL_LOCKS_EXCLUDED(async_tasks_mutex_);

  size_t GetSchedulingAccuracyPercentage() const;

  nlohmann::json ToJson() const;

  // This only counts the number of user scheduled async tasks.

  size_t GetAllUserTaskCount() const ABSL_LOCKS_EXCLUDED(async_tasks_mutex_);

  // This counts the number of all tasks including the internal scheduler tasks.

  size_t GetAllTaskCount() const ABSL_LOCKS_EXCLUDED(async_tasks_mutex_);

 protected:

  virtual absl::flat_hash_map<int, std::shared_ptr<Task>> GetSnapshot() const

      ABSL_LOCKS_EXCLUDED(async_tasks_mutex_);

  void CleanupInactiveTasks() ABSL_LOCKS_EXCLUDED(async_tasks_mutex_);

 private:

  enum class ExecutionMode : uint8_t {

    kRunImmediately = 0,

    kRunAfterScheduling,

  };

  int ScheduleAsyncInternal(

      absl::AnyInvocable<void(absl::AnyInvocable<void()>)> task,

      absl::Duration period, std::optional<absl::Duration> timeout,

      ExecutionMode execution_mode);

  struct AliveToken {

    absl::Mutex mu;

    bool alive ABSL_GUARDED_BY(mu) = true;

  };

  std::shared_ptr<AliveToken> alive_token_ = std::make_shared<AliveToken>();

  std::atomic<int> next_task_id_;

  ecclesia::Clock* clock_;

  std::unique_ptr<boost::asio::io_context> io_;

  boost::asio::executor_work_guard<boost::asio::io_context::executor_type>

      work_guard_;

  std::thread async_executor_thread_;

  mutable absl::Mutex thread_join_mutex_;

  mutable absl::Mutex async_tasks_mutex_;

  absl::flat_hash_map<int, std::shared_ptr<Task>> id_to_async_tasks_

      ABSL_GUARDED_BY(async_tasks_mutex_);

  std::atomic<bool> stop_ = false;

  std::atomic<int> cleanup_task_id_ = -1;

  absl::Duration default_task_timeout_;

};

}  // namespace milotic_tlbmc

#endif  // THIRD_PARTY_MILOTIC_EXTERNAL_CC_TLBMC_SCHEDULER_SCHEDULER_H_