/**
 * @file PowerManager.h
 * @author Phillip Kühne
 * @brief This component provides utilities for keeping track of power usage
 * consumption.
 * @version 0.1
 * @date 2024-11-23
 */

#include "PowerManager.h"

SemaphoreHandle_t PowerManager::powerMutex = NULL;

void vTaskUpdatePowerState(void *pvParameters) {
  for (;;) {
    ESP_LOGV(TAG, "Updating Power State...");
    TickType_t executionStart = xTaskGetTickCount();
    PowerManager::updatePowerStateHandler();
    vTaskDelayUntil(
        &executionStart,
        pdMS_TO_TICKS(PowerParameters::POWER_STATE_UPDATE_INTERVAL_MS));
  }
}

void PowerManager::begin() {
  // Check if another instance of us already initialized the power scheduler,
  // if not, we will do it.
  ESP_LOGI(TAG, "Initializing Power Management");

  // Create mutex if it doesn't exist
  if (powerMutex == NULL) {
    powerMutex = xSemaphoreCreateMutex();
    if (powerMutex == NULL) {
      ESP_LOGE(TAG, "Failed to create power mutex");
      Serial.println("Failed to create power mutex");
    }
  }

  // Initialize IO for battery voltage measurement
  pinMode(PowerParameters::PinConfig::BAT_ADC_EN, OUTPUT);
  pinMode(PowerParameters::PinConfig::BAT_ADC, INPUT);

  if (powerScheduler == nullptr) {
    ESP_LOGI(TAG, "Creating Power Scheduler");
    powerScheduler = &PowerScheduler::getPowerScheduler(
        PowerParameters::Battery::CELL_CURRENT_1C_MA,
        PowerParameters::Battery::CELL_CURRENT_2C_MA);
    PowerManager::initPowerState();
    PowerManager::recalculateCurrentBudgets();
    TaskHandle_t xHandle = NULL;
    xTaskCreate(vTaskUpdatePowerState, "vTaskPowerStateUpdate", 4096, NULL,
                tskIDLE_PRIORITY, &xHandle);
    configASSERT(xHandle);

    if (!(powerScheduler->tryAccquireCurrentAllowance(
            PowerParameters::PowerConsumers::ESP,
            PowerParameters::CurrentConsumptions::CURRENT_ESP_AVG))) {
      ESP_LOGE(TAG, "Could not get power for ESP");
      Serial.println("Could not get power for ESP");
    }
  }
}

float PowerManager::getFreeLimitCurrentBudget() {
  return powerScheduler->getFreeLimitCurrentBudget();
}

float PowerManager::getFreeMaximumCurrentBudget() {
  return powerScheduler->getFreeMaximumCurrentBudget();
}

bool PowerManager::tryAccquireCurrentAllowance(
    PowerParameters::PowerConsumers consumer, uint16_t neededCurrent,
    uint16_t requestedDurationMs) {
  return powerScheduler->tryAccquireCurrentAllowance(consumer, neededCurrent,
                                                     requestedDurationMs);
}

void PowerManager::releaseCurrent(PowerParameters::PowerConsumers consumer) {
  powerScheduler->releaseCurrent(consumer);
}

bool PowerManager::waitForCurrentAllowance(PowerParameters::PowerConsumers consumer,
                                    uint16_t neededCurrent,
                                    uint16_t maxSlackTimeMs,
                                    uint16_t requestedDurationMs) {
  return powerScheduler->waitForCurrentAllowance(
      consumer, neededCurrent, maxSlackTimeMs, requestedDurationMs);
}

void PowerManager::beginPermanentDeepSleep(void) {
  return powerScheduler->beginPermanentDeepSleep();
}

float PowerManager::getCurrentCurrent(void) {
  return powerScheduler->getCurrentCurrent();
}

float PowerManager::getBatteryCurrent(void) {
  const float i_3v3 = getCurrentCurrent();
  const float u_3v3 = 3.3;
  const float u_bat = getBatteryVoltage();
  const float eta = PowerParameters::BUCK_BOOST_EFFICIENCY;
  return (u_3v3 * i_3v3) / (u_bat * eta);
}

void PowerManager::recalculateCurrentBudgets(void) {
  return powerScheduler->recalculateCurrentBudgets();
}

float PowerManager::getConsumerCurrent(PowerParameters::PowerConsumers consumer) {
  return powerScheduler->getConsumerCurrent(consumer);
}

float PowerManager::getBatteryVoltage() {
  // Get the battery voltage from the ADC and convert it to a voltage
  // using the voltage divider.
  // Enable voltage divider
  digitalWrite(PowerParameters::PinConfig::BAT_ADC_EN, HIGH);
  // Allow voltage to stabilize
  delayMicroseconds(10);
  // Returns value between 0 and 4095 mapping to between 0 and 3.3V
  uint16_t batteryAdcValue = analogRead(PowerParameters::PinConfig::BAT_ADC);
  // Disable voltage divider
  digitalWrite(PowerParameters::PinConfig::BAT_ADC_EN, LOW);
  // Convert ADC value to voltage
  float batteryVoltage =
      (batteryAdcValue / 4095.0 * 3.3) *
      PowerParameters::Battery::BAT_ADC::VOLTAGE_DIVIDER_FACTOR;
  ESP_LOGD(TAG, "Battery ADC value: %d, Calculated voltage: %.2f",
           batteryAdcValue, batteryVoltage);
  return batteryVoltage;
}

float PowerManager::getBatteryChargePercent() { return percentRemaining; }

float PowerManager::getBatteryChargeCoulombs() { return coloumbsRemaining; }

float PowerManager::getBatteryVoltageChargePercent() {
  // Directly get the battery voltage, correct the curve with an offset and
  // calculate the charge state based on the discharge curve.
  float batteryVoltage = getBatteryVoltage() + PowerManager::fullVoltageOffset;
  float chargeState = 0;
  // Clamp edge cases
  if (batteryVoltage >=
      PowerParameters::Battery::DISCHARGE_CURVE::VOLTAGES[0]) {
    return PowerParameters::Battery::DISCHARGE_CURVE::CHARGE_STATES[0];
  }
  if (batteryVoltage <=
      PowerParameters::Battery::DISCHARGE_CURVE::VOLTAGES
          [PowerParameters::Battery::DISCHARGE_CURVE::NUM_POINTS - 1]) {
    return PowerParameters::Battery::DISCHARGE_CURVE::CHARGE_STATES
        [PowerParameters::Battery::DISCHARGE_CURVE::NUM_POINTS - 1];
  }
  float p1_x, p1_y, p2_x, p2_y;
  for (int i = 0; i < PowerParameters::Battery::DISCHARGE_CURVE::NUM_POINTS;
       i++) {
    if (batteryVoltage >=
        PowerParameters::Battery::DISCHARGE_CURVE::VOLTAGES[i]) {
      p1_y = PowerParameters::Battery::DISCHARGE_CURVE::CHARGE_STATES[i];
      p1_x = PowerParameters::Battery::DISCHARGE_CURVE::VOLTAGES[i];
      p2_y = PowerParameters::Battery::DISCHARGE_CURVE::CHARGE_STATES[i + 1];
      p2_x = PowerParameters::Battery::DISCHARGE_CURVE::VOLTAGES[i + 1];
      chargeState =
          ((p2_y - p1_y) / (p2_x - p1_x)) * (batteryVoltage - p1_x) + p1_y;
      return chargeState;
    }
  }
}

void PowerManager::updatePowerStateHandler() {

  // Update supply and charge state flags
  PowerManager::busPowered = !digitalRead(PowerParameters::PinConfig::VUSB_SENS);
  PowerManager::chargingState = digitalRead(PowerParameters::PinConfig::BAT_CHG_STAT);

  // If the battery is charging and fully charged
  if (PowerManager::busPowered && !PowerManager::chargingState) {
    // Calibrate voltage offset on full Battery
    PowerManager::fullVoltageOffset = PowerParameters::Battery::DISCHARGE_CURVE::VOLTAGES[0] - getBatteryVoltage();
  }

  ESP_LOGD(TAG, "Bus Powered: %d, Charging: %d", busPowered, chargingState);

  float currentCurrent = powerScheduler->getCurrentCurrent();
  int referenceCurrentMa =
      PowerParameters::Battery::DISCHARGE_CURVE::REFERENCE_CURRENT_A * 1000;

  float coloumbsConsumedSinceLastUpdate;

  // Calculate remaining battery charge in Coulombs based on current and time
  if (!PowerManager::busPowered) {
    coloumbsConsumedSinceLastUpdate =
        (currentCurrent / 1000) *
        ((pdTICKS_TO_MS(xTaskGetTickCount() - lastPowerStateUpdate)) / 1000.0);

    // Update coloumbs remaining
    coloumbsRemaining -= coloumbsConsumedSinceLastUpdate;
  }
  float chargeState;

  // If current flow is close enough to reference, get battery charge state via
  // voltage curve
  if (!PowerManager::busPowered) {
    if ((currentCurrent > (referenceCurrentMa * 0.6)) &&
        (currentCurrent < (referenceCurrentMa * 1.4))) {
      // Get battery charge state from voltage curve
      chargeState = getBatteryVoltageChargePercent();
      ESP_LOGD(TAG, "Charge state from voltage: %f", chargeState);
    } else {
      // Estimate battery charge state from charge consumption
      float oldChargeState = lastSOC[latestSoCIndex];
      chargeState = oldChargeState -
                    ((coloumbsConsumedSinceLastUpdate /
                      PowerParameters::Battery::CELL_CHARGE_FULL_COLOUMB) *
                     100);
      ESP_LOGD(TAG, "Charge state from current: %f", chargeState);
    }
  } else {
    // If we are charging, we can't estimate the charge state based on current
    // consumption
    chargeState = getBatteryVoltageChargePercent();
    ESP_LOGD(TAG, "Charge state from voltage (USB powered): %f", chargeState);
  }

  // Bound value between 0 and 100
  chargeState = chargeState > 100 ? 100 : chargeState;
  chargeState = chargeState < 0 ? 0 : chargeState;

  addSoCSample(chargeState);

  for (int i = 0; i < PowerParameters::Battery::AVERAGING_SAMPLES; i++) {
    ESP_LOGD(TAG, "SoC[%d]: %f", i, lastSOC[i]);
  }

  // Update percentage remaining based on charge state average
  float sampleSum = 0;
  for (int i = 0; i < PowerParameters::Battery::AVERAGING_SAMPLES; i++) {
    sampleSum += lastSOC[i];
  }
  PowerManager::percentRemaining =
      sampleSum / PowerParameters::Battery::AVERAGING_SAMPLES;

  // Update last update time
  PowerManager::lastPowerStateUpdate = xTaskGetTickCount();

  // Update the available current (changes based on battery state of charge)
  PowerManager::powerScheduler->recalculateCurrentBudgets();
  ESP_LOGD(TAG, "Current: %f mA, Charge: %f Coulombs, %f %%", currentCurrent,
           coloumbsRemaining, percentRemaining);

  return;
}

float PowerManager::getMax3V3Current() {
  // Conversion from Thesis
  float u_bat = getBatteryVoltage();
  float i_bat = PowerParameters::Battery::CELL_CURRENT_1C_MA;
  float eta = PowerParameters::BUCK_BOOST_EFFICIENCY;
  constexpr float u_3v3 = 3.3;
  return (u_bat * i_bat * eta) / u_3v3;
}

void PowerManager::addSoCSample(float soc) {
  PowerMutex lock(powerMutex);
  if (!lock.isLocked()) {
    ESP_LOGE(TAG, "Could not take power to add SoC sample");
    return;
  }
  latestSoCIndex =
      (latestSoCIndex + 1) % PowerParameters::Battery::AVERAGING_SAMPLES;
  lastSOC[latestSoCIndex] = soc;
}

void PowerManager::initPowerState(void) {
  // Initialize the power state
  lastPowerStateUpdate = xTaskGetTickCount();
  // TODO: Get initial battery charge state based on voltage, set coloumbs based
  // on that
  const float initialChargePercentages = getBatteryVoltageChargePercent();
  for (int i = 0; i < PowerParameters::Battery::AVERAGING_SAMPLES; i++) {
    lastSOC[i] = initialChargePercentages;
  }
  coloumbsRemaining = initialChargePercentages / 100 *
                      PowerParameters::Battery::CELL_CHARGE_FULL_COLOUMB;
  constexpr float fullColoumbs =
      PowerParameters::Battery::CELL_CHARGE_FULL_COLOUMB;
  percentRemaining = initialChargePercentages;
  // Set up flags and pins for them
  pinMode(PowerParameters::PinConfig::VUSB_SENS, INPUT_PULLUP);
  pinMode(PowerParameters::PinConfig::BAT_CHG_STAT, INPUT_PULLUP);
  busPowered = digitalRead(PowerParameters::PinConfig::VUSB_SENS);
  chargingState = digitalRead(PowerParameters::PinConfig::BAT_CHG_STAT);
}

void PowerManager::dumpPowerStatistics() {
  Serial.printf("======== Dezibot Power Statistics ========\r\n");
  Serial.printf("Current: %f mA\r\n", PowerManager::getCurrentCurrent());
  Serial.printf("Battery Voltage: %f V\r\n", PowerManager::getBatteryVoltage());
  Serial.printf("Battery Charge: %f %%\r\n", PowerManager::getBatteryChargePercent());
  Serial.printf("Battery Charge: %f Coulombs\r\n",
                PowerManager::getBatteryChargeCoulombs());
  Serial.printf("Max 3.3V Current in this state (1C, 2C): %f mA, %f mA \r\n",
                PowerManager::getMax3V3Current(), PowerManager::getMax3V3Current() * 2);
  Serial.printf("=========================================\r\n");
}

void PowerManager::dumpConsumerStatistics() {
  Serial.printf("======== Dezibot Consumer Statistics ========\r\n");
  Serial.printf("ESP: %f mA\r\n", PowerManager::getConsumerCurrent(
                                      PowerParameters::PowerConsumers::ESP));
  Serial.printf("WIFI: %f mA\r\n", PowerManager::getConsumerCurrent(
                                       PowerParameters::PowerConsumers::WIFI));
  Serial.printf("LED_RGB_TOP_LEFT: %f mA\r\n",
                PowerManager::getConsumerCurrent(
                    PowerParameters::PowerConsumers::LED_RGB_TOP_LEFT));
  Serial.printf("LED_RGB_TOP_RIGHT: %f mA\r\n",
                PowerManager::getConsumerCurrent(
                    PowerParameters::PowerConsumers::LED_RGB_TOP_RIGHT));
  Serial.printf("LED_RGB_BOTTOM: %f mA\r\n",
                PowerManager::getConsumerCurrent(
                    PowerParameters::PowerConsumers::LED_RGB_BOTTOM));
  Serial.printf(
      "RGBW_SENSOR: %f mA\r\n",
      PowerManager::getConsumerCurrent(PowerParameters::PowerConsumers::RGBW_SENSOR));
  Serial.printf("LED_IR_BOTTOM: %f mA\r\n",
                PowerManager::getConsumerCurrent(
                    PowerParameters::PowerConsumers::LED_IR_BOTTOM));
  Serial.printf(
      "LED_IR_FRONT: %f mA\r\n",
      PowerManager::getConsumerCurrent(PowerParameters::PowerConsumers::LED_IR_FRONT));
  Serial.printf(
      "PT_IR: %f mA\r\n",
      PowerManager::getConsumerCurrent(PowerParameters::PowerConsumers::PT_IR));
  Serial.printf(
      "PT_DL: %f mA\r\n",
      PowerManager::getConsumerCurrent(PowerParameters::PowerConsumers::PT_DL));
  Serial.printf(
      "LED_UV: %f mA\r\n",
      PowerManager::getConsumerCurrent(PowerParameters::PowerConsumers::LED_UV));
  Serial.printf(
      "DISPLAY_OLED: %f mA\r\n",
      PowerManager::getConsumerCurrent(PowerParameters::PowerConsumers::DISPLAY_OLED));
  Serial.printf(
      "MOTOR_LEFT: %f mA\r\n",
      PowerManager::getConsumerCurrent(PowerParameters::PowerConsumers::MOTOR_LEFT));
  Serial.printf(
      "MOTOR_RIGHT: %f mA\r\n",
      PowerManager::getConsumerCurrent(PowerParameters::PowerConsumers::MOTOR_RIGHT));
  Serial.printf("IMU: %f mA\r\n", PowerManager::getConsumerCurrent(
                                      PowerParameters::PowerConsumers::IMU));
  Serial.printf("=============================================\r\n");
}

bool PowerManager::isUSBPowered() { return busPowered; }

bool PowerManager::isBatteryPowered() { return !busPowered; }

bool PowerManager::isBatteryCharging() { return chargingState && busPowered; }

bool PowerManager::isBatteryDischarging() { return !chargingState && !busPowered; }

bool PowerManager::isBatteryFullyCharged() { return !chargingState && busPowered; }
int PowerManager::latestSoCIndex = 0;
float PowerManager::lastSOC[PowerParameters::Battery::AVERAGING_SAMPLES] = {0};
TickType_t PowerManager::lastPowerStateUpdate = 0;
float PowerManager::coloumbsRemaining =
    PowerParameters::Battery::CELL_CHARGE_FULL_COLOUMB;
float PowerManager::percentRemaining = 100.0;
PowerScheduler *PowerManager::powerScheduler = nullptr;

bool PowerManager::busPowered = false;

bool PowerManager::chargingState = false;

float PowerManager::fullVoltageOffset = 0;

PowerManager::PowerManager() {}

PowerManager::~PowerManager() {}