#include <SoftwareSerial.h>
SoftwareSerial mySerial(16, 17);  // RX, TX

const int red = 14;
const int yellow = 33;
const int green = 12;

const unsigned long greenDuration = 5000;
const unsigned long yellowDuration = 2000;
const unsigned long redDuration = 5000;
const unsigned long transitionDuration = 2000;
const unsigned long heartbeatInterval = 1000;
const unsigned long heartbeatTimeout = 3000;
const unsigned long blinkInterval = 500;

String receivedData = "";
bool receiving = false;

unsigned long previousMillis = 0;
unsigned long lastHeartbeatMillis = 0;
unsigned long lastBlinkMillis = 0;
bool blinkState = false;
bool isError = false;  

enum State { GREEN, YELLOW, RED, TRANSITION, ERROR };
State currentState = GREEN;

#define REPEAT_SEND false

void sendCommand(String cmd) {
    mySerial.print("<");
    mySerial.print(cmd);
    mySerial.println(">");
    Serial.print("Sent: ");
    Serial.println(cmd);
}

void sendHeartbeat() {
    mySerial.print("<HB>");
    Serial.println("Sent: Heartbeat");
}

void ledRed() {
    digitalWrite(red, HIGH);
    digitalWrite(yellow, LOW);
    digitalWrite(green, LOW);
  }
  
  void ledYellow() {
    digitalWrite(red, LOW);
    digitalWrite(yellow, HIGH);
    digitalWrite(green, LOW);
  }
  
  void ledGreen() {
    digitalWrite(red, LOW);
    digitalWrite(yellow, LOW);
    digitalWrite(green, HIGH);
  }
  
void updateLights() {
    switch (currentState) {
        case GREEN:
            ledGreen();
            sendCommand("G");
            break;
        case YELLOW:
            ledYellow();
            sendCommand("Y");
            break;
        case RED:
            ledRed();
            sendCommand("R");
            break;
        case TRANSITION:
            sendCommand("T");
            break;
        case ERROR:
            if (millis() - lastBlinkMillis >= blinkInterval) {
                blinkState = !blinkState;
                digitalWrite(yellow, blinkState ? HIGH : LOW);
                lastBlinkMillis = millis();
            }
            break;
    }
}

void setup() {
    pinMode(green, OUTPUT);
    pinMode(yellow, OUTPUT);
    pinMode(red, OUTPUT);
    mySerial.begin(115200);
    Serial.begin(115200);
    previousMillis = millis();
    lastHeartbeatMillis = millis();
    updateLights();
}

void loop() {
  unsigned long currentMillis = millis();

  // Send heartbeat regularly
  if (currentMillis - lastHeartbeatMillis >= heartbeatInterval) {
      sendHeartbeat();
      lastHeartbeatMillis = currentMillis;
  }

  // Heartbeat timeout check
  if (currentMillis - lastHeartbeatMillis >= heartbeatTimeout) {
      if (!isError) {
          isError = true;  
      }
  }

  // State transitions
  switch (currentState) {
      case GREEN:
          if (currentMillis - previousMillis >= greenDuration) {
              currentState = YELLOW;
              previousMillis = currentMillis;
              updateLights();
          }
          break;
      case YELLOW:
          if (currentMillis - previousMillis >= yellowDuration) {
              currentState = RED;
              previousMillis = currentMillis;
              updateLights();
          }
          break;
      case RED:
          if (currentMillis - previousMillis >= redDuration) {
              currentState = TRANSITION;
              previousMillis = currentMillis;
              updateLights();
          }
          break;
      case TRANSITION:
          if (currentMillis - previousMillis >= transitionDuration) {
              currentState = GREEN;
              previousMillis = millis();
              updateLights();
          }
          break;
      case ERROR:
          updateLights();
          break;
  }

  // Handle received data
  while (mySerial.available()) {
      char incomingByte = mySerial.read();

      if (incomingByte == '<') {
          receivedData = "";
          receiving = true;
      } else if (incomingByte == '>') {
          receiving = false;
          Serial.print("Received Data: ");
          Serial.println(receivedData);

          if (receivedData == "HB") {
              lastHeartbeatMillis = millis();
              if (isError) {
                  digitalWrite(yellow, LOW);  
              }
              isError = false;  
          }

          receivedData = "";
      } else if (receiving) {
          if (incomingByte >= 32 && incomingByte <= 126) {
              receivedData += incomingByte;
          }
      }
  }

  // Blinking yellow light in error state
  if (isError) {
      if (currentMillis - lastBlinkMillis >= blinkInterval) {
          blinkState = !blinkState;
          digitalWrite(yellow, blinkState ? HIGH : LOW);  
          lastBlinkMillis = currentMillis;
      }
  }
}