#ifndef THIRD_PARTY_MILOTIC_EXTERNAL_CC_TLBMC_COLLECTOR_THERMAL_COLLECTOR_H_
#define THIRD_PARTY_MILOTIC_EXTERNAL_CC_TLBMC_COLLECTOR_THERMAL_COLLECTOR_H_

#include <atomic>
#include <cstdint>
#include <memory>
#include <string>
#include <tuple>
#include <vector>

#include "absl/base/nullability.h"
#include "absl/base/thread_annotations.h"
#include "absl/container/flat_hash_map.h"
#include "absl/container/flat_hash_set.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"
#include "absl/synchronization/mutex.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/fru_collector.h"
#include "tlbmc/collector/sensor_collector_aggregator.h"
#include "tlbmc/configs/entity_config.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 "thermal_config.pb.h"
#include "virtual_sensor_config.pb.h"
#include "sensor.pb.h"
#include "tlbmc/scheduler/scheduler.h"
#include "tlbmc/thermal/controller/dff_controller.h"
#include "tlbmc/thermal/controller/eat_controller.h"
#include "tlbmc/thermal/controller/fan_pid_controller.h"
#include "tlbmc/thermal/controller/first_order_adrc_controller.h"
#include "tlbmc/thermal/controller/optimizer/pid/pid_optimizer.h"
#include "tlbmc/thermal/controller/pid_controller.h"
#include "tlbmc/thermal/controller/second_order_adrc_controller.h"
#include "tlbmc/thermal/controller/stepwise_controller.h"
#include "tlbmc/thermal/zone_manager.h"

namespace milotic_tlbmc {

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

  ~ThermalThreadManager() = 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 IDs of thermal zones to allow managing the
  // lifetime of the zone.
  absl::Mutex thermal_zone_id_to_task_id_mutex;
  absl::flat_hash_map<std::string, int> thermal_zone_id_to_task_id
      ABSL_GUARDED_BY(thermal_zone_id_to_task_id_mutex);
};

// TODO - b/459843950: `ConfigureCollection` will be needed for tuning mode.
class ThermalCollector : public Collector {
 public:
  // Parameters for creating a thermal collector.
  struct Params {
    // Thermal control related configs.
    ZoneManagerConfigs zone_manager_configs;
    FirstOrderAdrcControllerConfigs first_order_adrc_controller_configs;
    SecondOrderAdrcControllerConfigs second_order_adrc_controller_configs;
    DffControllerConfigs dff_controller_configs;
    PidControllerConfigs pid_controller_configs;
    StepwiseControllerConfigs stepwise_controller_configs;
    FanPidControllerConfigs fan_pid_controller_configs;
    // Supportive thermal control configs.
    SupportiveZoneManagerConfigs supportive_zone_manager_configs;
    EatControllerConfigs eat_controller_configs;
    // Failsafe logger config.
    FailsafeLoggerConfig failsafe_logger_config;
    // Tuning related configs.
    bool tuning_enabled = false;
    std::string sample_data_file_prefix = "/tmp/tlbmc_thermal_data";
    // Using a thread factory to allow us to switch to fibers in future when it
    // gets open sourced.
    ecclesia::ThreadFactoryInterface* thread_factory =
        ecclesia::GetDefaultThreadFactory();
    ecclesia::Clock* clock = ecclesia::Clock::RealClock();
  };

  ~ThermalCollector() override;

  // Creates a thermal collector.
  static absl::StatusOr<std::unique_ptr<ThermalCollector>> Create(
      const Params& params, SensorCollectorAggregator* absl_nonnull aggregator,
      FruCollector* absl_nonnull fru_collector,
      EntityConfig* absl_nonnull entity_config);

  // Returns the list of thermal zones sorted by their IDs.
  virtual std::vector<const thermal::ZoneManager*> GetAllThermalZones() const;
  // Returns the list of 1st-order ADRC controllers sorted by their IDs.
  virtual std::vector<const thermal::FirstOrderAdrcController*>
  GetAllFirstOrderAdrcControllers() const;
  // Returns the list of 2nd-order ADRC controllers sorted by their IDs.
  virtual std::vector<const thermal::SecondOrderAdrcController*>
  GetAllSecondOrderAdrcControllers() const;
  // Returns the list of PID controllers sorted by their IDs.
  virtual std::vector<const thermal::PidController*> GetAllPidControllers()
      const;
  // Returns the list of stepwise controllers sorted by their IDs.
  virtual std::vector<const thermal::StepwiseController*>
  GetAllStepwiseControllers() const;
  // Returns the list of fan PID controllers sorted by their IDs.
  virtual std::vector<const thermal::FanPidController*>
  GetAllFanPidControllers() const;
  // Returns the list of DFF controllers sorted by their IDs.
  virtual std::vector<const thermal::DffController*> GetAllDffControllers()
      const;
  // Returns the list of EAT controllers sorted by their IDs.
  virtual std::vector<const thermal::EatController*> GetAllEatControllers()
      const;

  virtual absl::StatusOr<const thermal::ZoneManager*> GetThermalZoneById(
      int zone_id) const;
  virtual absl::StatusOr<const thermal::FirstOrderAdrcController*>
  GetFirstOrderAdrcControllerById(absl::string_view controller_id) const;
  virtual absl::StatusOr<const thermal::SecondOrderAdrcController*>
  GetSecondOrderAdrcControllerById(absl::string_view controller_id) const;
  virtual absl::StatusOr<const thermal::PidController*> GetPidControllerById(
      absl::string_view controller_id) const;
  virtual absl::StatusOr<const thermal::StepwiseController*>
  GetStepwiseControllerById(absl::string_view controller_id) const;
  virtual absl::StatusOr<const thermal::FanPidController*>
  GetFanPidControllerById(absl::string_view controller_id) const;
  virtual absl::StatusOr<const thermal::DffController*> GetDffControllerById(
      absl::string_view controller_id) const;
  virtual absl::StatusOr<const thermal::EatController*> GetEatControllerById(
      absl::string_view controller_id) const;

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

  nlohmann::json ToJson() const override;

  virtual void StartCollection();

  // `StartThermalControl` will trigger infinite thermal control loops for every
  // zone manager.
  void StartThermalControl();

  void StopThermalControl() { stop_thermal_control_ = true; }

  // Enable/Disable fan manual mode for all zones.
  virtual void SetManualMode(bool enable);
  // Returns the map of zone ID to manual mode status.
  virtual absl::flat_hash_map<int, bool> GetZoneIdToManualModes() const;
  virtual void SetManualPwm(absl::string_view fan_name, double pwm);

  // Tuning request functions
  virtual absl::Status RequestPidControllerCoefficientsTuning(
      absl::string_view controller_id,
      const thermal::PidController::PidControllerCoefficients& coefficients);
  virtual absl::Status RequestFanPidControllerCoefficientsTuning(
      absl::string_view controller_id,
      const thermal::FanPidController::FanPidControllerCoefficients&
          coefficients);
  virtual bool HasPendingTuningRequest(absl::string_view controller_id);

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

  thermal::ZoneManager* GetNonConstThermalZoneById(int zone_id);

  // `ProcessOneThermalLoop` will perform one thermal control calculation for
  // all the thermal loops controlled by the specified zone.
  static void ProcessOneThermalLoop(thermal::ZoneManager* zone_manager);

  // `ProcessThermalControl` will continuously perform thermal loops, including
  // fan speed update for the specified zone.
  void ProcessThermalControl(
      thermal::ZoneManager* zone_manager,
      const std::shared_ptr<boost::asio::steady_timer>& timer,
      bool first_run = false, uint64_t ms_elapsed_in_current_cycle = 0);

  // Setters (e.g., for tuning purposes).
  virtual void ApplyPendingTuningRequests(int zone_id);

  // Creates a PID auto-tuner based on the PID controller config.
  static absl::StatusOr<thermal::PidOptimizerInfo> CreatePidAutoTuner(
      const PidControllerConfig& pid_controller_config);

 private:
  struct ControllerIdToController {
    absl::flat_hash_map<std::string, thermal::FirstOrderAdrcController*>
        controller_id_to_first_order_adrc_controller;
    absl::flat_hash_map<std::string, thermal::SecondOrderAdrcController*>
        controller_id_to_second_order_adrc_controller;
    absl::flat_hash_map<std::string, thermal::PidController*>
        controller_id_to_pid_controller;
    absl::flat_hash_map<std::string, thermal::StepwiseController*>
        controller_id_to_stepwise_controller;
    absl::flat_hash_map<std::string, thermal::FanPidController*>
        controller_id_to_fan_pid_controller;
    absl::flat_hash_map<std::string, thermal::DffController*>
        controller_id_to_dff_controller;
    absl::flat_hash_map<std::string, thermal::EatController*>
        controller_id_to_eat_controller;
  };

  ThermalCollector(
      absl::flat_hash_map<
          std::tuple</*zone_id*/ int, /*controller_id*/ std::string>,
          std::unique_ptr<thermal::FirstOrderAdrcController>>&&
          zone_id_and_controller_id_to_first_order_adrc_controller,
      absl::flat_hash_map<
          std::tuple</*zone_id*/ int, /*controller_id*/ std::string>,
          std::unique_ptr<thermal::SecondOrderAdrcController>>&&
          zone_id_and_controller_id_to_second_order_adrc_controller,
      absl::flat_hash_map<int, std::unique_ptr<thermal::ZoneManager>>&&
          zone_id_to_zone_manager,
      absl::flat_hash_map<
          std::tuple</*zone_id*/ int, /*controller_id*/ std::string>,
          std::unique_ptr<thermal::PidController>>&&
          zone_id_and_controller_id_to_pid_controller,
      absl::flat_hash_map<
          std::tuple</*zone_id*/ int, /*controller_id*/ std::string>,
          std::unique_ptr<thermal::StepwiseController>>&&
          zone_id_and_controller_id_to_stepwise_controller,
      absl::flat_hash_map<
          std::tuple</*zone_id*/ int, /*controller_id*/ std::string>,
          std::unique_ptr<thermal::FanPidController>>&&
          zone_id_and_controller_id_to_fan_pid_controller,
      absl::flat_hash_map<
          std::tuple</*zone_id*/ int, /*controller_id*/ std::string>,
          std::unique_ptr<thermal::DffController>>&&
          zone_id_and_controller_id_to_dff_controller,
      absl::flat_hash_map<std::string, std::vector<int>>&&
          controller_id_to_zone_ids,
      absl::flat_hash_map<int, std::unique_ptr<thermal::ZoneManager>>&&
          zone_id_to_supportive_zone_manager,
      absl::flat_hash_map<
          std::tuple</*zone_id*/ int, /*controller_id*/ std::string>,
          std::unique_ptr<thermal::EatController>>&&
          zone_id_and_controller_id_to_eat_controller,
      std::unique_ptr<ThermalThreadManager> thermal_thread_manager,
      ecclesia::ThreadFactoryInterface* thread_factory);

  // `Link<ControllerType>ToZoneAndCreateIdToControllerMap` will link the
  // controllers to their zone managers via `ZoneManager::AddThermalController`
  // and return a map of controller ID to controller for handling future Redfish
  // requests.
  absl::flat_hash_map<std::string, thermal::FirstOrderAdrcController*>
  LinkFirstOrderAdrcControllerToZoneAndCreateIdToControllerMap();
  absl::flat_hash_map<std::string, thermal::SecondOrderAdrcController*>
  LinkSecondOrderAdrcControllerToZoneAndCreateIdToControllerMap();
  absl::flat_hash_map<std::string, thermal::PidController*>
  LinkPidControllerToZoneAndCreateIdToControllerMap();
  absl::flat_hash_map<std::string, thermal::StepwiseController*>
  LinkStepwiseControllerToZoneAndCreateIdToControllerMap();
  absl::flat_hash_map<std::string, thermal::FanPidController*>
  LinkFanPidControllerToZoneAndCreateIdToControllerMap();
  absl::flat_hash_map<std::string, thermal::DffController*>
  LinkDffControllerToZoneAndCreateIdToControllerMap();
  absl::flat_hash_map<std::string, thermal::EatController*>
  LinkEatControllerToZoneAndCreateIdToControllerMap();

  // Make sure the accessor is destroyed after the thread manager.
  // In today's use case, we will never have dynamic set of thermal zones. All
  // zones are created during daemon start-up.
  absl::flat_hash_map<int, std::unique_ptr<thermal::ZoneManager>>
      zone_id_to_zone_manager_;
  absl::flat_hash_map<std::tuple<int, std::string>,
                      std::unique_ptr<thermal::FirstOrderAdrcController>>
      zone_id_and_controller_id_to_first_order_adrc_controller_;
  absl::flat_hash_map<std::tuple<int, std::string>,
                      std::unique_ptr<thermal::SecondOrderAdrcController>>
      zone_id_and_controller_id_to_second_order_adrc_controller_;
  absl::flat_hash_map<std::tuple<int, std::string>,
                      std::unique_ptr<thermal::PidController>>
      zone_id_and_controller_id_to_pid_controller_;
  const absl::flat_hash_map<std::tuple<int, std::string>,
                            std::unique_ptr<thermal::StepwiseController>>
      zone_id_and_controller_id_to_stepwise_controller_;
  absl::flat_hash_map<std::tuple<int, std::string>,
                      std::unique_ptr<thermal::FanPidController>>
      zone_id_and_controller_id_to_fan_pid_controller_;
  const absl::flat_hash_map<std::tuple<int, std::string>,
                            std::unique_ptr<thermal::DffController>>
      zone_id_and_controller_id_to_dff_controller_;
  const absl::flat_hash_map<std::string, std::vector<int>>
      controller_id_to_zone_ids_;

  // We may need to add `controller_id_to_supportive_zone_manager` in
  // the future, but not now, as the only controller type owned by such zones
  // is the EAT controller, which is expected to be unique (i.e., only one
  // instance).
  const absl::flat_hash_map<int, std::unique_ptr<thermal::ZoneManager>>
      zone_id_to_supportive_zone_manager_;
  const absl::flat_hash_map<std::tuple<int, std::string>,
                            std::unique_ptr<thermal::EatController>>
      zone_id_and_controller_id_to_eat_controller_;

  std::atomic<bool> stop_thermal_control_ = false;

  const ControllerIdToController controller_id_to_controller_;

  // Tuning mode related members.
  absl::Mutex tuning_mutex_;
  absl::flat_hash_set<int> zone_ids_with_pending_tuning_request_
      ABSL_GUARDED_BY(tuning_mutex_);
  // {<Controller ID> := <pending thermal coefficient>}
  absl::flat_hash_map<std::string,
                      thermal::PidController::PidControllerCoefficients>
      controller_id_to_pending_pid_coefficient_ ABSL_GUARDED_BY(tuning_mutex_);
  absl::flat_hash_map<std::string,
                      thermal::FanPidController::FanPidControllerCoefficients>
      controller_id_to_pending_fan_pid_coefficient_
          ABSL_GUARDED_BY(tuning_mutex_);

  std::vector<std::shared_ptr<boost::asio::steady_timer>>
      thermal_control_timers_;
  std::unique_ptr<ThermalThreadManager> thermal_thread_manager_;
  ecclesia::ThreadFactoryInterface* thread_factory_;
};

class EmptyThermalCollector final : public ThermalCollector {
 public:
  static std::unique_ptr<EmptyThermalCollector> Create();

  // Returns the list of thermal zones sorted by their IDs.
  std::vector<const thermal::ZoneManager*> GetAllThermalZones() const override;
  // Returns the list of 1st-order ADRC controllers sorted by their IDs.
  std::vector<const thermal::FirstOrderAdrcController*>
  GetAllFirstOrderAdrcControllers() const override;
  // Returns the list of 2nd-order ADRC controllers sorted by their IDs.
  std::vector<const thermal::SecondOrderAdrcController*>
  GetAllSecondOrderAdrcControllers() const override;
  // Returns the list of PID controllers sorted by their IDs.
  std::vector<const thermal::PidController*> GetAllPidControllers()
      const override;
  // Returns the list of stepwise controllers sorted by their IDs.
  std::vector<const thermal::StepwiseController*> GetAllStepwiseControllers()
      const override;
  // Returns the list of fan PID controllers sorted by their IDs.
  std::vector<const thermal::FanPidController*> GetAllFanPidControllers()
      const override;
  // Returns the list of DFF controllers sorted by their IDs.
  std::vector<const thermal::DffController*> GetAllDffControllers()
      const override;
  // Returns the list of EAT controllers sorted by their IDs.
  std::vector<const thermal::EatController*> GetAllEatControllers()
      const override;

  absl::StatusOr<const thermal::ZoneManager*> GetThermalZoneById(
      int zone_id) const override;
  absl::StatusOr<const thermal::FirstOrderAdrcController*>
  GetFirstOrderAdrcControllerById(
      absl::string_view controller_id) const override;
  absl::StatusOr<const thermal::SecondOrderAdrcController*>
  GetSecondOrderAdrcControllerById(
      absl::string_view controller_id) const override;
  absl::StatusOr<const thermal::PidController*> GetPidControllerById(
      absl::string_view controller_id) const override;
  absl::StatusOr<const thermal::StepwiseController*> GetStepwiseControllerById(
      absl::string_view controller_id) const override;
  absl::StatusOr<const thermal::FanPidController*> GetFanPidControllerById(
      absl::string_view controller_id) const override;
  absl::StatusOr<const thermal::DffController*> GetDffControllerById(
      absl::string_view controller_id) const override;
  absl::StatusOr<const thermal::EatController*> GetEatControllerById(
      absl::string_view controller_id) const override;

  nlohmann::json GetSchedulerStats() const override;

  nlohmann::json ToJson() const override;

  void StartCollection();

  void SetManualMode(bool enable) override;
  absl::flat_hash_map<int, bool> GetZoneIdToManualModes() const override;
  void SetManualPwm(absl::string_view fan_name, double pwm) override;

  absl::Status RequestPidControllerCoefficientsTuning(
      absl::string_view controller_id,
      const thermal::PidController::PidControllerCoefficients& coefficients)
      override;
  absl::Status RequestFanPidControllerCoefficientsTuning(
      absl::string_view controller_id,
      const thermal::FanPidController::FanPidControllerCoefficients&
          coefficients) override;
};

}  // namespace milotic_tlbmc

#endif  // THIRD_PARTY_MILOTIC_EXTERNAL_CC_TLBMC_COLLECTOR_THERMAL_COLLECTOR_H_