#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>

#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define OLED_RESET -1
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

#define TEST_BUTTON_PIN 46

const int sideA[] = {0, 1, 2, 3, 4, 5, 6, 7, 10, 11, 12, 13, 14, 15, 16, 17};
const int sideB[] = {45, 42, 41, 40, 39, 38, 37, 36, 35, 34, 33, 26, 21, 20, 19, 18};

bool testing = false;  // Flag to track if we are testing

void setup() {
  // Initialize Serial and OLED display
  Serial.begin(115200);
  if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    Serial.println(F("SSD1306 allocation failed"));
    for (;;);
  }
  display.clearDisplay();
  
  // Set Side A as outputs and Side B as inputs with pull-down resistors
  for (int i = 0; i < 16; i++) {
    pinMode(sideA[i], OUTPUT);
    pinMode(sideB[i], INPUT_PULLDOWN);
  }
  
  pinMode(TEST_BUTTON_PIN, INPUT_PULLUP);
  display.setTextSize(1);
  display.setTextColor(WHITE);

  // Display initial welcome screen
  showWelcomeScreen();
}

void showWelcomeScreen() {
  display.clearDisplay();
  display.setCursor(20, 25);
  display.print("Ribbon Tester");
  display.display();
}

void showResults(bool cutDetected[], bool shortDetected[]) {
  display.clearDisplay();
  
  // Left Column (Pins 1-8)
  for (int i = 0; i < 8; i++) {
    display.setCursor(0, i * 8);  // Each row is 8 pixels in height
    display.print("P ");
    display.print(i + 1);
    display.print(": ");
    if (cutDetected[i]) display.print("CUT");
    else if (shortDetected[i]) display.print("SRT"); 
    else display.print("OK");
  }

  // Right Column (Pins 9-16)
  for (int i = 8; i < 16; i++) {
    display.setCursor(64, (i - 8) * 8);  // Offset x to right column, same y offset
    display.print("P ");
    display.print(i + 1);
    display.print(": ");
    if (cutDetected[i]) display.print("CUT");
    else if (shortDetected[i]) display.print("SRT"); 
    else display.print("OK");
  }
  
  display.display();
}

void testRibbonCable() {
  bool cutDetected[16] = {false};
  bool shortDetected[16] = {false};

  // Test each pin
  for (int i = 0; i < 16; i++) {
    digitalWrite(sideA[i], HIGH);
    delay(10);  // Brief delay for stability

    // Check Side B for continuity and shorts
    for (int j = 0; j < 16; j++) {
      int readB = digitalRead(sideB[j]);
      if (i == j) {
        if (readB == LOW) {
          cutDetected[i] = true;  // Continuity check failed
        }
      } else if (readB == HIGH) {
        shortDetected[j] = true;  // Short detected
      }
    }

    digitalWrite(sideA[i], LOW);  // Reset pin on Side A
  }

  // Display results in two columns
  showResults(cutDetected, shortDetected);
}

void loop() {
  // Wait for button press to start or reset test
  if (digitalRead(TEST_BUTTON_PIN) == LOW) {
    delay(200);  // Debounce delay
    while (digitalRead(TEST_BUTTON_PIN) == LOW);  // Wait for button release

    if (!testing) {
      // Start testing
      testing = true;
      display.clearDisplay();
      display.setCursor(0, 0);
      display.print("Testing...");
      display.display();
      testRibbonCable();
    } else {
      // Reset to welcome screen after test
      testing = false;
      showWelcomeScreen();
    }
  }
}