#include <Arduino.h>
#include "MAGX_FSM.h"

// Define pin numbers for LEDs
const int RED_PIN = 13;
const int GREEN_PIN = 12;
const int YELLOW_PIN = 11;

// Define durations for each state in milliseconds
const int RED_DURATION = 5000;   // 5 seconds
const int GREEN_DURATION = 3000; // 3 seconds
const int YELLOW_DURATION = 1000; // 2 seconds

// Flags to control state suspension
bool redStateSuspended = false;
bool greenStateSuspended = false;
bool yellowStateSuspended = false;

// Function pointers for state logic
void redStateLogic(StateMachine &fsm);
void greenStateLogic(StateMachine &fsm);
void yellowStateLogic(StateMachine &fsm);

// State transition functions
void greenStateCallback(StateMachine &fsm) {
  fsm.setStateWithLogic("GREEN_STATE", greenStateLogic);
}
void yellowStateCallback(StateMachine &fsm) {
  fsm.setStateWithLogic("YELLOW_STATE", yellowStateLogic);
}
void redStateCallback(StateMachine &fsm) {
  fsm.setStateWithLogic("RED_STATE", redStateLogic);
}

// Common function to handle state suspension
void checkAndTransition(bool &isSuspended, void (*callback)(StateMachine &), StateMachine &fsm) {
  if (!isSuspended) {
    callback(fsm);
  }
}

void waitUntilElapsed(unsigned long duration, void (*callback)(StateMachine &), StateMachine &fsm) {
  static unsigned long previousMillis = 0; // Static variable to maintain state across function calls
  unsigned long currentMillis = millis();

  if (currentMillis - previousMillis >= duration) {
    // If the duration has elapsed, call the callback function
    callback(fsm);
    // Reset the timer for the next duration
    previousMillis = currentMillis;
  }
}

void redStateLogic(StateMachine &fsm) {
  // Set red light ON, green and yellow OFF
  digitalWrite(RED_PIN, HIGH);
  digitalWrite(GREEN_PIN, LOW);
  digitalWrite(YELLOW_PIN, LOW);
  waitUntilElapsed(RED_DURATION, greenStateCallback, fsm);
}

void greenStateLogic(StateMachine &fsm) {
  // Set green light ON, red and yellow OFF
  digitalWrite(RED_PIN, LOW);
  digitalWrite(GREEN_PIN, HIGH);
  digitalWrite(YELLOW_PIN, LOW);

  // Wait in this state for some time
  waitUntilElapsed(GREEN_DURATION, yellowStateCallback, fsm);
}

void yellowStateLogic(StateMachine &fsm) {
  // Set yellow light ON, red and green OFF
  digitalWrite(RED_PIN, LOW);
  digitalWrite(GREEN_PIN, LOW);
  digitalWrite(YELLOW_PIN, HIGH);

  // Wait in this state for some time
  waitUntilElapsed(YELLOW_DURATION, redStateCallback, fsm);
}

StateMachine fsm;

void setup() {
  Serial.begin(9600);
  pinMode(RED_PIN, OUTPUT);
  pinMode(GREEN_PIN, OUTPUT);
  pinMode(YELLOW_PIN, OUTPUT);

  // Start with RedState
  fsm.setStateWithLogic("RED_STATE", redStateLogic);
}

void loop() {
  // Check and transition to the next state if not suspended
  if (fsm.getCurrentStateName() == "RED_STATE") {
    checkAndTransition(redStateSuspended, redStateCallback, fsm);
  } else if (fsm.getCurrentStateName() == "GREEN_STATE") {
    checkAndTransition(greenStateSuspended, greenStateCallback, fsm);
  } else if (fsm.getCurrentStateName() == "YELLOW_STATE") {
    checkAndTransition(yellowStateSuspended, yellowStateCallback, fsm);
  }

  fsm.run(); // Start traffic light operation
}