#include <ESP8266WiFi.h>
#include <MQTT.h>
#include <AccelStepper.h>

#include<EEPROM.h>

WiFiClient net;
MQTTClient client;
AccelStepper stepper(AccelStepper::HALF4WIRE, D0, D2, D1, D3);

// Config
const char ssid[] = "";
const char wifi_pw[] = "";
const char mqtt_user[] = "";
const char mqtt_pw[] = "";
const char mqtt_addr[] = "";
const char client_name[] = "";

const int INIT_POS = 511;
const bool CLEAR_EEPROM = false;


// End Config

unsigned long lastMillis = 0;

// Backlash compensation and hold configuration
const int BACKLASH_STEPS = 400;              // Tune as needed for your linkage
const unsigned long HOLD_OUTPUTS_MS = 300;   // Hold time after moves to prevent settling

// Runtime variables for compensation and holding
int approachDir = 0;           // Desired final approach direction for current command (-1 or 1)
long prepTarget = 0;           // First phase target used to preload backlash
long finalMoveTarget = 0;      // Second phase target: the actual requested position
unsigned long holdUntil = 0;   // Timestamp when we should release outputs

struct Config {
  long openPosition;
  long closedPosition;
  long lastPosition;
  bool blindsOpen;
};


// States
const int SETUP = 0;
const int SLACK = 1;
const int OPENING_BLINDS = 2;
const int CLOSING_BLINDS = 3;
const int IDLE = 4;
const int ADJUSTING = 5;
const int HOLDING = 6;
const int OPENING_PREP = 7;
const int CLOSING_PREP = 8;

int state = SETUP;
int previousState = SETUP;

Config config;

void connect() {
  Serial.println("checking wifi...");
  while (WiFi.status() != WL_CONNECTED) {
    Serial.print(".");
    delay(1000);
  }

  Serial.println("\nconnecting...");
  while (!client.connect(client_name, mqtt_user, mqtt_pw)) {
    Serial.print(".");
    delay(1000);
  }

  Serial.println("\nconnected!");

  client.subscribe("/bedroom/blinds/2");
  client.subscribe("/bedroom/blinds/2/adjust");
}

void messageReceived(String &topic, String &payload) {
  Serial.println("incoming: " + topic + " - " + payload);

  if (topic.endsWith("adjust")) {
    int adjustment;
    sscanf(payload.c_str(), "%d", &adjustment);
    adjustClosedPosition(adjustment);
  } else {
    if (payload.indexOf("open") != -1) {
      openBlinds();
    } else if (payload.indexOf("close") != -1) {
      closeBlinds();
    } else if (payload.indexOf("reset") != -1) {
      Serial.println("Resetting current position and setting blinds state to closed");
      stepper.setCurrentPosition(0);
      config.blindsOpen = false;
    } else {
      Serial.println(payload);
    }
  }
}

void updateState(int newState) {
  previousState = state;
  state = newState;

  if(newState == IDLE) {
    config.lastPosition = stepper.currentPosition();

    Serial.print("Transitioning to idle (");
    Serial.print(config.lastPosition);
    Serial.println(")");
    EEPROM.put(0, config);
    EEPROM.commit();
  } else {
    Serial.print("Moving ");
    Serial.print(previousState);
    Serial.print(" -> ");
    Serial.println(newState);
  }
}

void openBlinds() {
  if (!config.blindsOpen && state == IDLE) {
    Serial.println("Opening Blinds");
    long baseTarget = config.openPosition;
    long delta = config.openPosition - config.closedPosition;
    int openDir = (delta > 0) ? 1 : ((delta < 0) ? -1 : -1); // default to -1 if equal (unlikely)
    approachDir = openDir;

    // Phase 1: move to preload point on the "closing" side of the final target
    prepTarget = baseTarget - (approachDir * BACKLASH_STEPS);
    finalMoveTarget = baseTarget;

    updateState(OPENING_PREP);
    stepper.enableOutputs();
    stepper.moveTo(prepTarget);
  }
}

void closeBlinds() {
  if (config.blindsOpen && state == IDLE) {
    Serial.println("Closing Blinds");
    long baseTarget = config.closedPosition;
    long delta = config.openPosition - config.closedPosition;
    int openDir = (delta > 0) ? 1 : ((delta < 0) ? -1 : -1);
    int closeDir = -openDir;
    approachDir = closeDir;

    // Phase 1: move to preload point on the "opening" side of the final target (opposite of close direction)
    prepTarget = baseTarget - (approachDir * BACKLASH_STEPS);
    finalMoveTarget = baseTarget;

    updateState(CLOSING_PREP);
    stepper.enableOutputs();
    stepper.moveTo(prepTarget);
  }
}

void adjustClosedPosition(int offset) {
  if(offset == 0 || state != IDLE) return;

  Serial.print("Adjust by: ");
  Serial.println(offset);

  if(config.blindsOpen) {
    config.openPosition += offset;
    stepper.moveTo(config.openPosition);

    Serial.print("New open position: ");
    Serial.println(config.openPosition);
  } else {
    config.closedPosition += offset;
    stepper.moveTo(config.closedPosition);

    Serial.print("New closed position: ");
    Serial.println(config.closedPosition);
  }

  updateState(ADJUSTING);
  stepper.enableOutputs();
}

void setup() {
  Serial.begin(115200);
  while(!Serial);

  WiFi.begin(ssid, wifi_pw);

  client.begin(mqtt_addr, 1883, net);
  client.onMessage(messageReceived);

  connect();

  EEPROM.begin(512);

  // Only clears if flag is set
  clearEeprom();

  if(EEPROM.read(INIT_POS) != 'T') {
    config.closedPosition = 0;
    config.openPosition = -12000;
    config.lastPosition = -3000;

    Serial.println("Save default config state");
    Serial.print("Open pos: ");
    Serial.println(config.openPosition);
    EEPROM.put(0, config);

    EEPROM.write(INIT_POS, 'T');
    if(!EEPROM.commit()) {
      Serial.println("ERROR WRITING CONFIG");
    }
  } else {
    Serial.println("Config already saved");
    EEPROM.get(0, config);
  }


  stepper.setMaxSpeed(600);
  stepper.setSpeed(100);
  stepper.setAcceleration(250);
}

void clearEeprom() {
  if(CLEAR_EEPROM) {
    for(int i = 0; i < 512; i++) {
      EEPROM.put(i, 0);
    }
    EEPROM.commit();
  }
}

void loop() {
  client.loop();

  if (!client.connected()) {
    connect();
  }

  switch (state) {
    case SETUP:
      setup_stepper();
      break;
    case IDLE:
      delay(10);  // <- fixes some issues with WiFi stability
      break;
    case SLACK:
      // Deprecated slack handling; transition to idle once any pending motion completes
      if (stepper.distanceToGo() == 0) {
        stepper.disableOutputs();
        updateState(IDLE);
      } else {
        stepper.run();
      }
      break;
    case OPENING_PREP:
    case CLOSING_PREP:
      if (stepper.distanceToGo() != 0) {
        stepper.run();
      } else {
         Serial.println("Open/Close Preparation complete");
        // Phase 2: approach target from a consistent direction to remove backlash
        if (state == OPENING_PREP) {
          updateState(OPENING_BLINDS);
        } else {
          updateState(CLOSING_BLINDS);
        }
        stepper.moveTo(finalMoveTarget);
      }
      break;
    case OPENING_BLINDS:
    case CLOSING_BLINDS:
      if (stepper.distanceToGo() != 0) {
        if (millis() - lastMillis > 2000) {
          Serial.print("Speed: ");
          Serial.print(stepper.speed());
          Serial.print(" Remaining Dist: ");
          Serial.println(stepper.distanceToGo());
          lastMillis = millis();
        }
        stepper.run();
      } else {
        bool open = state == OPENING_BLINDS;
        config.blindsOpen = open;
        sendBlindsUpdate(open);
        // Hold outputs briefly to prevent settling drift, then transition to IDLE
        holdUntil = millis() + HOLD_OUTPUTS_MS;
        updateState(HOLDING);
      }
      break;
    case HOLDING:
      if (millis() >= holdUntil) {
        Serial.println("Hold complete, transitioning to idle");
        stepper.disableOutputs();
        updateState(IDLE);
      }
      break;
    case ADJUSTING:
      if (millis() - lastMillis > 1000) {
        Serial.println("Adjusting...");
        lastMillis = millis();
      }

      stepper.run();

      if (stepper.distanceToGo() == 0) {
        Serial.print("Finished adjusting to: ");
        Serial.println(stepper.currentPosition());

        holdUntil = millis() + HOLD_OUTPUTS_MS;
        updateState(HOLDING);
      }

      break;
  }
}

void setup_stepper() {
  stepper.setCurrentPosition(config.lastPosition);
  updateState(IDLE);
}

void sendBlindsUpdate(bool isOpen) {
  client.publish("/bedroom/blinds/2/status", isOpen ? "true" : "false", true);
}