// --- Pin Definitions ---

// Array for the 20 puzzle buttons (Inputs)
const int buttonPins[20] = {
  11, 9, 7, 5, 3, 14, 16, 18, 20, 22,
  25, 29, 33, 37, 41, 45, 49, 53, A14, A12
};

// Array for the 20 corresponding relays (Outputs)
const int relayPins[20] = {
  12, 10, 8, 6, 4, 2, 15, 17, 19, 21,
  23, 27, 31, 35, 39, 43, 47, 51, A15, A13
};

// Control Buttons and Sensors (Inputs)
const int START_BUTTON_PIN = A0;
const int RESET_BUTTON_PIN = A1;
const int LOCK_SENSOR_PIN = A8;

// Relay and Status LED (Outputs)
const int RELAY_3_PIN = A2; // Puzzle Solved Relay
const int RELAY_4_PIN = A3; // Electronic Lock / Incorrect Relay
const int STATUS_LED_PIN = 13; // "Game Status" LED

// --- Constants ---
const int NUM_BUTTONS = 20;
const int NUM_BUTTONS_TO_CHANGE = 5; // Defines exactly how many buttons must be flipped
const long RELAY_4_DURATION = 1000; // 1 second for electronic lock
const long RELAY_3_DURATION = 5000; // 5 seconds for puzzle solved

// Relay States (Active LOW modules)
const int RELAY_ON = LOW;
const int RELAY_OFF = HIGH;

// Button States (internal pullup resistors)
const int BUTTON_PRESSED = LOW;
const int BUTTON_RELEASED = HIGH;

// --- Game States ---
enum GameState {
  IDLE,
  STARTING,
  ACTIVE,
  SOLVED
};
GameState currentState = IDLE;

// --- Global Variables ---
bool solution[NUM_BUTTONS];           // Stores the correct state for each button
bool initialButtonStates[NUM_BUTTONS]; // Stores the physical state of buttons at game start

// Timers for non-blocking delays
unsigned long relay4_timer_start = 0;
unsigned long relay3_timer_start = 0;

void setup() {
  Serial.begin(9600); // For debugging
  Serial.println("Serial monitor init");

  // Initialize pins for puzzle buttons and relays
  for (int i = 0; i < NUM_BUTTONS; i++) {
    pinMode(buttonPins[i], INPUT_PULLUP);
    pinMode(relayPins[i], OUTPUT);
    digitalWrite(relayPins[i], RELAY_OFF); // Ensure all relays start OFF
  }

  // Initialize control pins
  pinMode(START_BUTTON_PIN, INPUT_PULLUP);
  pinMode(RESET_BUTTON_PIN, INPUT_PULLUP);
  pinMode(LOCK_SENSOR_PIN, INPUT_PULLUP);

  // Initialize output pins
  pinMode(RELAY_3_PIN, OUTPUT);
  pinMode(RELAY_4_PIN, OUTPUT);
  pinMode(STATUS_LED_PIN, OUTPUT);

  // Ensure control relays are off at the start
  digitalWrite(RELAY_3_PIN, RELAY_OFF);
  digitalWrite(RELAY_4_PIN, RELAY_OFF);

  // Seed the random number generator from an unused analog pin
  randomSeed(analogRead(A6));

  // Set initial state
  enterIdleState();
}

void loop() {
  // Run functions based on the current game state
  switch (currentState) {
    case IDLE:
      handleIdleState();
      break;
    case STARTING:
      handleStartingState();
      break;
    case ACTIVE:
      handleActiveState();
      break;
    case SOLVED:
      handleSolvedState();
      break;
  }

  // Update timers for relays regardless of state
  updateRelayTimers();
}

// --- State Handling Functions ---

void enterIdleState() {
  Serial.println("Entering IDLE state.");
  currentState = IDLE;
  digitalWrite(STATUS_LED_PIN, LOW); // Turn off active status LED
  setAllRelays(RELAY_ON); // Turn all puzzle relays ON
}

void handleIdleState() {
  // Wait for the start button to be pressed
  if (digitalRead(START_BUTTON_PIN) == BUTTON_PRESSED) {
    delay(50); // Simple debounce
    if (digitalRead(START_BUTTON_PIN) == BUTTON_PRESSED) {
      currentState = STARTING;
    }
  }
}

void handleStartingState() {
  Serial.println("Entering STARTING state.");

  // Read the current physical state of all buttons
  Serial.println("Reading initial button states...");
  for (int i = 0; i < NUM_BUTTONS; i++) {
    initialButtonStates[i] = (digitalRead(buttonPins[i]) == BUTTON_PRESSED);
  }

  // 1. Blink all relays 3 times
  for (int i = 0; i < 3; i++) {
    setAllRelays(RELAY_ON);
    delay(250);
    setAllRelays(RELAY_OFF);
    delay(250);
  }

  // 2. Activate Relay 4 for 1 second
  activateRelay4();

  // 3. Generate the new puzzle solution that requires 5 buttons to be flipped
  generatePuzzle();

  // 4. Turn on the game status LED
  digitalWrite(STATUS_LED_PIN, HIGH);

  // 5. Update relays to reflect the starting position.
  updateRelayStatus();

  // 6. Transition to the active game state
  Serial.println("Entering ACTIVE state.");
  currentState = ACTIVE;
}

void handleActiveState() {
  // Check for reset button press
  if (digitalRead(RESET_BUTTON_PIN) == BUTTON_PRESSED) {
    delay(50); // Debounce
    if (digitalRead(RESET_BUTTON_PIN) == BUTTON_PRESSED) {
        enterIdleState();
        return;
    }
  }

  // Continuously check all buttons and update their corresponding relays
  int correctCount = updateRelayStatus();

  // Check for win/fail condition from the lock sensor
  if (digitalRead(LOCK_SENSOR_PIN) == BUTTON_PRESSED) {
    delay(50); // Debounce
    if (digitalRead(LOCK_SENSOR_PIN) == BUTTON_PRESSED) {
      if (correctCount == NUM_BUTTONS) {
        // --- WIN CONDITION ---
        Serial.println("Puzzle Solved!");
        currentState = SOLVED;
      } else {
        // --- FAIL CONDITION  ---
        Serial.println("Not all buttons are correct.");
        activateRelay4(); // Activate incorrect relay
        // Debounce lock
        while(digitalRead(LOCK_SENSOR_PIN) == BUTTON_PRESSED) {
            delay(10);
        }
      }
    }
  }
}

void handleSolvedState() {
  // This state starts the final reward sequence
  Serial.println("Entering SOLVED state.");
  digitalWrite(STATUS_LED_PIN, LOW); // Turn off active game LED
  
  // Turn on all puzzle relays to show completion
  setAllRelays(RELAY_ON);

  // Activate Relay 3 for 5 seconds
  activateRelay3();
  
  // After the solved sequence, wait until the relay is off before going to IDLE
  while (relay3_timer_start > 0) {
    updateRelayTimers();
  }

  enterIdleState();
}

void generatePuzzle() {
  Serial.println("Generating new 5-flip puzzle...");
  
  // The solution starts as the exact copy of the initial button states.
  for (int i = 0; i < NUM_BUTTONS; i++) {
    solution[i] = initialButtonStates[i];
  }

  // Create a list of button indices (0 to 19) to choose from.
  int buttonIndices[NUM_BUTTONS];
  for (int i = 0; i < NUM_BUTTONS; i++) {
    buttonIndices[i] = i;
  }

  // Shuffle the indices randomly
  for (int i = NUM_BUTTONS - 1; i > 0; i--) {
    int j = random(i + 1);
    int temp = buttonIndices[i];
    buttonIndices[i] = buttonIndices[j];
    buttonIndices[j] = temp;
  }

  // Take the first 5 indices from the shuffled list and flip their state in the solution array.
  Serial.print("Flipping buttons at indices: ");
  for (int i = 0; i < NUM_BUTTONS_TO_CHANGE; i++) {
    int indexToFlip = buttonIndices[i];
    solution[indexToFlip] = !solution[indexToFlip]; // Flip the boolean value
    Serial.print(indexToFlip);
    Serial.print(" ");
  }
  Serial.println();

  // For debugging: print the final, valid solution
  Serial.print("Final puzzle solution: ");
  for (int i = 0; i < NUM_BUTTONS; i++) {
    Serial.print(solution[i] ? "1" : "0");
  }
  Serial.println();
}

int updateRelayStatus() {
  int correctCount = 0;
  for (int i = 0; i < NUM_BUTTONS; i++) {
    bool isButtonOn = (digitalRead(buttonPins[i]) == BUTTON_PRESSED);
    
    // Check if the button's current state matches the required solution state
    if (isButtonOn == solution[i]) {
      digitalWrite(relayPins[i], RELAY_ON); // Turn relay ON for correct state
      correctCount++;
    } else {
      digitalWrite(relayPins[i], RELAY_OFF); // Turn relay OFF for incorrect state
    }
  }
  return correctCount;
}

void setAllRelays(int state) {
  for (int i = 0; i < NUM_BUTTONS; i++) {
    digitalWrite(relayPins[i], state);
  }
}

void activateRelay4() {
  digitalWrite(RELAY_4_PIN, RELAY_ON);
  relay4_timer_start = millis(); // Start the timer
}

void activateRelay3() {
  digitalWrite(RELAY_3_PIN, RELAY_ON);
  relay3_timer_start = millis(); // Start the timer
}

void updateRelayTimers() {
  // Check Timer for Relay 4
  if (relay4_timer_start > 0) {
    if (millis() - relay4_timer_start >= RELAY_4_DURATION) {
      digitalWrite(RELAY_4_PIN, RELAY_OFF);
      relay4_timer_start = 0; // Stop the timer
    }
  }

  // Check Timer for Relay 3
  if (relay3_timer_start > 0) {
    if (millis() - relay3_timer_start >= RELAY_3_DURATION) {
      digitalWrite(RELAY_3_PIN, RELAY_OFF);
      relay3_timer_start = 0; // Stop the timer
    }
  }
}