#include <elapsedMillis.h>
#include <pwmWrite.h>
#include <WiFi.h>
#include <PubSubClient.h>

// mqtt configuration variables
// WiFi
const char *ssid = "Wokwi-GUEST"; // Enter your Wi-Fi name
const char *password = "";  // Enter Wi-Fi password

// MQTT Broker
const char *mqtt_broker = "192.168.55.76";
const char *mqtt_username = "a";
const char *mqtt_password = "a";
const int mqtt_port = 1883;

WiFiClient espClient;
PubSubClient client(espClient);

// reload servo pin assignment
const int reloadServoPinM1 = 5;
const int reloadServoPinM2 = 7;
const int reloadServoPinM3 = 16;
const int reloadServoPinM4 = 18;

// dispense servo pin assignment
const int dispenseServoPinM1 = 14;
const int dispenseServoPinM2 = 12;
const int dispenseServoPinM3 = 10;
const int dispenseServoPinM4 = 46;

// reload sensor pin assignment
const int reloadSensorPinM1 = 4;
const int reloadSensorPinM2 = 6;
const int reloadSensorPinM3 = 15;
const int reloadSensorPinM4 = 17;

// dispense sensor pin assignment
const int dispenseSensorPinM1 = 13;
const int dispenseSensorPinM2 = 11;
const int dispenseSensorPinM3 = 9;
const int dispenseSensorPinM4 = 3;

const int posServoOpened = 180;
const int posServoClosed = 90;

const int clockPin = 21;  // SHCP
const int latchPin = 20;  // STCP
const int dataPin = 19;   // DS

enum StepperDirection {
  CounterClockWise, ClockWise
};

struct Stepper
{
  uint8_t step_pin;
  uint8_t dir_pin;
  StepperDirection direction;
  long target_steps;
  long cur_steps;
  bool enabled;
  elapsedMillis last_step_t;
};

class MultipleSteppersController {
    // Uses a 74HC595 IC to control a maximum of 4 steppers.
    // Each stepper uses 2 wires each.

  private:
    // IC mask is 0b00_00_00_00 , lsb first
    // 4 Steppers total, each having 2 pins
    Stepper steppers[4] = {};
    uint8_t mask = 0b00000000;

    // SPI pins
    int clockPin;
    int dataPin;
    int latchPin;

    void updateStepperMasks() {
      this->mask = 0b00000000;
      for (int i = 0; i < 4; i++) {
        if (
          !steppers[i].enabled ||
          steppers[i].target_steps - steppers[i].cur_steps == 0
        ) continue;

        bitWrite(mask, steppers[i].step_pin, true);
        steppers[i].cur_steps += 1;

        switch (steppers[i].direction) {
          case StepperDirection::CounterClockWise:
            bitWrite(this->mask, steppers[i].dir_pin, false);
            break;
          case StepperDirection::ClockWise:
            bitWrite(this->mask, steppers[i].dir_pin, true);
            break;
          default:
            break;
        }
      }
    }

    void triggerLatch() {
      digitalWrite(latchPin, HIGH);
      digitalWrite(latchPin, LOW);
      shiftOut(dataPin, clockPin, MSBFIRST, 0);
      digitalWrite(latchPin, HIGH);
      digitalWrite(latchPin, LOW);
    }

    void stepSteppers() {
      updateStepperMasks();
      shiftOut(dataPin, clockPin, MSBFIRST, mask);
      triggerLatch();
    }

  public:
    MultipleSteppersController(bool enabled_steppers_p[4], int spi[3]) {
      for (int i = 0; i < 4; i++) {
        int pin_selct = i * 2;
        if (enabled_steppers_p[i]) {
          steppers[i] = Stepper {
            // IC_MASK ^= (1 << stepper.step_pin); -> to step the mottor
            // IC_MASK ^= (1 << stepper.dir_pin); -> to direct the mottor;
            // all of this only if the stepper is enabled
            pin_selct + 1, pin_selct,
            StepperDirection::ClockWise,
            0, 0,
            true,
            0,
          };
        } else {
          steppers[i] = Stepper {
            pin_selct + 1, pin_selct,
            StepperDirection::ClockWise,
            0, 0,
            false,
            0
          };
        }
      }

      dataPin  = spi[0];
      clockPin = spi[1];
      latchPin = spi[2];
    }

    void enable() {
      pinMode(dataPin, OUTPUT);
      pinMode(clockPin, OUTPUT);
      pinMode(latchPin, OUTPUT);
    }

    void stepToN(uint8_t stepper, long steps) {
      // Invalid stepper indicies
      if (stepper > 3) return;
      // Check if stepper is enabled
      if (!steppers[stepper].enabled) return;

      if (steps < 0) {
        steps *= -1;
        steppers[stepper].direction = StepperDirection::CounterClockWise;
      } else if (steps == 0) {
        return;
      } else {
        steppers[stepper].direction = StepperDirection::ClockWise;
      }

      steppers[stepper].target_steps = steps;
      
      Serial.print("STEPS: ");
      Serial.println(steppers[stepper].cur_steps);

      if (steppers[stepper].cur_steps >= 200) {
        steppers[stepper].cur_steps = 0;
        steppers[stepper].target_steps = 0;
        stepSteppers();
        return;
      }

      if (steppers[stepper].last_step_t > 0) {
        steppers[stepper].last_step_t = 0;
        stepSteppers();
      }
    }

    void stepToNDegrees(uint8_t stepper, int degrees) {
      stepToN(stepper, degreesToSteps(degrees));
    }

    int degreesToSteps(float degrees_p) {
      float degs = 360.0 / 200;
      return degrees_p / degs;
    }

    int reloadAngles[10] = {0, 40, 80, 120, 160, 200, 240, 280, 320, 359};
    int dispenseAngles[9] = {80, 220, 260, 320, 340, 20, 60, 100, 140};

    void goTo(int mode, uint8_t stepper, int degrees) {
      if(mode == 0){
      stepToN(stepper, degreesToSteps(reloadAngles[degrees]));
      }
      else if(mode == 1){
      stepToN(stepper, degreesToSteps(dispenseAngles[degrees]));
      }
      else if(mode != 0 || mode!= 1){
        Serial.print("Mode not valid, chose 0 to reload or 1 to dispense");
      }
    }

    void drop(uint8_t stepper, int degrees) {
      stepToN(stepper, degreesToSteps(dispenseAngles[degrees]));
    }
};


bool enabled[4] = {true, true, true, true};
int spi[4] = {dataPin, clockPin, latchPin};

int inputDegrees1 = 0;
int inputDegrees2 = 0;
int inputDegrees3 = 0;
int inputDegrees4 = 0;

int mode1 = 0;
int mode2 = 0;
int mode3 = 0;
int mode4 = 0;


MultipleSteppersController msc(enabled, spi);

void callback(char *topic, byte *payload, unsigned int length) {
  String payloadStr;
  for (unsigned int i = 0; i < length; i++) {
    payloadStr += (char)payload[i];
  }

  Serial.print("Message arrived in topic: ");
  Serial.println(topic);
  Serial.print("Message: ");
  Serial.println(payloadStr);
  Serial.println("-----------------------");

  int separatorIndex = payloadStr.indexOf(',');
  if (separatorIndex != -1) {
    int mode = payloadStr.substring(0, separatorIndex).toInt();
    int value = payloadStr.substring(separatorIndex + 1).toInt();

    if (String(topic).equals("angle1")) {
      mode1 = mode;
      inputDegrees1 = value;
    } else if (String(topic).equals("angle2")) {
      mode2 = mode;
      inputDegrees2 = value;
    } else if (String(topic).equals("angle3")) {
      mode3 = mode;
      inputDegrees3 = value;
    } else if (String(topic).equals("angle4")) {
      mode4 = mode;
      inputDegrees4 = value;
    }
  }


  Serial.print("VALUE: ");
  Serial.println(inputDegrees1);

  Serial.print("MODE: ");
  Serial.println(mode1);

}

void setup() {
  Serial.begin(115200);
  Serial.println("Hello, ESP32-S3!");
  msc.enable();
  WiFi.begin(ssid, password);
  while (WiFi.status() != WL_CONNECTED) {
    delay(500);
    Serial.println("Connecting to WiFi..");
  }
  Serial.println("Connected to the Wi-Fi network");

  client.setServer(mqtt_broker, mqtt_port);
  client.setCallback(callback);
  while (!client.connected()) {
    String client_id = "esp32-client-";
    client_id += String(WiFi.macAddress());
    Serial.printf("The client %s connects to the public MQTT broker\n", client_id.c_str());
    if (client.connect(client_id.c_str(), mqtt_username, mqtt_password)) {
      Serial.println("Public EMQX MQTT broker connected");
    } else {
      Serial.print("failed with state ");
      Serial.print(client.state());
      delay(2000);
    }
  }

  client.subscribe("angle1");
  client.subscribe("angle2");
  client.subscribe("angle3");
  client.subscribe("angle4");
}

void reconnect() {
  String client_id = "esp32-client-";
  client_id += String(WiFi.macAddress());
  // Loop until we're reconnected
  while (!client.connected()) {
    Serial.print("Attempting MQTT connection...");
    // Attempt to connect
    if (client.connect(client_id.c_str())) {
      Serial.println("connected");
      // Subscribe
      client.subscribe("goTo1");
      client.subscribe("goTo2");
      client.subscribe("goTo3");
      client.subscribe("goTo4");
    } else {
      Serial.print("failed, rc=");
      Serial.print(client.state());
      Serial.println(" try again in 3 seconds");
      // Wait 5 seconds before retrying
      delay(3000);
    }
  }
}

void loop() {
  if (!client.connected()) {
    reconnect();
  }
  client.loop();

  Serial.print("Value: ");
  Serial.println(inputDegrees1);
  msc.goTo(mode1, 0, inputDegrees1);
  msc.goTo(mode2, 1, inputDegrees2);
  msc.goTo(mode3, 2, inputDegrees3);
  msc.goTo(mode4, 3, inputDegrees4);

/*
  msc.drop(0, dropInput0);
  msc.drop(1, dropInput1);
  msc.drop(2, dropInput2);
  msc.drop(3, dropInput3);
*/
}