/*
Copyright (c) 2020-2026, NVIDIA CORPORATION. All rights reserved.

NVIDIA CORPORATION and its licensors retain all intellectual property
and proprietary rights in and to this software, related documentation
and any modifications thereto. Any use, reproduction, disclosure or
distribution of this software and related documentation without an express
license agreement from NVIDIA CORPORATION is strictly prohibited.
*/
#ifndef NVIDIA_GXF_STD_ENTITY_EXECUTOR_HPP_
#define NVIDIA_GXF_STD_ENTITY_EXECUTOR_HPP_

#include <algorithm>
#include <atomic>
#include <chrono>
#include <map>
#include <memory>
#include <mutex>
#include <shared_mutex>  // NOLINT
#include <thread>
#include <utility>

#include "nvtx3/nvToolsExt.h"

#include "common/fixed_vector.hpp"
#include "common/logger.hpp"
#include "gxf/core/component.hpp"
#include "gxf/core/entity.hpp"
#include "gxf/core/handle.hpp"
#include "gxf/std/clock.hpp"
#include "gxf/std/codelet.hpp"
#include "gxf/std/controller.hpp"
#include "gxf/std/job_statistics.hpp"
#include "gxf/std/message_router.hpp"
#include "gxf/std/monitor.hpp"
#include "gxf/std/network_router.hpp"
#include "gxf/std/router.hpp"
#include "gxf/std/scheduling_term_combiner.hpp"
#include "gxf/std/scheduling_terms.hpp"
#include "gxf/std/system.hpp"

namespace nvidia {
namespace gxf {

#define GXF_STR_HELP(X) \
  case X:               \
    return #X;

// Executes entities.
class EntityExecutor {
 public:
  gxf_result_t initialize(Handle<Router> router,
                          Handle<MessageRouter> message_router,
                          Handle<NetworkRouter> network_router);

  gxf_result_t activate(gxf_context_t context, gxf_uid_t eid);

  gxf_result_t deactivate(gxf_uid_t eid);

  gxf_result_t deactivateAll();

  Expected<void> getEntities(FixedVectorBase<gxf_uid_t>& entities) const;

  // sets the clock in the router to be used to set pubtimestamps
  gxf_result_t setClock(Handle<Clock> clock);

  // Check if entity is ready to run and run it if possible. Returns scheduling conditions of the
  // entity anyway.
  Expected<SchedulingCondition> executeEntity(gxf_uid_t eid, int64_t timestamp);

  // Checks if entity is ready to run but not run it.
  Expected<SchedulingCondition> checkEntity(gxf_uid_t eid, int64_t timestamp);

  // Provides Statistics component to be updated when running codelets.
  Expected<void> addStatistics(Handle<JobStatistics> statistics);

  // Removes a statistics component from the executor
  Expected<void> removeStatistics(Handle<JobStatistics> statistics);

  // Adds a monitor to the executor
  Expected<void> addMonitor(Handle<Monitor> monitor);

  // Removes a monitor from the executor
  Expected<void> removeMonitor(Handle<Monitor> monitor);

  // A thread safe api to the current entity status
  Expected<void> getEntityStatus(gxf_uid_t eid, gxf_entity_status_t* entity_status);

  // A thread safe api to check if the entity can be executed.
  // An entity is said to be busy if it is in any one of these states
  // start pending, started, tick pending or ticking
  bool isEntityBusy(gxf_uid_t eid);

  gxf_result_t getEntityBehaviorStatus(gxf_uid_t eid, entity_state_t& behavior_status);

  // Checks if all network connections have been established
  // Returns true if all connections are ready or if no network router is configured
  bool areNetworkConnectionsReady() const;

  class EntityItem {
   public:
    Entity entity;

    // store the EntityItem::execute() result for behavior parent codelet to query
    entity_state_t behavior_status = GXF_BEHAVIOR_INIT;

    // the controller for this entity (cached during activation)
    Handle<Controller> controller;

    FixedVector<Handle<DownstreamReceptiveSchedulingTerm>, kMaxSchedulingTerms>
        downstream_receptive_terms;
    FixedVector<Handle<SchedulingTerm>, kMaxSchedulingTerms> terms;
    FixedVector<Handle<Codelet>, kMaxCodeletsPerEntity> codelets;
    FixedVector<Handle<SchedulingTermCombiner>, kMaxSchedulingTermCombiners> combiners;

    Expected<SchedulingCondition> check(int64_t timestamp) const;

    Expected<bool> activate(
        Entity other, MessageRouter* message_router,
        std::shared_ptr<FixedVector<Handle<JobStatistics>>> statistics,
        nvtxDomainHandle_t nvtx_domain, uint32_t nvtx_category_id);

    Expected<SchedulingCondition> execute(int64_t timestamp, Router* router,
                                          int64_t& ticking_variation);

    Expected<void> deactivate();

    Expected<gxf_entity_status_t> getEntityStatus();

   private:
    friend class EntityExecutor;

    Expected<void> start(int64_t timestamp);

    Expected<void> tick(int64_t timestamp, Router* router);

    Expected<void> stop();

    Expected<void> startCodelet(const Handle<Codelet>& codelet);
    Expected<void> tickCodelet(const Handle<Codelet>& codelet);
    Expected<void> stopCodelet(const Handle<Codelet>& codelet);

    Expected<void> setEntityStatus(const gxf_entity_status_t next_state);

    mutable std::mutex execution_mutex_;
    mutable std::mutex status_mutex_;

    // Timestamp at which the entity was executed most recently
    int64_t last_execution_timestamp_;
    // The current status of the entity
    std::atomic<gxf_entity_status_t> status_;
    // Flag used to mark nvtx range for state changes
    bool first_status_change_ = true;
    // shared_ptr to statistics_ instance for collecting codelet metrics
    std::shared_ptr<FixedVector<Handle<JobStatistics>>> statistics_;
    // Cached flag to avoid checking statistics_->empty() in hot path.
    // Atomic because it is read from execution threads without a lock and may be
    // written by addStatistics() on a different thread.
    std::atomic<bool> has_statistics_{false};

    // NVTX Domain handle used for profiling. All entities use one domain for GXF.
    nvtxDomainHandle_t nvtx_domain_;
    // NVTX Category id used for profiling. Each entity use a separate category.
    uint32_t nvtx_category_id_;
    // NVTX Range handles used for profiling
    nvtxRangeId_t nvtx_range_codelet_start_;
    nvtxRangeId_t nvtx_range_tick_codelet_;
    nvtxRangeId_t nvtx_range_codelet_stop_;
    nvtxRangeId_t nvtx_range_entity_state_;
    // Cached NVTX strings to avoid repeated allocations in check()
    std::string nvtx_check_string_;
    std::vector<std::string> nvtx_combiner_strings_;
    std::map<gxf_uid_t, std::string> nvtx_term_strings_;
    std::array<std::string, 5> nvtx_result_strings_;
    // Cached NVTX strings for codelet operations (start, tick, stop)
    std::vector<std::string> nvtx_codelet_start_strings_;
    std::vector<std::string> nvtx_codelet_tick_strings_;
    std::vector<std::string> nvtx_codelet_stop_strings_;
    // Cached NVTX string for scheduling term onExecute calls
    std::string nvtx_onexecute_string_;
    // Cached NVTX strings for entity state updates (indexed by gxf_entity_status_t).
    // Sized by GXF_ENTITY_MAX so the array automatically covers every enum value.
    std::array<std::string, GXF_ENTITY_MAX> nvtx_state_strings_;
    // Cached entity name (copy taken at setup to avoid dangling pointer from map rehashing)
    std::string cached_entity_name_;
  };

 private:
  mutable std::shared_timed_mutex mutex_;
  std::map<gxf_uid_t, std::unique_ptr<EntityItem>> items_;
  Handle<Router> router_;
  Handle<MessageRouter> message_router_;
  Handle<NetworkRouter> network_router_;

  mutable std::mutex statistics_mutex_;
  // List of statistics components that need to be updated when running entities
  std::shared_ptr<FixedVector<Handle<JobStatistics>>> statistics_;

  mutable std::mutex monitor_mutex_;
  // List of monitor components that perform a callback when an entity executes
  FixedVector<Handle<Monitor>, kMaxMonitors> monitors_;

  // NVTX Domain handle used for profiling. All entities use one domain for GXF.
  nvtxDomainHandle_t nvtx_domain_;
  // The last NVTX Category id we have used for an entity
  uint32_t nvtx_category_id_last_ = 0;
};

}  // namespace gxf
}  // namespace nvidia

#endif  // NVIDIA_GXF_STD_ENTITY_EXECUTOR_HPP_
