// Configure I/O ports
const boolean debug = true; // true allows serial port to function
const int MOTOR_pwm = 4; // GPIO to control motor speed
const int MOTOR_dir = 5; // GPIO to control motor direction
const int LED_headlight = 3; // GPIO for head light
const int LED_cablight = 2; // GPIO for cab light
const int LED_taillight = 1; // GPIO for tail light

// Configure timing and variables
String CTRL = ""; String lastCTRL = "";
int current_speed = 1;
int target_speed = 0;
int current_direction = 1;
int target_direction = 1;
unsigned long SPEED_timer = 0;
const int SPEED_interval = 133; // how frequently to update speed
const int SPEED_step = 3; // increment speed by this much
const int SPEED_adj = (PWMRANGE + 1) / 256;
const int fwd_min = 22; // minimum forward speed before motor turns
const int fwd_max = 255; // maximum formward speed
const int rev_min = 22; // minimum reverse speed before motor turns
const int rev_max = 128; // maximum reverse speed
const float SPEED_mph = 0.54; // conversion factor
unsigned long CTRL_timer = 0;
const int CTRL_timeout = 5100; // client not responding timeout (ms)

// Configure WIFI
#include 
const char* SSID = "sourpuss"; // widfi SSID
const char* PASS = "3037764157"; // wifi password
const char* LOCO = "DRGW101"; // Call sign for this loco

// Configure MQTT
#include 
const String lastwill = String(LOCO) + "/IO/online";
boolean publishInfo = true;
AsyncMqttClient mqttClient;

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

pinMode(LED_headlight, OUTPUT);
digitalWrite(LED_headlight, LOW);
pinMode(LED_cablight, OUTPUT);
digitalWrite(LED_cablight, LOW);
if (!debug) pinMode(LED_taillight, OUTPUT);
if (!debug) digitalWrite(LED_taillight, HIGH);

analogWriteFreq(200);
pinMode(MOTOR_pwm, OUTPUT);
pinMode(MOTOR_dir, OUTPUT);
digitalWrite(MOTOR_pwm, LOW);
digitalWrite(MOTOR_dir, LOW);

boolean cab = false;

WiFi.hostname(LOCO);
WiFi.mode(WIFI_STA);
WiFi.begin (SSID, PASS);
Serial.print("\r\nConnecting to ");
Serial.print(SSID);
while (WiFi.status() != WL_CONNECTED) {
delay(250);
cab = !cab;
Serial.print(".");
digitalWrite(LED_cablight, cab);
}
Serial.println(WiFi.localIP());
digitalWrite(LED_cablight, LOW);

mqttClient.onConnect(onMqttConnect);
mqttClient.onDisconnect(onMqttDisconnect);
mqttClient.onSubscribe(onMqttSubscribe);
mqttClient.onUnsubscribe(onMqttUnsubscribe);
mqttClient.onMessage(onMqttMessage);
mqttClient.onPublish(onMqttPublish);
mqttClient.setServer(IPAddress(10, 10, 1, 1), 1883);

//mqttClient.setKeepAlive(5).setCleanSession(false).setWill("/IO/online", 2, true, "0").setCredentials("username", "password").setClientId(LOCO);
mqttClient.setKeepAlive(5).setCleanSession(false).setWill(lastwill.c_str(), 2, true, "false").setClientId(LOCO);
Serial.println("Connecting to MQTT...");
mqttClient.connect();
} // setup

void loop() {
unsigned long NOW = millis();
unsigned int flag = 0;
String mqtt = "";
char buf[10];

if (NOW >= SPEED_timer) {
flag = 0;
if (current_direction != target_direction) {
current_speed -= SPEED_step;
if (current_speed <= fwd_min) current_speed = 0; flag=1; } else { if (target_speed < current_speed) { current_speed -= SPEED_step; if (current_speed <= fwd_min) current_speed = 0; if (current_speed < target_speed) target_speed = current_speed; flag=1; } if (target_speed > current_speed) {
if (current_speed <= fwd_min) current_speed = fwd_min; current_speed += SPEED_step; if (current_speed > target_speed) target_speed = current_speed;
flag=1;
}
}

//if (current_speed < 0) current_speed = 0; if (current_speed < fwd_min) { if (current_direction != target_direction) { mqtt = String(LOCO) + "/IO/dir"; sprintf (buf, "%d", target_direction); mqttClient.publish(mqtt.c_str(), 2, false, buf); } current_speed = 0; current_direction = target_direction; digitalWrite(LED_cablight, HIGH); } else digitalWrite(LED_cablight, LOW); if (current_direction < 0) { digitalWrite(LED_headlight, LOW); digitalWrite(LED_taillight, LOW); if (target_speed > rev_max) target_speed = rev_max;
if (current_speed > rev_max) current_speed = rev_max;
}
if (current_direction > 0) {
digitalWrite(LED_headlight, HIGH);
digitalWrite(LED_taillight, HIGH);
if (target_speed > fwd_max) target_speed = fwd_max;
if (current_speed > fwd_max) current_speed = fwd_max;
}

if (flag) { // Only update if something changed
if (current_direction >= 0) digitalWrite(MOTOR_dir, LOW);
else digitalWrite(MOTOR_dir, HIGH);
analogWrite(MOTOR_pwm, current_speed * SPEED_adj);

mqtt = String(LOCO) + "/IO/speed";
sprintf (buf, "%d", current_speed);
mqttClient.publish(mqtt.c_str(), 2, false, buf);
Serial.print("Speed: ");
Serial.println(current_speed);
}

SPEED_timer = NOW + SPEED_interval;
} // NOW >= SPEED_timer

if ((CTRL != "") && (NOW >= CTRL_timer)) {
// Release control from non-responding client
Serial.println("Client stopped responding");
CTRL = "";
target_speed = 0;
target_direction = 1;

mqtt = String(LOCO) + "/CLIENT/" + CTRL + "/#";
mqttClient.unsubscribe(mqtt.c_str());

mqtt = String(LOCO) + "/CLIENT/" + CTRL + "/has_ctrl";
mqttClient.publish(mqtt.c_str(), 2, false, "");

mqtt = String(LOCO) + "/IO/ctrl";
mqttClient.publish(mqtt.c_str(), 2, true, "true");
}
} // loop

void onMqttConnect(bool sessionPresent) {
Serial.println("** Connected to the broker **");
Serial.print("Session present: ");
Serial.println(sessionPresent);

String mqtt = "";
String mqttPub = "";
char buf[10];

mqtt = String(LOCO) + "/IO/#";
uint16_t packetIdSub = mqttClient.subscribe(mqtt.c_str(), 2);

mqtt = String(LOCO) + "/IO/online";
mqttClient.publish(mqtt.c_str(), 2, true, "true");

mqtt = String(LOCO) + "/IO/fwd_min";
sprintf (buf, "%d", fwd_min);
mqttClient.publish(mqtt.c_str(), 2, true, buf);

mqtt = String(LOCO) + "/IO/fwd_max";
sprintf (buf, "%d", fwd_max);
mqttClient.publish(mqtt.c_str(), 2, true, buf);

mqtt = String(LOCO) + "/IO/rev_min";
sprintf (buf, "%d", rev_min);
mqttClient.publish(mqtt.c_str(), 2, true, buf);

mqtt = String(LOCO) + "/IO/rev_max";
sprintf (buf, "%d", rev_max);
mqttClient.publish(mqtt.c_str(), 2, true, buf);

mqtt = String(LOCO) + "/IO/mph";
dtostrf(SPEED_mph, 9, 6, buf);
mqttClient.publish(mqtt.c_str(), 2, true, buf);

mqtt = String(LOCO) + "/IO/speed";
sprintf (buf, "%d", current_speed);
mqttClient.publish(mqtt.c_str(), 2, true, buf);

mqtt = String(LOCO) + "/IO/dir";
sprintf (buf, "%d", current_direction);
mqttClient.publish(mqtt.c_str(), 2, true, buf);

mqtt = String(LOCO) + "/IO/ctrl";
mqttClient.publish(mqtt.c_str(), 2, true, "true");
}

void onMqttDisconnect(AsyncMqttClientDisconnectReason reason) {
Serial.println("** Disconnected from the broker **");
Serial.println("Reconnecting to MQTT...");
mqttClient.connect();
}

void onMqttSubscribe(uint16_t packetId, uint8_t qos) {
Serial.println("** Subscribe acknowledged **");
Serial.print(" packetId: ");
Serial.println(packetId);
Serial.print(" qos: ");
Serial.println(qos);
}

void onMqttUnsubscribe(uint16_t packetId) {
Serial.println("** Unsubscribe acknowledged **");
Serial.print(" packetId: ");
Serial.println(packetId);
}

void onMqttMessage(char* topic, char* payload, AsyncMqttClientMessageProperties properties, size_t len, size_t index, size_t total) {
Serial.println("** Publish received **");
Serial.print(" topic: ");
Serial.println(topic);
/*
Serial.print(" qos: ");
Serial.println(properties.qos);
Serial.print(" dup: ");
Serial.println(properties.dup);
Serial.print(" retain: ");
Serial.println(properties.retain);
Serial.print(" len: ");
Serial.println(len);
Serial.print(" index: ");
Serial.println(index);
Serial.print(" total: ");
Serial.println(total);
Serial.print(" payload: '");
Serial.print(payload);
Serial.println("'");
*/
//String data = String(payload).substring(0,len);
char data[len+1];
strncpy(data, payload, len);
data[len] = '\0';

Serial.print(" data: '");
Serial.print(data);
Serial.println("'");

unsigned long NOW = millis();
String cmd = String(topic).substring(String(LOCO).length() + 1);
String mqtt = "";
char buf[4];

Serial.println("CMD: '" + cmd + "'");

if ((CTRL == "") && (cmd = "IO/ctrl") && (len == 23)) {
CTRL = String(data);
lastCTRL = CTRL;

mqtt = String(LOCO) + "/IO/ctrl";
mqttClient.publish(mqtt.c_str(), 2, true, "false");

mqtt = String(LOCO) + "/CLIENT/" + CTRL + "/#";
uint16_t packetIdSub = mqttClient.subscribe(mqtt.c_str(), 2);

mqtt = String(LOCO) + "/CLIENT/" + CTRL + "/has_ctrl";
mqttClient.publish(mqtt.c_str(), 2, false, "true");

CTRL_timer = NOW + CTRL_timeout;
}

if (cmd == "CLIENT/" + CTRL + "/ping") {
Serial.print("PING: ");
Serial.println(data);
CTRL_timer = NOW + CTRL_timeout;
}

if (cmd == "CLIENT/" + CTRL + "/has_ctrl") {
if (String(data) != "true") {
Serial.println("Client requested disconnect");
mqtt = String(LOCO) + "/CLIENT/" + CTRL + "/#";
mqttClient.unsubscribe(mqtt.c_str());
mqtt = String(LOCO) + "/IO/ctrl";
mqttClient.publish(mqtt.c_str(), 2, true, "true");
CTRL = "";
target_speed = 0;
target_direction = 1;
}
}

if (cmd == "CLIENT/" + CTRL + "/set_dir") {
target_direction = atoi(data);
//target_direction = data.toInt();
Serial.print("DIR: ");
Serial.println(target_direction);
CTRL_timer = NOW + CTRL_timeout;
}

if (cmd == "CLIENT/" + CTRL + "/set_spd") {
target_speed = atoi(data);
//target_speed = data.toInt();
Serial.print("SPEED: ");
Serial.println(target_speed);
CTRL_timer = NOW + CTRL_timeout;
}
}

void onMqttPublish(uint16_t packetId) {
Serial.println("** Publish acknowledged **");
Serial.print(" packetId: ");
Serial.println(packetId);
}