#pragma once

#include <smacc2/smacc.hpp>
#include <boost/mpl/list.hpp>
#include <rclcpp/rclcpp.hpp>
#include <rclcpp_action/rclcpp_action.hpp>
#include <atomic>
#include <string>
#include <functional>

namespace brain
{

namespace sc = boost::statechart;

// ----------------------------------------------------------------------------
// Runtime shared data
// ----------------------------------------------------------------------------
/**
 * @brief Shared runtime container passed implicitly between the state machine
 * (SMACC states) and the action server glue code.
 *
 * Fields are intentionally simple POD / atomics so they can be written from
 * multiple threads (worker threads, cancel callbacks, etc.) without heavy
 * synchronization. Lifecycle is a process‑wide singleton via runtime().
 */
struct CerebellumData
{
  /** @brief External preemption / cancel flag (set by on_cancel or timeout). */
  std::atomic<bool> preempt{false};            // external cancel / interruption flag
  /** @brief Human readable reason associated with last preempt / completion. */
  std::string preempt_reason;                  // reason / last message

  /**
   * @brief Result object produced by the asynchronous execute_fn functor.
   *
   * It aggregates overall success plus a textual summary and per‑sequence
   * counters so the state machine terminal states can publish action results.
   */
  struct ExecResult
  {
    bool success{true};                 ///< True if sequence finished without failure / cancel.
    std::string message;                ///< Final textual summary of the sequence.
    int32_t succeeded_skills{0};        ///< Number of skills that succeeded.
    int32_t total_skills{0};            ///< Total skills attempted / requested.
  };

  /**
   * @brief Bound by the action execute callback: performs the whole skill
   * sequence synchronously when invoked inside StExecuting supervision thread.
   */
  std::function<ExecResult(rclcpp::Logger)> execute_fn;          // installed by action execute
  /**
   * @brief Completion callback invoked by terminal states (Completed/Failed/Emergency)
   * to produce the actual action result (succeed/abort) back to the client.
   */
  std::function<void(bool, const std::string &)> complete_cb;    // called in terminal states

  /** @brief Latched last success flag (copied from ExecResult). */
  std::atomic<bool> last_success{true};
  /** @brief Latched total skills count from ExecResult. */
  std::atomic<int32_t> last_total_skills{0};
  /** @brief Latched succeeded skills count from ExecResult. */
  std::atomic<int32_t> last_succeeded_skills{0};
};

/**
 * @brief Singleton accessor for runtime shared data.
 * @return Reference to globally stored CerebellumData instance.
 */
inline CerebellumData & runtime()
{
  static CerebellumData inst; return inst;
}

// ----------------------------------------------------------------------------
// Events
// ----------------------------------------------------------------------------
/** @brief Event: a new goal has been received and accepted. */
struct EvGoalReceived : sc::event<EvGoalReceived> {};
/** @brief Event: asynchronous execution finished successfully. */
struct EvExecutionFinished : sc::event<EvExecutionFinished> {};
/** @brief Event: asynchronous execution ended with a failure (non emergency). */
struct EvExecutionFailed : sc::event<EvExecutionFailed> {};
/** @brief Event: unexpected exception / emergency condition occurred. */
struct EvEmergency : sc::event<EvEmergency> {};

// ----------------------------------------------------------------------------
// Forward declarations
// ----------------------------------------------------------------------------
struct StIdle; struct StExecuting; struct StCompleted; struct StFailed; struct StEmergency;

// ----------------------------------------------------------------------------
// State Machine
// ----------------------------------------------------------------------------
/**
 * @brief Root SMACC2 state machine for the cerebellum execution lifecycle.
 * Initial state is StIdle.
 */
struct SmCerebellum : smacc2::SmaccStateMachineBase<SmCerebellum, StIdle>
{ using SmaccStateMachineBase::SmaccStateMachineBase; };

// ----------------------------------------------------------------------------
// Idle
// ----------------------------------------------------------------------------
/**
 * @brief Idle state: waits for EvGoalReceived to start executing a skill sequence.
 */
struct StIdle : smacc2::SmaccState<StIdle, SmCerebellum>
{
  using SmaccState::SmaccState;
  typedef boost::mpl::list<sc::custom_reaction<EvGoalReceived>> reactions; ///< Supported events.
  /**
   * @brief Reaction: transition to StExecuting when a goal arrives.
   * @param unused Event payload (none).
   * @return Transition result to Executing.
   */
  sc::result react(const EvGoalReceived &)
  {
    RCLCPP_INFO(this->getLogger(), "[SM] Transition Idle -> Executing (goal received)");
    return transit<StExecuting>();
  }
};

// ----------------------------------------------------------------------------
// Executing
// ----------------------------------------------------------------------------
/**
 * @brief Executing state: waits until an execute_fn is installed, then runs it
 * in a detached supervision thread, posting outcome events.
 *
 * It periodically logs a warning if execute_fn has not yet been set (defensive
 * against race conditions). Exceptions are mapped to EvEmergency.
 */
struct StExecuting : smacc2::SmaccState<StExecuting, SmCerebellum>
{
  using SmaccState::SmaccState;
  typedef boost::mpl::list<sc::custom_reaction<EvExecutionFinished>, sc::custom_reaction<EvExecutionFailed>, sc::custom_reaction<EvEmergency>> reactions; ///< Reaction table.

  /** @brief State entry: spawns supervision thread and invokes execute_fn when ready. */
  void onEntry()
  {
    RCLCPP_INFO(this->getLogger(), "[SM] Enter Executing");
    // Supervising thread waits until execute_fn is set, then runs it.
    std::thread{[this]() {
        auto & d = runtime();
        using namespace std::chrono_literals;
        auto last_warn = std::chrono::steady_clock::now();
        while (!d.execute_fn) {
          std::this_thread::sleep_for(5ms);
          auto now = std::chrono::steady_clock::now();
            if (now - last_warn > 5s) {
              RCLCPP_WARN(this->getLogger(), "[SM] Still waiting for execute_fn to be set...");
              last_warn = now;
            }
        }
        try {
          auto res = d.execute_fn(this->getLogger());
          d.preempt.store(!res.success, std::memory_order_release);
          d.preempt_reason = res.message;
          d.last_success.store(res.success, std::memory_order_release);
          d.last_total_skills.store(res.total_skills, std::memory_order_release);
          d.last_succeeded_skills.store(res.succeeded_skills, std::memory_order_release);
          if (res.success) this->template postEvent<EvExecutionFinished>();
          else this->template postEvent<EvExecutionFailed>();
        } catch (const std::exception & e) {
          d.preempt.store(true, std::memory_order_release);
          d.preempt_reason = std::string("Exception: ") + e.what();
          this->template postEvent<EvEmergency>();
        } catch (...) {
          d.preempt.store(true, std::memory_order_release);
          d.preempt_reason = "Unknown exception";
          this->template postEvent<EvEmergency>();
        }
      }}.detach();
  }

  /**
   * @brief Reaction: successful completion -> StCompleted.
   * @return Transition to Completed state.
   */
  sc::result react(const EvExecutionFinished &)
  { RCLCPP_INFO(this->getLogger(), "[SM] Transition Executing -> Completed"); return transit<StCompleted>(); }
  /**
   * @brief Reaction: functional failure -> StFailed.
   * @return Transition to Failed state.
   */
  sc::result react(const EvExecutionFailed &)
  { RCLCPP_WARN(this->getLogger(), "[SM] Transition Executing -> Failed"); return transit<StFailed>(); }
  /**
   * @brief Reaction: exception / emergency -> StEmergency.
   * @return Transition to Emergency state.
   */
  sc::result react(const EvEmergency &)
  { RCLCPP_ERROR(this->getLogger(), "[SM] Transition Executing -> Emergency"); return transit<StEmergency>(); }
};

// ----------------------------------------------------------------------------
// Completed
// ----------------------------------------------------------------------------
/**
 * @brief Completed state: final success path; invokes completion callback and
 * resets runtime fields before optionally accepting a new goal.
 */
struct StCompleted : smacc2::SmaccState<StCompleted, SmCerebellum>
{
  using SmaccState::SmaccState;
  typedef boost::mpl::list<sc::custom_reaction<EvGoalReceived>> reactions; ///< Accepts new goals.
  /** @brief State entry: call complete_cb with success=true and clear runtime. */
  void onEntry()
  {
    auto & d = runtime();
    if (d.complete_cb) {
      bool success = d.last_success.load(std::memory_order_acquire) && !d.preempt.load();
      std::string msg = d.preempt_reason.empty() ? std::string("Done") : d.preempt_reason;
      d.complete_cb(success, msg);
    }
    d.preempt.store(false); d.preempt_reason.clear(); d.execute_fn = nullptr; d.complete_cb = nullptr;
    d.last_total_skills.store(0); d.last_succeeded_skills.store(0);
  }
  /**
   * @brief Reaction: new goal transitions back to Executing.
   * @return Transition to StExecuting.
   */
  sc::result react(const EvGoalReceived &)
  { RCLCPP_INFO(this->getLogger(), "[SM] Transition Completed -> Executing (new goal)"); return transit<StExecuting>(); }
};

// ----------------------------------------------------------------------------
// Failed
// ----------------------------------------------------------------------------
/**
 * @brief Failed state: sequence ended in a non‑emergency failure. Reports
 * failure via completion callback and resets runtime data.
 */
struct StFailed : smacc2::SmaccState<StFailed, SmCerebellum>
{
  using SmaccState::SmaccState;
  typedef boost::mpl::list<sc::custom_reaction<EvGoalReceived>> reactions; ///< Can accept a new goal.
  /** @brief State entry: invoke completion callback with failure status. */
  void onEntry()
  {
    auto & d = runtime();
    if (d.complete_cb) d.complete_cb(false, d.preempt_reason.empty()?"Failed":d.preempt_reason);
    d.execute_fn = nullptr; d.complete_cb = nullptr; d.preempt.store(false);
  }
  /**
   * @brief Reaction: accept new goal and transition to Executing.
   * @return Transition to StExecuting.
   */
  sc::result react(const EvGoalReceived &)
  { RCLCPP_INFO(this->getLogger(), "[SM] Transition Failed -> Executing (new goal)"); return transit<StExecuting>(); }
};

// ----------------------------------------------------------------------------
// Emergency
// ----------------------------------------------------------------------------
/**
 * @brief Emergency state: entered due to an unhandled exception in execute_fn.
 * Calls completion callback with an emergency prefix and resets runtime.
 */
struct StEmergency : smacc2::SmaccState<StEmergency, SmCerebellum>
{
  using SmaccState::SmaccState;
  typedef boost::mpl::list<sc::custom_reaction<EvGoalReceived>> reactions; ///< Can accept a new goal.
  /** @brief State entry: log emergency and invoke completion callback. */
  void onEntry()
  {
    auto & d = runtime();
    if (d.complete_cb) d.complete_cb(false, std::string("EMERGENCY: ")+ d.preempt_reason);
    RCLCPP_ERROR(this->getLogger(), "[SM] Emergency: %s", d.preempt_reason.c_str());
    d.execute_fn = nullptr; d.complete_cb = nullptr; d.preempt.store(false);
  }
  /**
   * @brief Reaction: accept new goal and transition to Executing.
   * @return Transition to StExecuting.
   */
  sc::result react(const EvGoalReceived &)
  { RCLCPP_INFO(this->getLogger(), "[SM] Transition Emergency -> Executing (new goal)"); return transit<StExecuting>(); }
};

} // namespace brain
