#pragma once

#include <Arduino.h>

#include "config.h"
#include "fan_step_response.h"
#include "pid_autotuner.h"
#include "pid_controller.h"
#include "settings_store.h"
#include "tuning_store.h"

class ThermalController {
public:
  enum class HeaterBlock : uint8_t { None, Cutoff, Corner, Autotune, StepResp };

  ThermalController()
      : heatPi_(HEAT_PI_KP, HEAT_PI_KI, 0.0f, 0.0f, 100.0f),
        mixPi_(MIX_PI_KP, MIX_PI_KI, 0.0f, 0.0f, 255.0f),
        autotuner_(),
        stepresp_(),
        targetTempC_(TARGET_TEMP_C),
        heaterDutyPercent_(0.0f),
        heaterAllowancePercent_(100.0f),
        cornerSpreadC_(0.0f),
        lastMaxTempC_(0.0f),
        regulatingFanPwm_(0),
        fanPwm_(0),
        tuningLoaded_(false),
        fanIdleOverride_(false),
        sensorWarmValid_(false),
        cutoffActive_(false),
        failSafeActive_(true),
        heaterCycleStartMs_(0),
        lastHeaterUpdateMs_(0),
        heaterOn_(false),
        lastAvgTempC_(0.0f),
        heaterBlock_(HeaterBlock::None),
        fanTestActive_(false),
        fanTestPwm_(0),
        fanTestEndMs_(0) {
    mixPi_.setSetpoint(0.0f);
  }

  void begin() {
    pinMode(FAN_PIN, OUTPUT);
    pinMode(HEATER_PIN, OUTPUT);
    digitalWrite(HEATER_PIN, LOW);
    writeFan(0);

    heatPi_.setSetpoint(targetTempC_);
    heatPi_.reset();
    mixPi_.reset();
    heaterCycleStartMs_ = millis();
    lastHeaterUpdateMs_ = 0;
    failSafeActive_ = true;
    cornerSpreadC_ = 0.0f;
    lastMaxTempC_ = 0.0f;
    sensorWarmValid_ = false;
    heaterAllowancePercent_ = 100.0f;
    if (isIdle()) {
      forceHeaterOff();
    }
    applyFan(millis());

    TuningData stored;
    if (tuningLoad(stored)) {
      applyTuning(stored);
      Serial.println(F("Loaded learned PI from EEPROM"));
      printTuning();
    }

    SettingsData settings;
    if (settingsLoad(settings) && settings.targetC >= TARGET_MIN_C &&
        settings.targetC <= TARGET_MAX_C) {
      setTarget(settings.targetC, false);
    }
  }

  void applyTuning(const TuningData &data) {
    heatPi_.setTunings(data.heatKp, data.heatKi, 0.0f);
    mixPi_.setTunings(data.mixKp, data.mixKi, 0.0f);
    tuningLoaded_ = true;
  }

  void clearTuning() {
    tuningLoaded_ = false;
    heatPi_.setTunings(HEAT_PI_KP, HEAT_PI_KI, 0.0f);
    mixPi_.setTunings(MIX_PI_KP, MIX_PI_KI, 0.0f);
    tuningClear();
    heatPi_.reset();
    mixPi_.reset();
    Serial.println(F("PI reset to defaults"));
  }

  void setMixTunings(float kp, float ki) {
    mixPi_.setTunings(kp, ki, 0.0f);
    mixPi_.reset();
    Serial.print(F("Mix PI Kp="));
    Serial.print(kp, 3);
    Serial.print(F(" Ki="));
    Serial.println(ki, 4);
  }

  void resetMixTunings() {
    mixPi_.setTunings(MIX_PI_KP, MIX_PI_KI, 0.0f);
    mixPi_.reset();
    Serial.println(F("Mix PI reset to defaults"));
  }

  void printTuning() const {
    Serial.print(F("Heat PI Kp="));
    Serial.print(heatPi_.kp(), 3);
    Serial.print(F(" Ki="));
    Serial.print(heatPi_.ki(), 4);
    Serial.print(F("  Mix PI Kp="));
    Serial.print(mixPi_.kp(), 3);
    Serial.print(F(" Ki="));
    Serial.print(mixPi_.ki(), 4);
    Serial.print(F("  tuned="));
    Serial.println(tuningLoaded_ ? F("yes") : F("no"));
  }

  void saveTuningToEeprom() {
    TuningData data;
    data.magic = TUNING_MAGIC;
    data.heatKp = heatPi_.kp();
    data.heatKi = heatPi_.ki();
    data.mixKp = mixPi_.kp();
    data.mixKi = mixPi_.ki();
    tuningSave(data);
    tuningLoaded_ = true;
    Serial.println(F("Saved PI to EEPROM"));
  }

  float heatKp() const { return heatPi_.kp(); }
  float heatKi() const { return heatPi_.ki(); }
  float mixKp() const { return mixPi_.kp(); }
  float mixKi() const { return mixPi_.ki(); }

  bool isTuningLoaded() const { return tuningLoaded_; }

  bool startAutotune(float setpointC) {
    if (autotuner_.isActive() || stepresp_.isActive()) {
      return false;
    }
    cutoffActive_ = false;
    heatPi_.reset();
    mixPi_.reset();
    return autotuner_.start(setpointC);
  }

  void stopAutotune() { autotuner_.abort(); }

  bool isAutotuning() const { return autotuner_.isActive(); }

  uint32_t autotuneElapsedMs(uint32_t nowMs) const { return autotuner_.elapsedMs(nowMs); }

  const char *autotunePhaseName() const { return autotuner_.phaseName(); }

  uint8_t autotuneCycleCount() const { return autotuner_.cycleCount(); }

  uint8_t autotunePeriodCount() const { return autotuner_.periodCount(); }

  float autotunePreheatTargetC() const { return autotuner_.preheatTargetC(); }

  bool startStepResponse(float targetC, float heaterPct) {
    if (autotuner_.isActive() || stepresp_.isActive()) {
      return false;
    }
    stopFanTest();
    cutoffActive_ = false;
    heatPi_.reset();
    mixPi_.reset();
    setTarget(targetC, false);
    if (!stepresp_.start(targetC, heaterPct)) {
      return false;
    }
    writeFan(0);
    return true;
  }

  void stopStepResponse() {
    stepresp_.abort();
    writeFan(0);
  }

  bool isStepResponseActive() const { return stepresp_.isActive(); }

  uint32_t stepResponseElapsedMs(uint32_t nowMs) const { return stepresp_.elapsedMs(nowMs); }

  const char *stepResponsePhaseName() const { return stepresp_.phaseName(); }

  uint8_t stepResponseStepIndex() const { return stepresp_.stepIndex(); }

  uint8_t stepResponseStepCount() const { return stepresp_.stepCount(); }

  float stepResponseHeaterPct() const { return stepresp_.heaterPct(); }

  void logStepResponseIfDue(const float *sensorTemps, const bool *sensorValid, uint8_t sensorCount,
                            float avgTempC, float minTempC, float maxTempC, float spreadC,
                            uint32_t nowMs) {
    stepresp_.logIfDue(sensorTemps, sensorValid, sensorCount, avgTempC, minTempC, maxTempC,
                       spreadC, nowMs);
  }

  bool commitAutotuneIfDone() {
    if (autotuner_.phase() != PidAutotuner::Phase::Done) {
      return false;
    }

    heatPi_.setTunings(autotuner_.resultKp(), autotuner_.resultKi(), 0.0f);
    heatPi_.reset();

    TuningData data;
    data.magic = TUNING_MAGIC;
    data.heatKp = autotuner_.resultKp();
    data.heatKi = autotuner_.resultKi();
    data.mixKp = mixPi_.kp();
    data.mixKi = mixPi_.ki();
    tuningSave(data);
    tuningLoaded_ = true;
    autotuner_.reset();
    Serial.println(F("Saved heat PI to EEPROM (mix PI unchanged)"));
    return true;
  }

  void setTarget(float targetC, bool persist = true) {
    targetTempC_ = targetC;
    heatPi_.setSetpoint(targetC);
    heatPi_.reset();
    mixPi_.reset();
    cutoffActive_ = false;
    if (targetC > 0.0f) {
      fanIdleOverride_ = false;
    } else {
      forceHeaterOff();
      fanIdleOverride_ = false;
    }
    applyFan(millis());
    if (persist && targetC >= TARGET_MIN_C && targetC <= TARGET_MAX_C) {
      settingsSaveTarget(targetC);
    }
  }

  void noteSensorMax(float maxTempC) {
    lastMaxTempC_ = maxTempC;
    sensorWarmValid_ = true;
  }

  bool setFanOff() {
    if (!isIdle()) {
      return false;
    }
    fanIdleOverride_ = false;
    applyFan(millis());
    return true;
  }

  void setFanIdle() {
    if (!isIdle()) {
      return;
    }
    if (!sensorWarmValid_ || lastMaxTempC_ >= IDLE_AUTO_FAN_OFF_TEMP_C) {
      return;
    }
    fanIdleOverride_ = true;
    applyFan(millis());
  }

  bool isFanOff() const {
    return isIdle() && sensorWarmValid_ && lastMaxTempC_ < IDLE_AUTO_FAN_OFF_TEMP_C &&
           !fanIdleOverride_;
  }

  bool isIdleCooling() const {
    return isIdle() &&
           (!sensorWarmValid_ || lastMaxTempC_ >= IDLE_AUTO_FAN_OFF_TEMP_C);
  }

  float target() const { return targetTempC_; }

  bool isIdle() const { return targetTempC_ <= 0.0f; }

  float cutoffThreshold() const {
    if (isIdle()) {
      return INFINITY;
    }
    return EMERGENCY_MAX_TEMP_C;
  }

  bool isCutoffActive() const { return cutoffActive_; }

  bool isFailSafeActive() const { return failSafeActive_; }

  float heaterDutyPercent() const { return heaterDutyPercent_; }

  float heaterAllowance() const { return heaterAllowancePercent_; }

  float cornerSpread() const { return cornerSpreadC_; }

  uint8_t fanPwm() const { return fanPwm_; }

  bool isHeaterSsrOn() const { return heaterOn_; }

  bool isFanTestActive(uint32_t nowMs) const {
    return fanTestActive_ && nowMs < fanTestEndMs_;
  }

  bool startFanTest(uint8_t pwm, uint32_t nowMs, uint32_t durationMs = 15000) {
    fanTestPwm_ = pwm;
    fanTestEndMs_ = nowMs + durationMs;
    fanTestActive_ = true;
    writeFan(pwm);
    return true;
  }

  void stopFanTest() { fanTestActive_ = false; }

  float maxHeatStopAt(float avgTempC) const {
    if (isIdle()) {
      return INFINITY;
    }
    return maxHeatStopTemp(avgTempC);
  }

  const char *heaterBlockReason() const {
    switch (heaterBlock_) {
    case HeaterBlock::Cutoff:
      return "cutoff";
    case HeaterBlock::Corner:
      return "corner";
    case HeaterBlock::Autotune:
      return "autotune";
    case HeaterBlock::StepResp:
      return "stepresp";
    default:
      return "none";
    }
  }

  const char *regulatingModeName() const {
    return tuningLoaded_ ? "regulating" : "manual";
  }

  void update(float avgTempC, float maxTempC, float cornerSpreadC, uint32_t nowMs) {
    failSafeActive_ = false;
    noteSensorMax(maxTempC);
    lastAvgTempC_ = avgTempC;
    heaterBlock_ = HeaterBlock::None;

    cornerSpreadC_ =
        SPREAD_EMA_ALPHA * cornerSpreadC_ +
        (1.0f - SPREAD_EMA_ALPHA) * cornerSpreadC;

    if (stepresp_.isActive()) {
      updateStepResponse(avgTempC, maxTempC, nowMs);
      return;
    }

    if (autotuner_.isActive()) {
      updateAutotune(avgTempC, maxTempC, nowMs);
      return;
    }

    if (isIdle()) {
      forceHeaterOff();
      cutoffActive_ = false;
      heatPi_.reset();
      mixPi_.reset();
      lastHeaterUpdateMs_ = nowMs;
      applyFan(nowMs);
      return;
    }

    updateRegulating(avgTempC, maxTempC, nowMs);
    lastHeaterUpdateMs_ = nowMs;
    applyHeaterBurst(nowMs);
    writeFan(regulatingFanPwm_);
  }

  void enterFailSafe() {
    failSafeActive_ = true;
    cutoffActive_ = false;
    forceHeaterOff();
    applyFan(millis());
    heatPi_.reset();
    mixPi_.reset();
    autotuner_.abort();
    stepresp_.abort();
  }

  void forceHeaterOff() {
    heaterDutyPercent_ = 0.0f;
    heaterAllowancePercent_ = 0.0f;
    heaterOn_ = false;
    digitalWrite(HEATER_PIN, LOW);
  }

  void writeFan(uint8_t pwm) {
    fanPwm_ = pwm;
    pinMode(FAN_PIN, OUTPUT);

    if (FAN_PWM_INVERT) {
      if (pwm == 0) {
        digitalWrite(FAN_PIN, HIGH);
        return;
      }
      if (pwm >= 254) {
        digitalWrite(FAN_PIN, LOW);
        return;
      }
      analogWrite(FAN_PIN, static_cast<uint8_t>(255 - pwm));
      return;
    }

    if (pwm == 0) {
      digitalWrite(FAN_PIN, LOW);
      return;
    }
    if (pwm >= 254) {
      digitalWrite(FAN_PIN, HIGH);
      return;
    }
    analogWrite(FAN_PIN, pwm);
  }

private:
  void updateStepResponse(float avgTempC, float maxTempC, uint32_t nowMs) {
    float duty = 0.0f;
    uint8_t fan = 0;
    stepresp_.update(avgTempC, maxTempC, cornerSpreadC_, nowMs, duty, fan);

    heaterDutyPercent_ = duty;
    heaterAllowancePercent_ = duty;
    heaterBlock_ = duty > 0.0f ? HeaterBlock::StepResp : HeaterBlock::None;
    applyHeaterBurst(nowMs);
    writeFan(fan);
    lastHeaterUpdateMs_ = nowMs;

    if (stepresp_.phase() == FanStepResponse::Phase::Done ||
        stepresp_.phase() == FanStepResponse::Phase::Failed) {
      stepresp_.reset();
    }
  }

  void updateAutotune(float avgTempC, float maxTempC, uint32_t nowMs) {
    float duty = 0.0f;
    uint8_t fan = AUTOTUNE_PREHEAT_FAN_PWM;
    autotuner_.update(avgTempC, maxTempC, cornerSpreadC_, nowMs, duty, fan);

    heaterDutyPercent_ = duty;
    heaterAllowancePercent_ = duty;
    heaterOn_ = duty >= 50.0f;
    digitalWrite(HEATER_PIN, heaterOn_ ? HIGH : LOW);
    heaterBlock_ = duty > 0.0f ? HeaterBlock::None : HeaterBlock::Autotune;
    writeFan(fan);
    lastHeaterUpdateMs_ = nowMs;
    commitAutotuneIfDone();
  }

  void updateRegulating(float avgTempC, float maxTempC, uint32_t nowMs) {
    if (maxTempC >= EMERGENCY_MAX_TEMP_C) {
      cutoffActive_ = true;
      heaterDutyPercent_ = 0.0f;
      heaterAllowancePercent_ = 0.0f;
      regulatingFanPwm_ = FAN_MAX_PWM;
      heaterBlock_ = HeaterBlock::Cutoff;
      heatPi_.reset();
      mixPi_.reset();
      return;
    }

    if (cutoffActive_ && maxTempC < EMERGENCY_MAX_TEMP_C - 5.0f) {
      cutoffActive_ = false;
      heatPi_.reset();
      mixPi_.reset();
    }

    if (cutoffActive_) {
      heaterBlock_ = HeaterBlock::Cutoff;
      regulatingFanPwm_ = FAN_MAX_PWM;
      return;
    }

    heaterAllowancePercent_ = allowanceFromMaxCorner(maxTempC, avgTempC);

    float duty = heatPi_.compute(avgTempC, nowMs);
    const float below = targetTempC_ - avgTempC;
    if (below > 8.0f) {
      const float floor = below > 15.0f ? 75.0f : 60.0f;
      if (duty < floor) {
        duty = floor;
      }
    }
    const float maxDuty = heaterMaxDuty(avgTempC);
    if (duty > maxDuty) {
      duty = maxDuty;
    }
    if (duty > heaterAllowancePercent_) {
      duty = heaterAllowancePercent_;
      if (heaterAllowancePercent_ < 100.0f) {
        heaterBlock_ = HeaterBlock::Corner;
      }
    }
    heaterDutyPercent_ = applyHeaterRamp(duty, avgTempC, nowMs);

    const float mixInput = SPREAD_TARGET_C - cornerSpreadC_;
    float fanOut = mixPi_.compute(mixInput, nowMs);
    uint8_t fanPwm = static_cast<uint8_t>(fanOut + 0.5f);
    if (avgTempC < targetTempC_ - HEAT_UP_BAND_C && fanPwm > FAN_COLD_CAP_PWM) {
      fanPwm = FAN_COLD_CAP_PWM;
    }
    regulatingFanPwm_ = fanPwm;
  }

  static float clampPercent(float value) {
    if (value < 0.0f) {
      return 0.0f;
    }
    if (value > 100.0f) {
      return 100.0f;
    }
    return value;
  }

  bool shouldLimitMaxCorner(float avgTempC) const {
    return avgTempC >= targetTempC_ - CORNER_LIMIT_BAND_C;
  }

  float maxHeatStopTemp(float avgTempC) const {
    if (!shouldLimitMaxCorner(avgTempC)) {
      return EMERGENCY_MAX_TEMP_C;
    }
    return EMERGENCY_MAX_TEMP_C - CORNER_STOP_MARGIN_C;
  }

  float allowanceFromMaxCorner(float maxTempC, float avgTempC) const {
    if (!shouldLimitMaxCorner(avgTempC)) {
      return 100.0f;
    }

    const float stopAt = maxHeatStopTemp(avgTempC);
    if (maxTempC >= stopAt) {
      return 0.0f;
    }

    const float headroom = stopAt - maxTempC;
    if (headroom >= MAX_TEMP_HEADROOM_C) {
      return 100.0f;
    }

    return clampPercent((headroom / MAX_TEMP_HEADROOM_C) * 100.0f);
  }

  float heaterMaxDuty(float avgTempC) const {
    if (avgTempC >= targetTempC_) {
      return HEATER_MAX_DUTY_NEAR;
    }

    const float below = targetTempC_ - avgTempC;
    if (below >= HEATER_COLD_BELOW_C) {
      return HEATER_MAX_DUTY_COLD;
    }
    if (below >= HEATER_WARM_BELOW_C) {
      return HEATER_MAX_DUTY_MID;
    }
    return HEATER_MAX_DUTY_NEAR;
  }

  float applyHeaterRamp(float requestedDuty, float avgTempC, uint32_t nowMs) {
    const float maxDuty = heaterMaxDuty(avgTempC);
    if (requestedDuty > maxDuty) {
      requestedDuty = maxDuty;
    }

    if (lastHeaterUpdateMs_ > 0 && requestedDuty > heaterDutyPercent_) {
      const float dt = static_cast<float>(nowMs - lastHeaterUpdateMs_) / 1000.0f;
      const float maxUp = heaterDutyPercent_ + HEATER_SLEW_UP_PER_S * dt;
      if (requestedDuty > maxUp) {
        requestedDuty = maxUp;
      }
    }

    return requestedDuty;
  }

  void applyHeaterBurst(uint32_t nowMs) {
    if (heaterDutyPercent_ <= 0.0f) {
      forceHeaterOff();
      return;
    }

    if (nowMs - heaterCycleStartMs_ >= HEATER_CYCLE_MS) {
      heaterCycleStartMs_ = nowMs;
    }

    const float onFraction = heaterDutyPercent_ / 100.0f;
    const uint32_t onTimeMs = static_cast<uint32_t>(HEATER_CYCLE_MS * onFraction);
    const bool shouldHeat = (nowMs - heaterCycleStartMs_) < onTimeMs;

    if (shouldHeat != heaterOn_) {
      heaterOn_ = shouldHeat;
      digitalWrite(HEATER_PIN, heaterOn_ ? HIGH : LOW);
    }
  }

  void applyFan(uint32_t nowMs) {
    if (fanTestActive_) {
      if (nowMs < fanTestEndMs_) {
        writeFan(fanTestPwm_);
        return;
      }
      fanTestActive_ = false;
    }

    if (isIdle()) {
      if (sensorWarmValid_ && lastMaxTempC_ >= IDLE_AUTO_FAN_OFF_TEMP_C) {
        writeFan(FAN_MAX_PWM);
      } else if (fanIdleOverride_) {
        writeFan(FAN_IDLE_PWM);
      } else {
        writeFan(0);
      }
      return;
    }

    if (failSafeActive_ || cutoffActive_) {
      writeFan(FAN_MAX_PWM);
      return;
    }
  }

  PidController heatPi_;
  PidController mixPi_;
  PidAutotuner autotuner_;
  FanStepResponse stepresp_;
  float targetTempC_;
  float heaterDutyPercent_;
  float heaterAllowancePercent_;
  float cornerSpreadC_;
  float lastMaxTempC_;
  uint8_t regulatingFanPwm_;
  uint8_t fanPwm_;
  bool tuningLoaded_;
  bool fanIdleOverride_;
  bool sensorWarmValid_;
  bool cutoffActive_;
  bool failSafeActive_;
  uint32_t heaterCycleStartMs_;
  uint32_t lastHeaterUpdateMs_;
  bool heaterOn_;
  float lastAvgTempC_;
  HeaterBlock heaterBlock_;
  bool fanTestActive_;
  uint8_t fanTestPwm_;
  uint32_t fanTestEndMs_;
};