#ifndef THIRD_PARTY_MILOTIC_EXTERNAL_CC_TLBMC_COLLECTOR_SENSOR_COLLECTOR_H_
#define THIRD_PARTY_MILOTIC_EXTERNAL_CC_TLBMC_COLLECTOR_SENSOR_COLLECTOR_H_

#include <memory>
#include <optional>
#include <string>
#include <utility>
#include <vector>

#include "absl/base/thread_annotations.h"
#include "absl/container/flat_hash_map.h"
#include "absl/functional/any_invocable.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/synchronization/mutex.h"
#include "absl/types/span.h"
#include "boost/asio.hpp"  //NOLINT: boost::asio is commonly used in BMC
#include "thread/thread.h"
#include "time/clock.h"
#include "nlohmann/json.hpp"
#include "tlbmc/collector/collector.h"
#include "tlbmc/collector/peci_scanner.h"
#include "fan_controller_config.pb.h"
#include "fan_pwm_config.pb.h"
#include "fan_tach_config.pb.h"
#include "hwmon_temp_sensor_config.pb.h"
#include "intel_cpu_sensor_config.pb.h"
#include "nic_telemetry_config.pb.h"
#include "psu_sensor_config.pb.h"
#include "shared_mem_sensor_config.pb.h"
#include "virtual_sensor_config.pb.h"
#include "tlbmc/hal/nic_veeprom/interface.h"
#include "tlbmc/hal/peci/peci_access_interface.h"
#include "tlbmc/hal/sysfs/i2c.h"
#include "tlbmc/hal/sysfs/i3c.h"
#include "tlbmc/hal/sysfs/peci.h"
#include "sensor.pb.h"
#include "tlbmc/scheduler/scheduler.h"
#include "tlbmc/sensors/sensor.h"

namespace milotic_tlbmc {

struct ThreadManager {
  explicit ThreadManager(ecclesia::Clock* clock)
      : task_scheduler(std::make_unique<TaskScheduler>(clock)) {}

  ~ThreadManager() = default;

  std::vector<std::unique_ptr<ecclesia::ThreadInterface>> threads;
  std::vector<std::shared_ptr<boost::asio::io_context>> io_contexts;
  std::vector<
      boost::asio::executor_work_guard<boost::asio::io_context::executor_type>>
      work_guards;
  std::unique_ptr<TaskScheduler> task_scheduler;
  // This is used to store the task id of the sensor to allow managing the
  // lifetime of the sensor.
  absl::Mutex key_to_task_id_mutex;
  absl::flat_hash_map<std::string, int> sensor_key_to_task_id
      ABSL_GUARDED_BY(key_to_task_id_mutex);
};

// Register a callback to be invoked when a collector refreshes data.
// This class is thread safe.
// Used mostly for testing.
class SensorNotification {
 public:
  explicit SensorNotification(
      absl::AnyInvocable<void(const std::shared_ptr<const SensorValue>&) const>
          callback)
      : callback_(std::move(callback)) {}

  void NotifyWithData(const std::shared_ptr<const SensorValue>& sensor_data)
      const ABSL_LOCKS_EXCLUDED(mutex_) {
    absl::MutexLock lock(&mutex_);
    callback_(sensor_data);
  }

 private:
  mutable absl::Mutex mutex_;
  const absl::AnyInvocable<void(const std::shared_ptr<const SensorValue>&)
                               const>
      callback_ ABSL_GUARDED_BY(mutex_);
};

class SensorCollector : public Collector {
 public:
  struct SensorConfigs {
    absl::Span<const HwmonTempSensorConfig> hwmon_temp_sensor_configs;
    absl::Span<const PsuSensorConfig> psu_sensor_configs;
    absl::Span<const FanControllerConfig> fan_controller_configs;
    absl::Span<const FanPwmConfig> fan_pwm_configs;
    absl::Span<const FanTachConfig> fan_tach_configs;
    absl::Span<const SharedMemSensorConfig> shared_mem_sensor_configs;
    absl::Span<const VirtualSensorConfig> virtual_sensor_configs;
    absl::Span<const IntelCpuSensorConfig> intel_cpu_sensor_configs;
    absl::Span<const NicTelemetryConfig> nic_telemetry_configs;
  };

  // Parameters for creating a sensor collector.
  struct Params {
    SensorConfigs sensor_configs;
    I2cSysfs& i2c_sysfs;
    I3cSysfs& i3c_sysfs;
    PeciSysfs& peci_sysfs;
    PeciAccessInterface& peci_access;
    // Using a thread factory to allow us to switch to fibers in future when it
    // gets open sourced.
    ecclesia::ThreadFactoryInterface* thread_factory =
        ecclesia::GetDefaultThreadFactory();
    SensorNotification* refresh_notification = nullptr;
    ecclesia::Clock* clock = ecclesia::Clock::RealClock();
    // Typically used when we want to override the default sensor sampling
    // interval.
    std::optional<int> override_sensor_sampling_interval_ms = std::nullopt;
    // Factory for creating NIC accessors, allowing mocks in tests.
    absl::AnyInvocable<std::unique_ptr<nic_veeprom::Accessor>(
        nic_veeprom::NicTelemetryVersion version, int bus, int address) const>
        nic_accessor_factory = &nic_veeprom::Accessor::Create;
  };

  ~SensorCollector() override;

  // Creates a sensor collector.
  static absl::StatusOr<std::unique_ptr<SensorCollector>> Create(
      const Params& params);

  // Returns the sorted list of sensor names contained by the given config key.
  virtual std::vector<std::string> GetAllSensorKeysByConfigKey(
      const std::string& board_config_key) const;

  // Reinitializes all sensors for the given config key.
  void ReinitializeAndScheduleAllSensorsForConfigKey(
      const std::string& board_config_key);

  // Returns all the sensors sorted by sensor name.
  virtual std::vector<std::shared_ptr<const Sensor>> GetAllSensors() const;

  // Returns the sensor for the given sensor key.
  virtual std::shared_ptr<const Sensor> GetSensorBySensorKey(
      const std::string& sensor_key) const;

  // Configures the sensor collector.
  // A key usage of this method is to configure the sampling interval of the
  // data.
  absl::Status ConfigureCollection(const Config& config) const override;

  // Returns the sensor for the given board config key and given sensor
  // key.
  virtual std::shared_ptr<const Sensor> GetSensorByConfigKeyAndSensorKey(
      const std::string& board_config_key, const std::string& sensor_key) const;

  void ReinitializeSensorByConfigKeyAndSensorKey(
      const std::string& board_config_key, const std::string& sensor_key) const;

  nlohmann::json GetSchedulerStats() const override {
    return thread_manager_->task_scheduler->ToJson();
  }

  nlohmann::json ToJson() const override;

  const ThreadManager* GetThreadManager() const {
    return thread_manager_.get();
  }

  virtual void StartCollection();

 protected:
  // Protected constructor to allow mocking.
  SensorCollector() = default;

  void InitializePeciScanner(const Params& params);

 private:
  SensorCollector(
      std::vector<std::shared_ptr<Sensor>>&& sensors,
      absl::flat_hash_map<int, std::unique_ptr<nic_veeprom::Accessor>>&&
          bus_to_v1_accessor,
      absl::flat_hash_map<int, std::unique_ptr<nic_veeprom::Accessor>>&&
          bus_to_v2_accessor,
      std::unique_ptr<ThreadManager> thread_manager,
      const SensorNotification* refresh_notification,
      std::optional<int> override_sensor_sampling_interval_ms,
      ecclesia::ThreadFactoryInterface* thread_factory,
      std::unique_ptr<PeciScanner> peci_scanner);

  // Make sure the accessor is destroyed after the thread manager.
  absl::flat_hash_map<int, std::unique_ptr<nic_veeprom::Accessor>>
      bus_to_v1_accessor_;
  absl::flat_hash_map<int, std::unique_ptr<nic_veeprom::Accessor>>
      bus_to_v2_accessor_;

  // In today's use case, we will never have dynamic set of sensors. All sensors
  // are created during daemon start-up.
  // First dimension: board_config_key
  // Second dimension: sensor_key
  const absl::flat_hash_map<
      std::string, absl::flat_hash_map<std::string, std::shared_ptr<Sensor>>>
      sensor_table_;
  // Order is important here, PeciScanner needs to be destroyed before
  // ThreadManager. It contains a pointer to the TaskScheduler in the
  // ThreadManager.
  std::unique_ptr<ThreadManager> thread_manager_;
  std::unique_ptr<PeciScanner> peci_scanner_;
  const SensorNotification* refresh_notification_;
  std::optional<int> override_sensor_sampling_interval_ms_;
  ecclesia::ThreadFactoryInterface* thread_factory_;
};

class EmptySensorCollector final : public SensorCollector {
 public:
  static std::unique_ptr<EmptySensorCollector> Create();

  // Returns the sorted list of sensor names contained by the given config key.
  std::vector<std::string> GetAllSensorKeysByConfigKey(
      const std::string& board_config_key) const override;

  // Returns all the sensors sorted by sensor name.
  std::vector<std::shared_ptr<const Sensor>> GetAllSensors() const override;

  // Returns the sensor for the given sensor key.
  std::shared_ptr<const Sensor> GetSensorBySensorKey(
      const std::string& sensor_key) const override;

  absl::Status ConfigureCollection(const Config& config) const override;

  std::shared_ptr<const Sensor> GetSensorByConfigKeyAndSensorKey(
      const std::string& board_config_key,
      const std::string& sensor_key) const override;

  nlohmann::json GetSchedulerStats() const override;

  nlohmann::json ToJson() const override;

  void StartCollection();
};

}  // namespace milotic_tlbmc

#endif  // THIRD_PARTY_MILOTIC_EXTERNAL_CC_TLBMC_COLLECTOR_SENSOR_COLLECTOR_H_