/**
 * @file PowerScheduler.cpp
 * @author Phillip Kühne
 * @brief The actual power scheduler class, which keeps track of the power
 * budget and allocates power to different components.
 * @version 0.1
 * @date 2024-12-21
 *
 * @copyright (c) 2024
 *
 */

#include "PowerScheduler.h"
#include "PowerManager.h"

bool PowerScheduler::tryAccquireCurrentAllowance(
    PowerParameters::PowerConsumers consumer, float neededCurrent,
    uint16_t requestedDurationMs) {
      float existingConsumption = getConsumerCurrent(consumer);
      const bool currentAvailableBelowLimit =
      this->freeLimitCurrentBudget + existingConsumption > 0;
      const bool currentAvailableBelowMaximum =
      this->freeMaximumCurrentBudget + existingConsumption >= neededCurrent;
      const bool currentIsInsignificant = neededCurrent < 1;
      if (currentIsInsignificant ||
        (currentAvailableBelowLimit && currentAvailableBelowMaximum)) {
    if (existingConsumption > 0) {
      releaseCurrent(consumer);
    }
    PowerSchedulerMutex lock(powerSchedulerMutex);
    if (lock.isLocked() == false) {
      ESP_LOGE(TAG, "Failed to Acquire PowerScheduler Mutex during Current Allocation");
      return false;
    }
    this->currentAllowances.push_back(PowerScheduler::CurrentAllowance{
        .consumer = consumer,
        .maxSlackTimeMs = 0,
        .requestedDurationMs = requestedDurationMs,
        .taskHandle = xTaskGetCurrentTaskHandle(),
        .neededCurrent = neededCurrent,
        .requestedAt = xTaskGetTickCount(),
        .grantedAt = xTaskGetTickCount(),
        .granted = true});
    this->recalculateCurrentBudgets();
    lock.unlock();
    return true;
  } else {
    ESP_LOGI(TAG,
             "Task %p denied %f mA of power;"
             "currently allocated: %f mA;"
             "total available (2C): %f mA",
             xTaskGetCurrentTaskHandle(), neededCurrent, getCurrentCurrent(),
             maximumCurrent);
    return false;
  }
}

void PowerScheduler::releaseCurrent(PowerParameters::PowerConsumers consumer) {
  PowerSchedulerMutex lock(powerSchedulerMutex);
  if (lock.isLocked() == false) {
    return;
  }
  for (auto it = currentAllowances.begin(); it != currentAllowances.end();
       ++it) {
    if (it->consumer == consumer) {
      currentAllowances.erase(it);
      break;
    }
  }
  lock.unlock();
  recalculateCurrentBudgets();
  // Check if there are tasks waiting for power
  checkWaitingTasks();
}

bool PowerScheduler::waitForCurrentAllowance(
    PowerParameters::PowerConsumers consumer, float neededCurrent,
    uint16_t maxSlackTimeMs, uint16_t requestedDurationMs) {
  if (tryAccquireCurrentAllowance(consumer, neededCurrent,
                                  requestedDurationMs)) {
    return true;
  } else {

    // Suspend the task while waiting for power to be available
    TaskHandle_t currentTask = xTaskGetCurrentTaskHandle();
    TickType_t initialTickCount = xTaskGetTickCount();
    PowerScheduler::CurrentAllowance newAllowance = {
        .consumer = consumer,
        .maxSlackTimeMs = maxSlackTimeMs,
        .requestedDurationMs = requestedDurationMs,
        .taskHandle = currentTask,
        .neededCurrent = neededCurrent,
        .requestedAt = initialTickCount,
        .granted = false};
    PowerSchedulerMutex lock(powerSchedulerMutex);
    if (lock.isLocked() == false) {
      return false;
    }
    this->currentAllowances.push_back(newAllowance);
    lock.unlock();
    uint32_t notificationValue;
    BaseType_t notificationStatus = xTaskNotifyWait(
        0, 0, &notificationValue, pdMS_TO_TICKS(maxSlackTimeMs));
    // Code below will be executed after the task is woken up
    while (notificationStatus == pdPASS) {
      if (notificationValue ==
          PowerScheduler::PowerWakeupReasons::POWER_AVAILABLE) {
        // We were woken up because new power is available, check if it is
        // enough
        // Exclude existing consumption from the same consumer, as it will be
        // gone if this is granted
        float existingConsumption = getConsumerCurrent(consumer);
        const bool currentAvailableBelowLimit =
            this->freeLimitCurrentBudget + existingConsumption > 0;
        const bool currentAvailableBelowMaximum =
            this->freeMaximumCurrentBudget + existingConsumption >=
            neededCurrent;
        const bool currentIsInsignificant = neededCurrent < 1;
        if (currentIsInsignificant ||
            (currentAvailableBelowLimit && currentAvailableBelowMaximum)) {
          PowerSchedulerMutex lock(powerSchedulerMutex);
          if (lock.isLocked() == false) {
            return false;
          }
          if (existingConsumption > 0) {
            releaseCurrent(consumer);
          }
          this->currentAllowances.push_back(
              {.consumer = consumer,
               .maxSlackTimeMs = 0,
               .requestedDurationMs = requestedDurationMs,
               .taskHandle = xTaskGetCurrentTaskHandle(),
               .neededCurrent = neededCurrent,
               .requestedAt = initialTickCount,
               .grantedAt = xTaskGetTickCount(),
               .granted = true});
          ESP_LOGV(TAG, "%d mA granted to consumer %d after %d ms",
                   neededCurrent, static_cast<int>(consumer),
                   xTaskGetTickCount() - initialTickCount);
          return true;
        } else {
          // Still not enough power available for us. Wait the remaining ticks.
          xTaskNotifyWait(0, 0, &notificationValue,
                          pdMS_TO_TICKS(maxSlackTimeMs) -
                              (xTaskGetTickCount() - initialTickCount));
        }
      }
    }
    if (notificationStatus == pdFALSE) {
      // We waited long enough...
      // Remove the task from the list of waiting tasks
      PowerSchedulerMutex lock(powerSchedulerMutex);
      if (lock.isLocked() == false) {
        return false;
      }
      for (auto it = currentAllowances.begin(); it != currentAllowances.end();
           ++it) {
        if (it->consumer == consumer && it->requestedAt == initialTickCount) {
          currentAllowances.erase(it);
          break;
        }
      }
      ESP_LOGW(TAG,
               "Task %p timed out waiting for %f mA of power;"
               "currently allocated: %f mA;"
               "total available (2C): %f mA",
               xTaskGetCurrentTaskHandle(), neededCurrent, getCurrentCurrent(),
               this->maximumCurrent);
      return false;
    } else {
      // Should be impossible to reach
      throw "Reached impossible state";
    }
  }
}

void PowerScheduler::checkWaitingTasks(void) {
  // If there are requested allowances, try to grant the one expiring next
  if (this->currentAllowances.size() > 0) {

    PowerScheduler::CurrentAllowance *nextAllowance =
        getNextExpiringAllowance();
    if (nextAllowance != nullptr) {
      xTaskNotify(nextAllowance->taskHandle,
                  PowerScheduler::PowerWakeupReasons::POWER_AVAILABLE,
                  eSetValueWithOverwrite);
    }
  }
}

void PowerScheduler::recalculateCurrentBudgets(void) {

  // Get the respective maximums and subtract currently flowing currents
  ESP_LOGI(TAG, "Recalculating current budgets...");
  float tempFreeLimitCurrentBudget = PowerManager::getMax3V3Current();
  ESP_LOGI(TAG, "Got max 3V3 current: %.2f", tempFreeLimitCurrentBudget);
  float tempFreeMaximumCurrentBudget = PowerManager::getMax3V3Current() * 2;
  PowerSchedulerMutex lock(powerSchedulerMutex);
  if (lock.isLocked() == false) {
    return;
  }
  for (auto &allowance : currentAllowances) {
    if (allowance.granted) {
      tempFreeLimitCurrentBudget -= allowance.neededCurrent;
      tempFreeMaximumCurrentBudget -= allowance.neededCurrent;
    }
  }
  this->freeLimitCurrentBudget = tempFreeLimitCurrentBudget;
  this->freeMaximumCurrentBudget = tempFreeMaximumCurrentBudget;
  ESP_LOGV(TAG, "Current budgets recalculated: %d mA, %d mA",
           this->freeLimitCurrentBudget, this->freeMaximumCurrentBudget);
}

PowerScheduler::CurrentAllowance *
PowerScheduler::getCurrentAllowance(PowerParameters::PowerConsumers consumer) {
  PowerSchedulerMutex lock(powerSchedulerMutex);
  if (lock.isLocked() == false) {
    return nullptr;
  }
  for (auto &allowance : currentAllowances) {
    if (allowance.consumer == consumer) {
      return &allowance;
    }
  }
  return nullptr;
}
PowerScheduler::CurrentAllowance *
PowerScheduler::getCurrentAllowance(TaskHandle_t taskHandle) {
  PowerSchedulerMutex lock(powerSchedulerMutex);
  if (lock.isLocked() == false) {
    return nullptr;
  }
  for (auto &allowance : currentAllowances) {
    if (allowance.taskHandle == taskHandle) {
      return &allowance;
    }
  }
  return nullptr;
}
PowerScheduler::CurrentAllowance *
PowerScheduler::getNextExpiringAllowance(void) {
  PowerSchedulerMutex lock(powerSchedulerMutex);
  TickType_t minTicks = UINT32_MAX;
  CurrentAllowance *nextAllowance = nullptr;
  if (lock.isLocked() == false) {
    return nullptr;
  }
  for (auto &allowance : currentAllowances) {
    if (!(allowance.granted)) {
      TickType_t ticks =
          allowance.requestedAt + pdMS_TO_TICKS(allowance.maxSlackTimeMs);
      if (ticks < minTicks) {
        minTicks = ticks;
        nextAllowance = &allowance;
      }
    }
  }
  // Will be nullptr if no allowance was found
  return nextAllowance;
}

PowerScheduler &PowerScheduler::getPowerScheduler(float i_limit_ma,
                                                  float i_max_ma) {
  if (powerSchedulerInstance == nullptr) {
      powerSchedulerInstance = new PowerScheduler(i_limit_ma, i_max_ma);
    }
  return *powerSchedulerInstance;
}

float PowerScheduler::getCurrentCurrent(void) {
  float currentSum = 0;
  PowerSchedulerMutex lock(powerSchedulerMutex);
  if (lock.isLocked() == false) {
    return false;
  }
  for (auto &allowance : currentAllowances) {
    if (allowance.granted) {
      currentSum += allowance.neededCurrent;
    }
  }
  return currentSum;
}

float PowerScheduler::getFreeLimitCurrentBudget(void) {
  return this->freeLimitCurrentBudget;
}

float PowerScheduler::getFreeMaximumCurrentBudget(void) {
  return this->freeMaximumCurrentBudget;
}

float PowerScheduler::getConsumerCurrent(
    PowerParameters::PowerConsumers consumer) {
  PowerSchedulerMutex lock(powerSchedulerMutex);
  if (lock.isLocked() == false) {
    return false;
  }
  float currentSum = 0;
  for (auto &allowance : currentAllowances) {
    if (allowance.consumer == consumer && allowance.granted) {
      currentSum += allowance.neededCurrent;
    }
  }
  return currentSum;
}

PowerScheduler::PowerScheduler(float i_limit_ma, float i_max_ma) {
  this->limitCurrent = i_limit_ma;
  this->maximumCurrent = i_max_ma;
  this->currentAllowances = std::vector<CurrentAllowance>();
  this->powerSchedulerMutex = xSemaphoreCreateMutex();
}

PowerScheduler::~PowerScheduler() {
  if (powerSchedulerMutex != NULL) {
    vSemaphoreDelete(powerSchedulerMutex);
  }
}