#pragma once

#include "esphome/core/component.h"
#include "esphome/core/helpers.h"
#include "esphome/components/sensor/sensor.h"

#include <functional>

namespace esphome::combination {

class CombinationComponent : public Component, public sensor::Sensor {
 public:
  /// @brief Logs all source sensor's names
  virtual void log_source_sensors() = 0;

 protected:
  /// @brief Logs the sensor for use in dump_config
  /// @param combo_type Name of the combination operation
  void log_config_(const LogString *combo_type);
};

/// @brief Base class for operations that do not require an extra parameter to compute the combination
class CombinationNoParameterComponent : public CombinationComponent {
 public:
  /// @brief Adds a callback to each source sensor
  void setup() override;

  void add_source(Sensor *sensor);

  /// @brief Computes the combination
  /// @param value Newest sensor measurement
  virtual void handle_new_value(float value) = 0;

  /// @brief Logs all source sensor's names in sensors_
  void log_source_sensors() override;

 protected:
  std::vector<Sensor *> sensors_;
};

// Base class for opertions that require one parameter to compute the combination
class CombinationOneParameterComponent : public CombinationComponent {
 public:
  void set_source_count(size_t count) { this->sensor_sources_.init(count); }
  void add_source(Sensor *sensor, std::function<float(float)> const &compute);
  void add_source(Sensor *sensor, float value);

  /// @brief Logs all source sensors' names in sensor_sources_
  void log_source_sensors() override;

 protected:
  struct SensorSource {
    sensor::Sensor *sensor;
    std::function<float(float)> compute;
    CombinationOneParameterComponent *parent;
  };

  FixedVector<SensorSource> sensor_sources_;
};

class KalmanCombinationComponent final : public CombinationOneParameterComponent {
 public:
  void dump_config() override;
  void setup() override;

  void set_process_std_dev(float process_std_dev) {
    this->update_variance_value_ = process_std_dev * process_std_dev * 0.001f;
  }
  void set_std_dev_sensor(Sensor *sensor) { this->std_dev_sensor_ = sensor; }

 protected:
  void update_variance_();
  void correct_(float value, float stddev);

  // Optional sensor for publishing the current error
  sensor::Sensor *std_dev_sensor_{nullptr};

  // Tick of the last update
  uint32_t last_update_{0};
  // Change of the variance, per ms
  float update_variance_value_{0.f};

  // Best guess for the state and its variance
  float state_{NAN};
  float variance_{INFINITY};
};

class LinearCombinationComponent final : public CombinationOneParameterComponent {
 public:
  void dump_config() override { this->log_config_(LOG_STR("linear")); }
  void setup() override;

  void handle_new_value(float value);
};

class MaximumCombinationComponent final : public CombinationNoParameterComponent {
 public:
  void dump_config() override { this->log_config_(LOG_STR("max")); }

  void handle_new_value(float value) override;
};

class MeanCombinationComponent final : public CombinationNoParameterComponent {
 public:
  void dump_config() override { this->log_config_(LOG_STR("mean")); }

  void handle_new_value(float value) override;
};

class MedianCombinationComponent final : public CombinationNoParameterComponent {
 public:
  void dump_config() override { this->log_config_(LOG_STR("median")); }

  void handle_new_value(float value) override;
};

class MinimumCombinationComponent final : public CombinationNoParameterComponent {
 public:
  void dump_config() override { this->log_config_(LOG_STR("min")); }

  void handle_new_value(float value) override;
};

class MostRecentCombinationComponent final : public CombinationNoParameterComponent {
 public:
  void dump_config() override { this->log_config_(LOG_STR("most_recently_updated")); }

  void handle_new_value(float value) override;
};

class RangeCombinationComponent final : public CombinationNoParameterComponent {
 public:
  void dump_config() override { this->log_config_(LOG_STR("range")); }

  void handle_new_value(float value) override;
};

class SumCombinationComponent final : public CombinationNoParameterComponent {
 public:
  void dump_config() override { this->log_config_(LOG_STR("sum")); }

  void handle_new_value(float value) override;
};

}  // namespace esphome::combination
