/*
 * SPDX-FileCopyrightText: Copyright (c) 2023-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: Apache-2.0
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#ifndef NVIDIA_GXF_STD_EVENT_BASED_SCHEDULER_HPP_
#define NVIDIA_GXF_STD_EVENT_BASED_SCHEDULER_HPP_

#include <algorithm>
#include <atomic>
#include <chrono>
#include <condition_variable>
#include <map>
#include <memory>
#include <mutex>
#include <set>
#include <shared_mutex>  // NOLINT
#include <string>
#include <thread>
#include <unordered_map>
#include <utility>
#include <vector>

#include "cuda.h"

#include "common/fixed_map.hpp"
#include "gxf/core/entity.hpp"
#include "gxf/core/handle.hpp"
#include "gxf/std/clock.hpp"
#include "gxf/std/cpu_thread.hpp"
#include "gxf/std/cuda_green_context_pool.hpp"
#include "gxf/std/gems/event_list/unique_event_list.hpp"
#include "gxf/std/gems/staging_queue/staging_queue.hpp"
#include "gxf/std/gems/timed_job_list/timed_job_list.hpp"
#include "gxf/std/resources.hpp"
#include "gxf/std/scheduler.hpp"
#include "gxf/std/scheduling_condition.hpp"

namespace nvidia {
namespace gxf {

// Forward declarations
class EntityExecutor;

/// @brief Event based scheduler
///
/// The scheduler requires a Clock to keep track of time. It spawns a dispatcher thread and several
/// worker threads as configured. Dispatcher thread checks condition of entities and place them in
/// worker queue if they are ready to run or have expected time to run.
class EventBasedScheduler : public Scheduler {
 public:
  gxf_result_t registerInterface(Registrar* registrar) override;
  gxf_result_t initialize() override;
  gxf_result_t deinitialize() override;
  gxf_result_t prepare_abi(EntityExecutor* executor) override;
  gxf_result_t schedule_abi(gxf_uid_t eid) override;
  gxf_result_t unschedule_abi(gxf_uid_t eid) override;
  gxf_result_t runAsync_abi() override;
  gxf_result_t stop_abi() override;
  gxf_result_t wait_abi() override;
  gxf_result_t event_notify_abi(gxf_uid_t eid, gxf_event_t event) override;
  gxf_result_t notifyDispatcher(gxf_uid_t eid = kNullUid);

 private:
  // Enum to indicate if a worker thread has acquired an entity for execution or not
  enum class EntityOwnership : uint8_t {
    kFree = 0,
    kAcquired
  };

  class ScheduleEntity {
   public:
    ScheduleEntity(gxf_uid_t eid, const char * name) {
      eid_ = eid;
      name_ = std::string(name);
      condition_ = SchedulingCondition{SchedulingConditionType::READY, 0};
      ownership_ = EntityOwnership::kFree;
      queue_index_ = -1;
      thread_pool_id_ = -1;
      thread_id_ = kDefaultThreadPoolThreadId;
      unschedule_ = false;
      // is_present_in_ready_queue_ initialized via member initializer (atomic<bool>{false})
      cuda_context_ = nullptr;
      cuda_green_context_pool_ = nullptr;
    }

    void updateCondition() {
      condition_ = SchedulingCondition{SchedulingConditionType::READY, 0};
    }

    Expected<void> tryToAcquire(void) {
      EntityOwnership free = EntityOwnership::kFree;
      if (ownership_.compare_exchange_strong(free, EntityOwnership::kAcquired)) {
        return Success;
      }
      return Unexpected{GXF_FAILURE};
    }

    Expected<void> releaseOwnership(void) {
      EntityOwnership acquired = EntityOwnership::kAcquired;
      if (ownership_.compare_exchange_strong(acquired, EntityOwnership::kFree)) {
        return Success;
      }
      return Unexpected{GXF_FAILURE};
    }

    // Atomically try to mark entity as present in ready queue.
    // Returns true if successfully marked (was not in queue), false if already in queue.
    bool tryMarkInReadyQueue() {
      bool expected = false;
      return is_present_in_ready_queue_.compare_exchange_strong(
          expected, true, std::memory_order_acq_rel);
    }

    // Atomically mark entity as not present in ready queue.
    void markNotInReadyQueue() {
      is_present_in_ready_queue_.store(false, std::memory_order_release);
    }

    // Check if entity is currently in the ready queue.
    bool isInReadyQueue() const {
      return is_present_in_ready_queue_.load(std::memory_order_acquire);
    }

    // Atomically try to mark entity as NOT in ready queue (for removal).
    // Returns true if successfully marked (was in queue), false if not in queue.
    bool tryMarkNotInReadyQueue() {
      bool expected = true;
      return is_present_in_ready_queue_.compare_exchange_strong(
          expected, false, std::memory_order_acq_rel);
    }

    int32_t queue_index_;
    gxf_uid_t eid_;
    SchedulingCondition condition_;
    std::string name_;
    std::atomic<EntityOwnership> ownership_;
    int32_t thread_pool_id_;
    gxf_uid_t thread_id_;
    bool unschedule_;
    std::atomic<bool> is_present_in_ready_queue_{false};
    CUcontext cuda_context_ = nullptr;
    CudaGreenContextPool* cuda_green_context_pool_ = nullptr;
  };

  std::unordered_map<gxf_uid_t, std::shared_ptr<ScheduleEntity>> entities_;

  struct TimeStat {
    std::atomic<uint64_t> execTime = {0};
    std::atomic<uint64_t> waitTime = {0};
    std::atomic<uint64_t> execCount = {0};
  };

  struct TimeStat dispatcherStats;
  struct TimeStat workerStats;           // operations from multiple worker threads

  struct PerfCounters {
    std::atomic<uint64_t> notify_calls{0};
    std::atomic<uint64_t> notify_calls_with_entity{0};
    std::atomic<uint64_t> notify_skipped_never{0};
    std::atomic<uint64_t> notify_inserted{0};
    std::atomic<uint64_t> notify_duplicates{0};
    std::atomic<uint64_t> notify_signals_sent{0};
    std::atomic<uint64_t> dispatcher_loop_iterations{0};
    std::atomic<uint64_t> dispatcher_events_popped{0};
    std::atomic<uint64_t> dispatcher_events_dispatched{0};
    std::atomic<uint64_t> dispatcher_busy_or_readyqueue_skips{0};
    std::atomic<uint64_t> worker_waitforjob_calls{0};
    std::atomic<uint64_t> worker_waitforjob_total_us{0};
    std::atomic<uint64_t> worker_waitforjob_wakeups{0};
    std::atomic<uint64_t> worker_execute_calls{0};
    std::atomic<uint64_t> worker_execute_total_us{0};
    std::atomic<uint64_t> steal_attempts{0};
    std::atomic<uint64_t> steal_successes{0};
    std::atomic<uint64_t> worker_postcheck_calls{0};
    std::atomic<uint64_t> worker_postcheck_fastpath_ready{0};
    std::atomic<uint64_t> worker_postcheck_fastpath_wait_time{0};
    std::atomic<uint64_t> worker_postcheck_fallback_notify{0};
    std::atomic<uint64_t> worker_postcheck_fallback_skip_notify{0};
    std::atomic<uint64_t> worker_postcheck_fallback_periodic_notify{0};
  };

  PerfCounters perf_counters_;

  static constexpr int64_t kMsToNs = 1'000'000l;      // Convenient constant of 1 ms = 1e6 ns
  static constexpr int64_t kMaxSlipNs = 1 * kMsToNs;  // Max slip tolerance set to 1 ms
  static constexpr size_t kMaxThreads = 128;          // Max number of worker threads

  // Common thread id for all the worker threads in default thread pool
  static constexpr int64_t kDefaultThreadPoolThreadId = -1;
  // Entrance for dispatcher threads
  void dispatcherThreadEntrance();
  // notify the dispatcher thread
  // Move entity to some queue
  void dispatchEntity(std::shared_ptr<ScheduleEntity> e);
  void dispatchEntityAsync(std::shared_ptr<ScheduleEntity> e);
  // Entrance for async event handler thread
  void asyncEventThreadEntrance();
  // Entrance for worker threads
  void workerThreadEntrance(ThreadPool* pool, int64_t thread_number,
                            gxf_uid_t cpu_thread_cid = kUnspecifiedUid);

  // Checks if need to stop due to no active entities or expiration
  bool checkEndingCriteria(int64_t timestamp);

  // Updates condition of specific entity id and keep track of how many entities are good to run
  void updateCondition(std::shared_ptr<ScheduleEntity> e,
   const SchedulingCondition& next_condition);

  // stops all async threads and deactivates all the entities
  gxf_result_t stopAllThreads();

  // stops all jobs queued in dispatcher and workers
  void stopAllJobs();

  // Logs optional phase-0 instrumentation counters.
  void logPerfStats();
  // Pushes an internal notification into a shard queue; returns true on insert.
  bool pushInternalNotification(gxf_uid_t eid);
  // Removes an internal notification from the owning shard queue.
  void removeInternalNotification(gxf_uid_t eid);
  // Pops one internal notification from shard queues (round-robin across shards).
  bool popInternalNotification(gxf_uid_t& eid);
  // Pops up to max_count internal notifications from one shard queue.
  bool popInternalNotifications(std::vector<gxf_uid_t>& eids, size_t max_count);
  // Returns true if any internal shard queue has pending notifications.
  bool hasInternalNotifications() const;
  // Returns total pending notifications across internal shard queues.
  uint64_t internalNotificationCount() const;
  // WAIT_EVENT shard helpers.
  void addEventWaiting(gxf_uid_t eid);
  void removeEventWaiting(gxf_uid_t eid);
  uint64_t eventWaitingCount() const;
  // WAIT shard helpers.
  void addWaiting(gxf_uid_t eid);
  void removeWaiting(gxf_uid_t eid);
  uint64_t waitingCount() const;

  // returns the current state of the scheduler as string
  const char* schedulerStateString();
  // cache thread info for pinned job into resources_ and set cuda context if provided
  gxf_result_t prepareResourceMapStrict(std::shared_ptr<ScheduleEntity> e);
  // index of the ready queue in the ready_wait_time_jobs_ vector where entity is supposed to go

  uint64_t getReadyCount();

  // Helper functions for configuring worker threads
  gxf_result_t configureThread(gxf_uid_t thread_uid, gxf_uid_t cpu_thread_cid);
  gxf_result_t setCPUAffinity(const std::vector<uint32_t>& cpu_cores);
  bool isRealtimeCPUThread(gxf_uid_t cpu_thread_cid);

  // Load CUDA Driver API functions through runtime API
  gxf_result_t loadDriverFunctions();

  // Parameters
  Parameter<Handle<Clock>> clock_;
  Parameter<int64_t> max_duration_ms_;
  Parameter<bool> stop_on_deadlock_;
  Parameter<int64_t> stop_on_deadlock_timeout_;
  Parameter<int64_t> network_connection_timeout_;
  Parameter<int64_t> worker_thread_number_;
  Parameter<bool> thread_pool_allocation_auto_;
  Parameter<std::string> worker_thread_name_id_;
  // CPU core IDs to pin the worker threads to (empty means no core pinning)
  Parameter<std::vector<uint32_t>> pin_cores_;
  // Enables worker-side queue stealing for default worker queues.
  Parameter<bool> enable_queue_stealing_;
  // Maximum number of victim queues scanned per steal attempt (0 = scan all).
  Parameter<int64_t> steal_scan_limit_;
  // Non-pinned entity assignment mode for default worker queues:
  // true = round-robin, false = hash(eid)-based.
  Parameter<bool> default_queue_assign_round_robin_;
  // If true, workers run a fresh checkEntity() after executeEntity() and can update scheduling
  // condition locally to reduce dispatcher round-trips.
  Parameter<bool> enable_worker_postcheck_fastpath_;
  // Periodic dispatcher wake interval for non-ready worker postcheck fallback path.
  // 0 disables periodic wake-ups (only notify when no workers are running).
  Parameter<int64_t> postcheck_fallback_notify_interval_;
  // Enable periodic fallback dispatcher wake-ups only when worker_thread_number is at least this
  // value.
  Parameter<int64_t> postcheck_fallback_notify_min_workers_;
  // Minimum global time spacing (ns) between periodic fallback dispatcher wake-ups.
  // 0 disables global time-based throttling.
  Parameter<int64_t> postcheck_fallback_notify_min_period_ns_;
  // Number of internal notification shards (0 = auto = worker_thread_number).
  Parameter<int64_t> internal_event_shard_count_;
  // Maximum internal notifications drained per dispatcher pop batch (>=1).
  Parameter<int64_t> dispatcher_internal_pop_batch_size_;
  // Number of shards for WAIT_EVENT / WAIT tracking lists (default 1).
  Parameter<int64_t> wait_state_shard_count_;
  // Emits EventBasedScheduler phase-0 instrumentation counters at deinitialize() when true.
  Parameter<bool> log_perf_stats_;

  EntityExecutor* executor_ = nullptr;

  // Track when scheduler started to enforce minimum grace period for network connections
  int64_t scheduler_start_timestamp_ = -1;

  // Cache for network connection readiness (once true, stays true)
  bool network_connections_established_ = false;

  Entity default_thread_pool_entity_;
  ThreadPool* default_thread_pool_;

  // thread pool set including default pool and added pools
  std::set<ThreadPool*> thread_pool_set_;

  // map for storing the relevant ready_wait_queue
  FixedMap<gxf_uid_t, int> thread_queue_mapping_;
  // Queue indices owned by default worker threads, used for non-pinned entity assignment.
  std::vector<int> default_queue_indices_;
  // Round-robin cursor for assigning non-pinned entities across default worker queues.
  std::atomic<size_t> next_default_queue_rr_{0};

  gxf_result_t thread_error_code_;

  // A thread to dispatch jobs to worker pool
  std::thread dispatcher_thread_;

  // A thread to trigger exit after max time duration
  std::thread max_duration_thread_;

  // Mutex to synchronize dispatcher thread
  std::mutex dispatcher_sync_mutex_;

  // Mutex to synchronize max duration thread
  std::mutex max_duration_sync_mutex_;

  // instances of dispatcher/worker threads. 0 index would be async event handler thread.
  std::vector<std::thread> async_threads_;

  // Keep track of the timestamp that the scheduler started working for checking expiration
  int64_t start_timestamp_ = 0;

  // Queue with execution time for jobs to execute
  std::vector<std::unique_ptr<TimedJobList<gxf_uid_t>>> ready_wait_time_jobs_;

  // synchronize entities
  std::mutex external_event_notification_mutex_;
  std::mutex internal_event_notification_mutex_;

  // Mutex to synchronize worker thread pool
  std::mutex thread_sync_mutex_;
  std::condition_variable thread_sync_cv_;

  // Keep track of asynchronous external events
  std::unique_ptr<UniqueEventList<gxf_uid_t>> external_event_notified_;
  // Holds entities waiting for some event (sharded).
  std::vector<std::unique_ptr<UniqueEventList<gxf_uid_t>>> event_waiting_shards_;
  // Holds entities waiting for custom conditions (sharded).
  std::vector<std::unique_ptr<UniqueEventList<gxf_uid_t>>> waiting_shards_;

  // Sharded internal event notifications to reduce notify-path lock contention.
  std::vector<std::unique_ptr<UniqueEventList<gxf_uid_t>>> internal_event_notified_shards_;
  size_t next_internal_event_shard_rr_{0};

  // Running count of entities in flight / execution
  std::atomic<uint8_t> running_threads_{0};

  // Flag indicating if dispatcher thread is waiting on condition variable.
  // Used to skip unnecessary cv.notify_one() syscalls when dispatcher is busy.
  std::atomic<bool> dispatcher_waiting_{false};

  std::condition_variable external_event_notification_cv_;
  std::condition_variable internal_event_notification_cv_;

  std::condition_variable max_duration_thread_cv_;

  // enum type to represent scheduler state changes
  enum class State : uint8_t {
    kNotStarted = 0,       // Scheduler has not started execution yet, starts during runAsync_abi()
    kRunning,              // Execution threads are running
    kStopping,             // Ending criteria reached, execution threads are doing pending jobs
    kStopped,              // Scheduler has stopped
  };

  std::atomic<State> state_{State::kNotStarted};
  // Global throttle for periodic fallback dispatcher wake-ups.
  std::atomic<int64_t> next_postcheck_periodic_notify_ts_ns_{0};
  std::mutex state_change_mutex_;
  std::condition_variable work_done_cv_;

  // latest timestamp from last should_stop == false
  int64_t last_no_stop_ts_ = 0;
  // maintain last no stop timestamp, and check if need to update should_stop
  gxf_result_t stop_on_deadlock_timeout(const int64_t timeout,
    const int64_t now, bool& should_stop);

  // Deadlock detection grace period end timestamp
  int64_t grace_end_timestamp_ = 0;
  // Armed flag to require two consecutive deadlock checks before exiting
  bool deadlock_armed_ = false;

  // CUDA Driver API
  CUresult (*fnCuCtxPushCurrent)(CUcontext) = nullptr;
  CUresult (*fnCuCtxPopCurrent)(CUcontext*) = nullptr;
};

}  // namespace gxf
}  // namespace nvidia

#endif  // NVIDIA_GXF_STD_EVENT_BASED_SCHEDULER_HPP_
