// DCCControllerSimple V1 Software Ben Forster 2023

// Import Libraries
#include <WiFi.h>
#include <WiFiMulti.h>
#include <WiFiClientSecure.h>
#include <WebSocketsClient.h>
#include <SocketIOclient.h>

WiFiMulti WiFiMulti;
SocketIOclient socketIO;

// Define The Input Pins
#define ADDRESS_BIT_1 1
#define ADDRESS_BIT_2 2
#define ADDRESS_BIT_3 3
#define ADDRESS_BIT_4 4
#define SPEED_INCREASE_BTN 8
#define CUSTOM_BTN_3 9
#define FORWARD_BTN 10
#define LED_REGISTER_CLOCK 11
#define LED_SHIFT_CLOCK 12
#define LED_SERIAL_IN 13
#define LED_OUTPUT_ENABLE_ACTIVE_LOW 14
#define DIGIT_2_LATCH_ENABLE 15
#define DIGIT_2_BLANKING_ACTIVE_LOW 16
#define BCD_BIT_1 17
#define BCD_BIT_4 21
#define BCD_BIT_3 33
#define BCD_BIT_2 34
#define DIGIT_1_LATCH_ENABLE 35
#define DIGIT_1_BLANKING_ACTIVE_LOW 36
#define REVERSE_BTN 37
#define CUSTOM_BTN_1 38
#define CUSTOM_BTN_2 39
#define SPEED_DECREASE_BTN 40
#define EMERGENCY_STOP_BTN 41
#define ROTARY_A 42
#define ROTARY_B 45
#define ROTARY_SWITCH 46

// then define which led is at which index in the shift register
#define EMERGENCY_STOP_BTN_LED 0
#define SPEED_DECREASE_BTN_LED 1
#define CUSTOM_BTN_1_LED 2
#define REVERSE_BTN_LED 3
#define CUSTOM_BTN_2_LED 4
#define FORWARD_BTN_LED 5
#define CUSTOM_BTN_3_LED 6
#define SPEED_INCREASE_BTN_LED 7

// then define the wifi credentials
#define WIFI_SSID "DCC_CONTROLLERS"
#define WIFI_PASSWORD "dcc_controllers"
#define SERVER_IP_ADDRESS "ws://192.168.0.2"
#define SERVER_PORT 81
#define SERVER_URL "/socket.io"
#define RECONNECTION_TIME 5000

// then define the variables for locomotives
volatile int speed = 0; // locomotive speed from 0 to +28
volatile bool customBtn1On = false;
volatile bool customBtn2On = false;
volatile bool customBtn3On = false;
volatile bool isForward = true;
volatile bool localEmergencyStop = false;

// then define the variables for the actual controller
bool emergencyStopFromServer = false;
volatile bool activateEmergencyStop = false;

// then define the loco address
int locoAddressBit1;
int locoAddressBit2;
int locoAddressBit3;
int locoAddressBit4;

volatile int lastStateRotaryA;
volatile int currentStateRotaryA;

volatile bool resendPacket = false;

// then define the array for use with the shift register
int shiftRegisterLedArray[8] = {1,0,0,0,0,1,0,0};

// error number 99 88 77 66 55
int errors[5] = {0, 0, 0, 0, 0};

// then define the errors function
void ShowError(int errorCode) {
    // then check what error code it is
    if(errorCode == 99) {
        // then we have the disconnection error to display
        errors[0] = 1;
    }
    else if(errorCode == 88) {
        // then we have the general error to display
        errors[1] = 1;
    }
    else if(errorCode == 77) {
        // then we have the not connecting error to display
        errors[2] = 1;
    }
    else if(errorCode == 66) {
        // then we have the emergency stop error to display
        errors[3] = 1;
    }
    else if(errorCode == 55) {
        // then we have the track short error to display 
        errors[4] = 1;
    }
}

void CancelError(int errorCode) {
    // then check what error code it is
    if(errorCode == 99) {
        // then we have the disconnection error to remove
        errors[0] = 1;
    }
    else if(errorCode == 88) {
        // then we have the general error to remove
        errors[1] = 1;
    }
    else if(errorCode == 77) {
        // then we have the not connecting error to remove
        errors[2] = 1;
    }
    else if(errorCode == 66) {
        // then we have the emergency stop error to remove
        errors[3] = 1;
    }
    else if(errorCode == 55) {
        // then we have the track short error to remove
        errors[4] = 1;
    }
}

// then define the shift register functions
void updateShiftRegisterArray(int index, int value) {
    shiftRegisterLedArray[index] = value;
    
}
void sendArrayToShiftRegister() {
    // then make a for loop that will run 8 times
    for(int i = 0; i < sizeof(shiftRegisterLedArray); i++) {
        // then pull the clock low
        digitalWrite(LED_SHIFT_CLOCK, LOW);
        // then write the data to the SR DATA IN line
        digitalWrite(LED_SERIAL_IN, shiftRegisterLedArray[i]);
        delay(0.5);
        // then pull the clock high to clock it in
        digitalWrite(LED_SHIFT_CLOCK, HIGH);
    }
    // then once it has finished, pull the clock low and then latch the data by pulling the register clock high and low
    digitalWrite(LED_SHIFT_CLOCK, LOW);
    digitalWrite(LED_REGISTER_CLOCK, HIGH);
    delay(0.5);
    digitalWrite(LED_REGISTER_CLOCK, LOW);
    // then make sure the output enable is active so pull it low
    digitalWrite(LED_OUTPUT_ENABLE_ACTIVE_LOW, LOW);
}

void autoUpdateShiftRegisterArray() {
    // then find the data and the index and pass them to the updateShiftRegisterArray function
    updateShiftRegisterArray(CUSTOM_BTN_1_LED, customBtn1On);
    updateShiftRegisterArray(CUSTOM_BTN_2_LED, customBtn2On);
    updateShiftRegisterArray(CUSTOM_BTN_3_LED, customBtn3On);
    // then do the direction
    if(isForward == true) {
        updateShiftRegisterArray(FORWARD_BTN_LED, 1);
        updateShiftRegisterArray(REVERSE_BTN_LED, 0);
    }
    else {
        updateShiftRegisterArray(FORWARD_BTN_LED, 0);
        updateShiftRegisterArray(REVERSE_BTN_LED, 1);
    }

    // then send the array to the shift register
    sendArrayToShiftRegister();
    
    
}

// then define the 7 segment display function
void displayNumberOn7Segment(int number) {
    // then do a validity check on the number to make sure its not less than zero and not more than 99
    if(number >= 0 && number <= 99) {
        // then the number is in range so split it into digits
        int digit1;
        int digit2 = number % 10;
        if(number > 9) {
            digit1 = number / 10 % 10;
        }
        else {
            digit1 = 0;
        }
        // then do digit 1 first
        int digit1binary[4] = {0,0,0,0};
        // then divide the number into the array
        int newNumber1 = digit1;
        digit1binary[3] = newNumber1 % 2;
        newNumber1 = digit1 / 2;
        digit1binary[2] = newNumber1 % 2;
        newNumber1 = newNumber1 / 2;
        digit1binary[1] = newNumber1 % 2;
        newNumber1 = newNumber1 / 2;
        digit1binary[0] = newNumber1 % 2;
        // then do digit 2
        int digit2binary[4] = {0,0,0,0};
        // then divide the number into the array
        int newNumber2 = digit2;
        digit2binary[3] = newNumber2 % 2;
        newNumber2 = newNumber2 / 2;
        digit2binary[2] = newNumber2 % 2;
        newNumber2 = newNumber2 / 2;
        digit2binary[1] = newNumber2 % 2;
        newNumber2 = newNumber2 / 2;
        digit2binary[0] = newNumber2 % 2;

        // then update the 7 segment display
        // first pull both latch enable pins low
        digitalWrite(DIGIT_1_LATCH_ENABLE, LOW);
        digitalWrite(DIGIT_2_LATCH_ENABLE, LOW);

        // then put digit 1's data on the address lines
        digitalWrite(BCD_BIT_1, digit1binary[0]);
        digitalWrite(BCD_BIT_2, digit1binary[1]);
        digitalWrite(BCD_BIT_3, digit1binary[2]);
        digitalWrite(BCD_BIT_4, digit1binary[3]);

        // then latch the data into the 7 segment
        digitalWrite(DIGIT_1_LATCH_ENABLE, HIGH);
        delay(0.5);
        digitalWrite(DIGIT_1_LATCH_ENABLE, LOW);

        // then put digit 2's data on the address lines
        digitalWrite(BCD_BIT_1, digit2binary[0]);
        digitalWrite(BCD_BIT_2, digit2binary[1]);
        digitalWrite(BCD_BIT_3, digit2binary[2]);
        digitalWrite(BCD_BIT_4, digit2binary[3]);

        // then latch the data into the 7 segment
        digitalWrite(DIGIT_2_LATCH_ENABLE, HIGH);
        delay(0.5);
        digitalWrite(DIGIT_2_LATCH_ENABLE, LOW);

        // then pull the address lines low
        digitalWrite(BCD_BIT_1, LOW);
        digitalWrite(BCD_BIT_2, LOW);
        digitalWrite(BCD_BIT_3, LOW);
        digitalWrite(BCD_BIT_4, LOW);

    }
}

// readMessage function - takes the payload from the webSocketEvent
void readMessage(uint8_t payload[28], int length) {
    // then read off each bit from the char array and check if anything needs to be updated
    // global emergency stop
    if((int)payload[0] == 1) {
        // then flash the emergency stop error
        ShowError(66);
    }
    else {
        // cancel the error
        CancelError(66);
    }

    // global track short
    if((int)payload[1] == 1) {
        // then flash the track short error
        ShowError(55);
    }
    else {
        // cancel the error
        CancelError(55);
    }

    // then check the loco address bits to see if we need to read on
    // check the controller bit 1 matches the packet bit 1
    if((int)payload[5] == locoAddressBit1 && (int)payload[4] == locoAddressBit2 && (int)payload[3] == locoAddressBit3 && (int)payload[2] == locoAddressBit4) {
        // then the address matches so read the packet

        // then read the speed bcd bits and update the speed on the controller without sending another packet
        int newSpeed = 0;
        // 5th bit
        if((int)payload[7] == 1) {
            // then increment the speed by 16
            newSpeed = newSpeed + 16;
        }
        // 4th bit
        if((int)payload[8] == 1) {
            // then increment the speed by 8
            newSpeed = newSpeed + 8;
        }
        // 3rd bit
        if((int)payload[9] == 1) {
            // then increment the speed by 4
            newSpeed = newSpeed + 4;
        }
        // 2nd bit
        if((int)payload[10] == 1) {
            // then increment the speed by 2
            newSpeed = newSpeed + 2;
        }
        // 1st bit
        if((int)payload[11] == 1) {
            // then increment the speed by 1
            newSpeed = newSpeed + 1;
        }

        // then set the speed to the newSpeed
        speed = newSpeed;

        // then update the direction
        isForward = (int)payload[12];

        // then check the 3 functions that this device can handle
        customBtn1On = (int)payload[13];
        customBtn2On = (int)payload[14];
        customBtn3On = (int)payload[15];

        // then update the shift register
        autoUpdateShiftRegisterArray();

        // then update the 7 segment display
        displayNumberOn7Segment(speed);

    }
}


// then define the function that handles any events
void socketIOEvent(socketIOmessageType_t type, uint8_t * payload, size_t length) {
    // then check what even happened
    switch(type) {
        case sIOtype_DISCONNECT:
            // disconnected so give the error 99
            ShowError(99);
            break;
        case sIOtype_CONNECT:
            // connected
            socketIO.send(sIOtype_CONNECT, "/");
            break;
        case sIOtype_ERROR:
            // error so give the error code 88
            ShowError(88);
            break;
        case sIOtype_EVENT:
            // then send the message to the readMessage function
            readMessage(payload, length);
            break;
    }
}


// ################### INTERRUPT FUNCTIONS ########################

void interrupt_speedIncreaseBtn() {
    // then increment the speed if speed is not 28 or more
    if(speed < 28) {
        // then increment the speed
        speed = speed + 1;
        // then set the resendPacket bool to true
        resendPacket = true;
    }
}

void interrupt_customBtn3() {
    // then flip the state
    customBtn3On = !customBtn3On;
    // then set the resendPacket bool to true
    resendPacket = true;
}

void interrupt_forwardBtn() {
    // then set the isForward boolean to true
    isForward = true;
    // then set the resendPacket bool to true
    resendPacket = true;
}

void interrupt_reverseBtn() {
    // then set the isForward boolean to false
    isForward = false;
    // then set the resendPacket bool to true
    resendPacket = true;
}

void interrupt_customBtn1() {
    // then flip the state
    customBtn1On = !customBtn1On;
    // then set the resendPacket bool to true
    resendPacket = true;
}

void interrupt_customBtn2() {
    // then flip the state
    customBtn2On = !customBtn2On;
    // then set the resendPacket bool to true
    resendPacket = true;
}

void interrupt_speedDecreaseBtn() {
    // then decrement the speed if speed is not 0 or below
    if(speed > 0) {
        // then decrement the speed
        speed = speed - 1;
        // then set the resendPacket bool to true
        resendPacket = true;
    }
}

void interrupt_emergencyStopBtn() {
    // then set the speed to 0
    speed = 0;
    // then set the localEmergencyStop to true
    localEmergencyStop = true;
    // then set the resendPacket bool to true
    resendPacket = true;
}

void interrupt_rotarySwitch() {
    // this function checks if the rotary switch has been pressed
    if(activateEmergencyStop == false) {
        // then set it to true
        activateEmergencyStop = true;
        // then set the resendPacket bool to true
        resendPacket = true;
    }
    else {
        // then cancel the emergency stop
        activateEmergencyStop = false;
        // then set the resendPacket bool to true
        resendPacket = true;
    }
}

// then the rotary encoder interrupt function
void interrupt_checkRotaryEncoder() {
    // check the state of ROTARY_A
    currentStateRotaryA = digitalRead(ROTARY_A);
    // then check if it is different to the previous state
    if(currentStateRotaryA != lastStateRotaryA && currentStateRotaryA == 0) {
        // then check if the Rotary B state is different than the Rotary A state
        if(digitalRead(ROTARY_B) != currentStateRotaryA) {
            // then it is rotating clockwise
            // so increment the speed if its not more than 27
            if(speed < 28) {
                speed = speed + 1;
                // then set the resendPacket bool to true
                resendPacket = true;
            }      
        }
        else {
            // then it is rotating counterclockwise
            // so decrement the speed if its not less than 1
            if(speed > 1) {
                speed = speed - 1;
                // then set the resendPacket bool to true
                resendPacket = true;
            }
        }
    }
    // then update the last Rotary A state
    lastStateRotaryA = currentStateRotaryA;
}

void setup() {

    // then initialise the wifi connection
    WiFiMulti.addAP(WIFI_SSID, WIFI_PASSWORD);

    // then make sure its connected
    while(WiFiMulti.run() != WL_CONNECTED) {
        delay(100);
        ShowError(77);
    }

    CancelError(77);

    // then start the websocket connection
    socketIO.begin(SERVER_IP_ADDRESS, SERVER_PORT, SERVER_URL);

    // then attach the event handler
    socketIO.onEvent(socketIOEvent);

    // then set the reconnection time
    //webSocket.setReconnectInterval(RECONNECTION_TIME);

    // PIN MODES
    // set all the pin modes
    pinMode(ADDRESS_BIT_1, INPUT_PULLUP);
    pinMode(ADDRESS_BIT_2, INPUT_PULLUP);
    pinMode(ADDRESS_BIT_3, INPUT_PULLUP);
    pinMode(ADDRESS_BIT_4, INPUT_PULLUP);
    pinMode(SPEED_INCREASE_BTN, INPUT_PULLUP);
    pinMode(CUSTOM_BTN_3, INPUT_PULLUP);
    pinMode(FORWARD_BTN, INPUT_PULLUP);
    pinMode(LED_REGISTER_CLOCK, OUTPUT);
    pinMode(LED_SHIFT_CLOCK, OUTPUT);
    pinMode(LED_SERIAL_IN, OUTPUT);
    pinMode(LED_OUTPUT_ENABLE_ACTIVE_LOW, OUTPUT);
    pinMode(DIGIT_2_LATCH_ENABLE, OUTPUT);
    pinMode(DIGIT_2_BLANKING_ACTIVE_LOW, OUTPUT);
    pinMode(BCD_BIT_1, OUTPUT);
    pinMode(BCD_BIT_4, OUTPUT);
    pinMode(BCD_BIT_3, OUTPUT);
    pinMode(BCD_BIT_2, OUTPUT);
    pinMode(DIGIT_1_LATCH_ENABLE, OUTPUT);
    pinMode(DIGIT_1_BLANKING_ACTIVE_LOW, OUTPUT);
    pinMode(REVERSE_BTN, INPUT_PULLUP);
    pinMode(CUSTOM_BTN_1, INPUT_PULLUP);
    pinMode(CUSTOM_BTN_2, INPUT_PULLUP);
    pinMode(SPEED_DECREASE_BTN, INPUT_PULLUP);
    pinMode(EMERGENCY_STOP_BTN, INPUT_PULLUP);
    pinMode(ROTARY_A, INPUT_PULLUP);
    pinMode(ROTARY_B, INPUT_PULLUP);
    pinMode(ROTARY_SWITCH, INPUT_PULLUP);

    // attach interrupts to all the input pins exept the address bits as these are not important
    attachInterrupt(SPEED_INCREASE_BTN, interrupt_speedIncreaseBtn, FALLING);
    attachInterrupt(CUSTOM_BTN_3, interrupt_customBtn3 , FALLING);
    attachInterrupt(FORWARD_BTN, interrupt_forwardBtn, FALLING);
    attachInterrupt(REVERSE_BTN, interrupt_reverseBtn, FALLING);
    attachInterrupt(CUSTOM_BTN_1, interrupt_customBtn1 , FALLING);
    attachInterrupt(CUSTOM_BTN_2, interrupt_customBtn2, FALLING);
    attachInterrupt(SPEED_DECREASE_BTN, interrupt_speedDecreaseBtn, FALLING);
    attachInterrupt(EMERGENCY_STOP_BTN, interrupt_emergencyStopBtn , FALLING);
    attachInterrupt(ROTARY_A, interrupt_checkRotaryEncoder, CHANGE);
    attachInterrupt(ROTARY_B, interrupt_checkRotaryEncoder, CHANGE);
    attachInterrupt(ROTARY_SWITCH, interrupt_rotarySwitch, FALLING);

    // then get the current state of rotary a and store it
    lastStateRotaryA = digitalRead(ROTARY_A);

    // then read the loco address from the address bits and flip them because pullup resistors have been used
    locoAddressBit1 = !digitalRead(ADDRESS_BIT_1);
    locoAddressBit2 = !digitalRead(ADDRESS_BIT_2);
    locoAddressBit3 = !digitalRead(ADDRESS_BIT_3);
    locoAddressBit4 = !digitalRead(ADDRESS_BIT_4);

    // then pull the 7 segment display blanking lines high
    digitalWrite(DIGIT_1_BLANKING_ACTIVE_LOW, HIGH);
    digitalWrite(DIGIT_2_BLANKING_ACTIVE_LOW, HIGH);


    // then send 7 more clock pulses to the led shift register
    digitalWrite(LED_SHIFT_CLOCK, HIGH);
    delay(0.5);
    digitalWrite(LED_SHIFT_CLOCK, LOW);
    digitalWrite(LED_SHIFT_CLOCK, HIGH);
    delay(0.5);
    digitalWrite(LED_SHIFT_CLOCK, LOW);
    digitalWrite(LED_SHIFT_CLOCK, HIGH);
    delay(0.5);
    digitalWrite(LED_SHIFT_CLOCK, LOW);
    digitalWrite(LED_SHIFT_CLOCK, HIGH);
    delay(0.5);
    digitalWrite(LED_SHIFT_CLOCK, LOW);
    digitalWrite(LED_SHIFT_CLOCK, HIGH);
    delay(0.5);
    digitalWrite(LED_SHIFT_CLOCK, LOW);
    digitalWrite(LED_SHIFT_CLOCK, HIGH);
    delay(0.5);
    digitalWrite(LED_SHIFT_CLOCK, LOW);
    digitalWrite(LED_SHIFT_CLOCK, HIGH);
    delay(0.5);
    digitalWrite(LED_SHIFT_CLOCK, LOW);

}

void sendMessage(char payload[16]) {
    socketIO.sendEVENT(payload);
}

void prepareMessage() {
    // then create a char array with all the bits in to send to the host
    char message[16];
    
    // global emergency stop
    message[0] = (int)activateEmergencyStop;
    // global track short - host to controller only
    message[1] = 0;

    // loco address 4th bit
    message[2] = locoAddressBit4;
    // loco address 3rd bit
    message[3] = locoAddressBit3;
    // loco address 2nd bit
    message[4] = locoAddressBit2;
    // loco address 1st bit
    message[5] = locoAddressBit1;

    // then do the local emergency stop (just the current loco)
    message[6] = (int)localEmergencyStop;

    // then do the current speed

    int currentSpeedArray[5] = {0,0,0,0,0};
    // then divide the number into the array
    int newNumber1 = speed;
    currentSpeedArray[4] = newNumber1 % 2;
    newNumber1 = newNumber1 / 2;
    currentSpeedArray[3] = newNumber1 % 2;
    newNumber1 = newNumber1 / 2;
    currentSpeedArray[2] = newNumber1 % 2;
    newNumber1 = newNumber1 / 2;
    currentSpeedArray[1] = newNumber1 % 2;
    newNumber1 = newNumber1 / 2;
    currentSpeedArray[0] = newNumber1 % 2;

    message[7] = currentSpeedArray[4];
    message[8] = currentSpeedArray[3];
    message[9] = currentSpeedArray[2];
    message[10] = currentSpeedArray[1];
    message[11] = currentSpeedArray[0];

    // then send the direction
    message[12] = (int)isForward;

    // then send the 3 custom buttons that this hardware is capable of

    message[13] = (int)customBtn1On;
    message[14] = (int)customBtn2On;
    message[15] = (int)customBtn3On;

    // then call the sendMessage function
    sendMessage(message);
}

void loop() {
    // then go in the webSocket loop
    socketIO.loop();
    // then check if there are any errors to display
    if(errors[4] == 1) {
        displayNumberOn7Segment(55);
    }
    if(errors[3] == 1) {
        displayNumberOn7Segment(66);
    }
    if(errors[2] == 1) {
        displayNumberOn7Segment(77);
    }
    if(errors[1] == 1) {
        displayNumberOn7Segment(88);
    }
    if(errors[0] == 1) {
        displayNumberOn7Segment(99);
    }

    // then check if there are no errors and the speed can be put on the display
    if(errors[0] == 0 && errors[1] == 0 && errors[2] == 0 && errors[3] == 0 && errors[4] == 0) {
        displayNumberOn7Segment(speed);
    }
    // then update the lights
    autoUpdateShiftRegisterArray();
    // then check if the packet needs resending
    if(resendPacket == true) {
        prepareMessage();
        resendPacket = false;
    }
}