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

// --- Pin Definitions ---
#define ENCODER_PIN_A 2
#define ENCODER_PIN_B 3
#define ENCODER_BTN_PIN 4

// --- OLED Display Configuration ---
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define OLED_RESET -1 // Reset pin # (or -1 if sharing Arduino reset pin)
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

// --- Timing Configuration ---
const unsigned long BLINK_INTERVAL = 500; // Blinking interval in ms
const unsigned long DEBOUNCE_TIME = 50;   // Button debounce time in ms
const unsigned long RETURN_DELAY = 300;   // Delay after returning to menu in ms

// --- Menu Configuration ---
const uint8_t MENU_ITEM_COUNT = 3; // W, V, A
const int MENU_CIRCLE_X_OFFSET = -28;
const int MENU_CIRCLE_RADIUS = 30;
const int MENU_BOX_WIDTH = 24;
const int MENU_BOX_HEIGHT = 20;
const int MENU_BOX_X_OFFSET = -5;
const int MENU_TEXT_X_OFFSET = 7;
const int MENU_TEXT_Y_OFFSET = 2;
const int MENU_BOX_Y_TOP = 2;

// --- Graph Configuration ---
#define GRAPH_WIDTH (SCREEN_WIDTH / 2) // Common width for all graphs
#define GRAPH_HEIGHT SCREEN_HEIGHT     // Common height for all graphs
#define GRAPH_POINTS 64                // Number of data points for graphs

// --- Back Arrow Configuration ---
const int BACK_ARROW_BOX_WIDTH = 18;
const int BACK_ARROW_BOX_HEIGHT = 18;
const int BACK_ARROW_BOX_X_OFFSET = -2;
const int BACK_ARROW_BOX_Y = 2;
const int BACK_ARROW_TEXT_X_OFFSET = 3;
const int BACK_ARROW_TEXT_Y_OFFSET = 2;

// --- Info Panel Configuration ---
const int INFO_PANEL_X_OFFSET = 4;
const int INFO_PANEL_Y = 24;
const int INFO_PANEL_LINE1_Y = 4;
const int INFO_PANEL_LINE2_Y = 16;
const int INFO_PANEL_VALUE_Y_OFFSET = 6;
const int INFO_PANEL_VALUE_LINE_SPACING = 12;


// --- Screen States ---
enum ScreenState {
    SCREEN_MENU,
    SCREEN_W_GRAPH,
    SCREEN_V_GRAPH,
    SCREEN_A_GRAPH
};
ScreenState currentScreen = SCREEN_MENU;

// --- Global Variables ---
volatile long encoderPosition = 0;
volatile uint8_t lastEncoded = 0; // Use uint8_t as it only needs 2 bits
uint8_t menuIndex = 0;            // 0:W, 1:V, 2:A
ezButton encoderBtn(ENCODER_BTN_PIN);

// Graph data arrays and indices
float wValues[GRAPH_POINTS] = {0};
uint8_t wIndex = 0;
float vValues[GRAPH_POINTS] = {0};
uint8_t vIndex = 0;
float aValues[GRAPH_POINTS] = {0};
uint8_t aIndex = 0;

// State management variables
bool displayNeedsUpdate = true; // Flag to control display updates
bool blinkState = false;
unsigned long lastBlinkTime = 0;
bool justEnteredScreen = false; // To prevent immediate exit from graph screens
unsigned long screenTransitionTime = 0; // For non-blocking return delay

// --- Function Prototypes ---
void handleEncoder();
void updateMenuScreen();
void updateWGraphScreen();
void updateVGraphScreen();
void updateAGraphScreen();
void drawMenu();
void drawWGraphScreen();
void drawVGraphScreen();
void drawAGraphScreen();
void drawGraphAxes(uint8_t graphWidth, uint8_t graphHeight);
void drawGraphData(float values[], uint8_t index, uint8_t points, uint8_t graphWidth, uint8_t graphHeight);
void drawBackArrowBox(bool blink);
void drawInfoPanel(float w, float v, float a, const char* formulaLine1, const char* formulaLine2, uint8_t primaryValueIndex);


// --- Interrupt Service Routine ---
void handleEncoder() {
    uint8_t MSB = digitalRead(ENCODER_PIN_A);
    uint8_t LSB = digitalRead(ENCODER_PIN_B);
    uint8_t encoded = (MSB << 1) | LSB;
    uint8_t sum = (lastEncoded << 2) | encoded;

    // Determine direction: CLOCKWISE
    if (sum == 0b1101 || sum == 0b0100 || sum == 0b0010 || sum == 0b1011) {
        encoderPosition++;
    }
    // Determine direction: COUNTER-CLOCKWISE
    if (sum == 0b1110 || sum == 0b0111 || sum == 0b0001 || sum == 0b1000) {
        encoderPosition--;
    }
    lastEncoded = encoded; // Store this state
}

// --- Setup Function ---
void setup() {
    pinMode(ENCODER_PIN_A, INPUT_PULLUP);
    pinMode(ENCODER_PIN_B, INPUT_PULLUP);

    // Read initial state for encoder
    uint8_t MSB = digitalRead(ENCODER_PIN_A);
    uint8_t LSB = digitalRead(ENCODER_PIN_B);
    lastEncoded = (MSB << 1) | LSB;

    // Attach interrupts
    attachInterrupt(digitalPinToInterrupt(ENCODER_PIN_A), handleEncoder, CHANGE);
    attachInterrupt(digitalPinToInterrupt(ENCODER_PIN_B), handleEncoder, CHANGE);

    // Initialize display
    if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
        Serial.println(F("SSD1306 allocation failed"));
        for (;;); // Don't proceed, loop forever
    }
    display.clearDisplay();
    display.display(); // Show initial clear screen

    // Initialize button
    encoderBtn.setDebounceTime(DEBOUNCE_TIME);

    // Initial draw
    drawMenu();
    displayNeedsUpdate = true; // Ensure the first draw happens in loop
}

// --- Main Loop ---
void loop() {
    unsigned long currentTime = millis();
    encoderBtn.loop(); // Update button state

    // --- Blinking Logic ---
    if (currentTime - lastBlinkTime >= BLINK_INTERVAL) {
        blinkState = !blinkState;
        lastBlinkTime = currentTime;
        displayNeedsUpdate = true; // Need to redraw blinking elements
    }

    // --- State Machine ---
    switch (currentScreen) {
        case SCREEN_MENU:
            updateMenuScreen(currentTime);
            break;
        case SCREEN_W_GRAPH:
            updateWGraphScreen(currentTime);
            break;
        case SCREEN_V_GRAPH:
            updateVGraphScreen(currentTime);
            break;
        case SCREEN_A_GRAPH:
            updateAGraphScreen(currentTime);
            break;
    }

    // --- Display Update ---
    // Only call display.display() if something changed
    if (displayNeedsUpdate) {
        // Clear display before redrawing based on current state
        display.clearDisplay();

        switch (currentScreen) {
            case SCREEN_MENU:
                drawMenu();
                break;
            case SCREEN_W_GRAPH:
                drawWGraphScreen();
                break;
            case SCREEN_V_GRAPH:
                drawVGraphScreen();
                break;
            case SCREEN_A_GRAPH:
                drawAGraphScreen();
                break;
        }
        display.display();
        displayNeedsUpdate = false; // Reset flag after update
    }
}

// --- Screen Update Functions ---

void updateMenuScreen(unsigned long currentTime) {
    static long lastEncoderPos = 0;

    // Prevent interaction shortly after returning from a graph screen
    if (currentTime - screenTransitionTime < RETURN_DELAY) {
        return;
    }

    // Handle encoder rotation for menu navigation
    // Divide by 4 for debouncing/filtering encoder steps
    long currentEncoderPos = encoderPosition / 4;
    if (currentEncoderPos != lastEncoderPos) {
        if (currentEncoderPos > lastEncoderPos) {
            menuIndex = (menuIndex + 1) % MENU_ITEM_COUNT;
        } else {
            // Modulo arithmetic for wrapping around backwards
            menuIndex = (menuIndex + MENU_ITEM_COUNT - 1) % MENU_ITEM_COUNT;
        }
        lastEncoderPos = currentEncoderPos;
        displayNeedsUpdate = true; // Menu selection changed
    }

    // Handle button press to enter graph screens
    if (encoderBtn.isPressed()) {
        justEnteredScreen = true; // Set flag for the new screen
        screenTransitionTime = currentTime; // Record time for potential immediate exit check

        switch (menuIndex) {
            case 0: currentScreen = SCREEN_W_GRAPH; break;
            case 1: currentScreen = SCREEN_V_GRAPH; break;
            case 2: currentScreen = SCREEN_A_GRAPH; break;
        }
        displayNeedsUpdate = true; // Screen changed
    }
}

void updateWGraphScreen(unsigned long currentTime) {
    // Sample new data (replace with actual sensor reading)
    float w = (sin(currentTime / 1000.0) + 1.0) / 2.0; // Example: 0 to 1 sine wave
    wValues[wIndex] = w;
    wIndex = (wIndex + 1) % GRAPH_POINTS;
    displayNeedsUpdate = true; // Data changed

    // Handle exit logic
    if (justEnteredScreen) {
        // Wait for button release after entering
        if (encoderBtn.isReleased()) {
            justEnteredScreen = false;
        }
    } else {
        // If button is pressed and released (after initial entry), return to menu
        if (encoderBtn.isReleased()) {
            currentScreen = SCREEN_MENU;
            screenTransitionTime = currentTime; // Record time for return delay
            displayNeedsUpdate = true; // Screen changed
        }
    }
}

void updateVGraphScreen(unsigned long currentTime) {
    // Sample new data (replace with actual sensor reading)
    float v = (cos(currentTime / 1200.0) + 1.0) / 2.0; // Example: 0 to 1 cosine wave
    vValues[vIndex] = v;
    vIndex = (vIndex + 1) % GRAPH_POINTS;
    displayNeedsUpdate = true; // Data changed

    // Handle exit logic (same as W graph)
    if (justEnteredScreen) {
        if (encoderBtn.isReleased()) {
            justEnteredScreen = false;
        }
    } else {
        if (encoderBtn.isReleased()) {
            currentScreen = SCREEN_MENU;
            screenTransitionTime = currentTime;
            displayNeedsUpdate = true;
        }
    }
}

void updateAGraphScreen(unsigned long currentTime) {
    // Sample new data (replace with actual sensor reading)
    float a = (sin(currentTime / 800.0) * 0.5 + 0.5); // Example: 0 to 1 sine wave (different freq)
    aValues[aIndex] = a;
    aIndex = (aIndex + 1) % GRAPH_POINTS;
    displayNeedsUpdate = true; // Data changed

    // Handle exit logic (same as W graph)
    if (justEnteredScreen) {
        if (encoderBtn.isReleased()) {
            justEnteredScreen = false;
        }
    } else {
        if (encoderBtn.isReleased()) {
            currentScreen = SCREEN_MENU;
            screenTransitionTime = currentTime;
            displayNeedsUpdate = true;
        }
    }
}


// --- Drawing Functions ---

void drawMenu() {
    // ----- Left side: Formula Circle -----
    int cx = GRAPH_WIDTH + MENU_CIRCLE_X_OFFSET; // Centered relative to graph width
    int cy = SCREEN_HEIGHT / 2;
    display.drawCircle(cx, cy, MENU_CIRCLE_RADIUS, SSD1306_WHITE);
    // Horizontal line (diameter)
    display.drawLine(cx - MENU_CIRCLE_RADIUS + 1, cy, cx + MENU_CIRCLE_RADIUS - 1, cy, SSD1306_WHITE);
    // Vertical line (radius)
    display.drawLine(cx, cy, cx, cy + MENU_CIRCLE_RADIUS -1 , SSD1306_WHITE);

    display.setTextSize(2);
    display.setTextColor(SSD1306_WHITE);
    // Position text relative to center
    display.setCursor(cx - 5, cy - 20); // V (Top)
    display.print("V");
    display.setCursor(cx - 15, cy + 5); // I (Bottom Left)
    display.print("I"); // Using I for current as is common in formulas
    display.setCursor(cx + 7, cy + 5);  // R (Bottom Right)
    display.print("R");

    // ----- Right side: W/V/A Selection Boxes -----
    int boxX = SCREEN_WIDTH - MENU_BOX_WIDTH + MENU_BOX_X_OFFSET;
    int yMid = MENU_BOX_Y_TOP + MENU_BOX_HEIGHT;
    int yBot = yMid + MENU_BOX_HEIGHT;

    // W Box
    bool isWSelected = (menuIndex == 0);
    display.setTextColor(isWSelected ? SSD1306_BLACK : SSD1306_WHITE);
    if (isWSelected && blinkState) { // Only fill selected box when blinking 'on'
         display.fillRect(boxX, MENU_BOX_Y_TOP, MENU_BOX_WIDTH, MENU_BOX_HEIGHT, SSD1306_WHITE);
    } else {
        display.drawRect(boxX, MENU_BOX_Y_TOP, MENU_BOX_WIDTH, MENU_BOX_HEIGHT, SSD1306_WHITE);
    }
    display.setCursor(boxX + MENU_TEXT_X_OFFSET, MENU_BOX_Y_TOP + MENU_TEXT_Y_OFFSET);
    display.print("W");

    // V Box
    bool isVSelected = (menuIndex == 1);
    display.setTextColor(isVSelected ? SSD1306_BLACK : SSD1306_WHITE);
     if (isVSelected && blinkState) {
         display.fillRect(boxX, yMid, MENU_BOX_WIDTH, MENU_BOX_HEIGHT, SSD1306_WHITE);
    } else {
        display.drawRect(boxX, yMid, MENU_BOX_WIDTH, MENU_BOX_HEIGHT, SSD1306_WHITE);
    }
    display.setCursor(boxX + MENU_TEXT_X_OFFSET, yMid + MENU_TEXT_Y_OFFSET);
    display.print("V");

    // A Box
    bool isASelected = (menuIndex == 2);
    display.setTextColor(isASelected ? SSD1306_BLACK : SSD1306_WHITE);
    if (isASelected && blinkState) {
         display.fillRect(boxX, yBot, MENU_BOX_WIDTH, MENU_BOX_HEIGHT, SSD1306_WHITE);
    } else {
        display.drawRect(boxX, yBot, MENU_BOX_WIDTH, MENU_BOX_HEIGHT, SSD1306_WHITE);
    }
    display.setCursor(boxX + MENU_TEXT_X_OFFSET, yBot + MENU_TEXT_Y_OFFSET);
    display.print("A");

    // Restore default text color for subsequent draws if needed elsewhere
    display.setTextColor(SSD1306_WHITE);
}


void drawWGraphScreen() {
    // Draw graph area
    drawGraphAxes(GRAPH_WIDTH, GRAPH_HEIGHT);
    drawGraphData(wValues, wIndex, GRAPH_POINTS, GRAPH_WIDTH, GRAPH_HEIGHT);

    // Draw back arrow (blinks when button not held after entry)
    drawBackArrowBox(!justEnteredScreen && blinkState);

    // --- Right Panel: Info ---
    // Get current values (replace with actual sensor reads/calculations)
    unsigned long currentTime = millis();
    float v = 3.45 + 0.5 * sin(currentTime / 1500.0); // Example V
    float a = 2.10 + 0.3 * cos(currentTime / 1200.0); // Example A
    float w = v * a;                                  // Calculated W

    drawInfoPanel(w, v, a, "W =", "V x A", 0); // 0 indicates W is primary
}

void drawVGraphScreen() {
    // Draw graph area
    drawGraphAxes(GRAPH_WIDTH, GRAPH_HEIGHT);
    drawGraphData(vValues, vIndex, GRAPH_POINTS, GRAPH_WIDTH, GRAPH_HEIGHT);

    // Draw back arrow
    drawBackArrowBox(!justEnteredScreen && blinkState);

    // --- Right Panel: Info ---
    unsigned long currentTime = millis();
    float w = 7.00 + 2.0 * sin(currentTime / 1500.0); // Example W
    float a = 2.10 + 0.3 * cos(currentTime / 1200.0); // Example A
    float v = (a != 0) ? w / a : 0;                   // Calculated V

    drawInfoPanel(w, v, a, "V =", "W / A", 1); // 1 indicates V is primary
}

void drawAGraphScreen() {
    // Draw graph area
    drawGraphAxes(GRAPH_WIDTH, GRAPH_HEIGHT);
    drawGraphData(aValues, aIndex, GRAPH_POINTS, GRAPH_WIDTH, GRAPH_HEIGHT);

    // Draw back arrow
    drawBackArrowBox(!justEnteredScreen && blinkState);

    // --- Right Panel: Info ---
    unsigned long currentTime = millis();
    float w = 7.00 + 2.0 * sin(currentTime / 1500.0); // Example W
    float v = 3.45 + 0.5 * cos(currentTime / 1200.0); // Example V
    float a = (v != 0) ? w / v : 0;                   // Calculated A

    drawInfoPanel(w, v, a, "A =", "W / V", 2); // 2 indicates A is primary
}


// --- Drawing Helper Functions ---

void drawGraphAxes(uint8_t graphWidth, uint8_t graphHeight) {
    // Draw border/box for the graph area
    display.drawRect(0, 0, graphWidth, graphHeight, SSD1306_WHITE);
    // Explicit X axis line (optional, as rect bottom edge covers it)
    // display.drawLine(0, graphHeight - 1, graphWidth - 1, graphHeight - 1, SSD1306_WHITE);
    // Explicit Y axis line (optional, as rect left edge covers it)
    // display.drawLine(0, 0, 0, graphHeight - 1, SSD1306_WHITE);
}

void drawGraphData(float values[], uint8_t index, uint8_t points, uint8_t graphWidth, uint8_t graphHeight) {
    int innerHeight = graphHeight - 2; // Available vertical pixels inside axes
    for (int i = 1; i < points && i < graphWidth; i++) { // Ensure we don't draw outside width
        // Calculate buffer indices, wrapping around
        int idx0 = (index + i - 1 + points) % points;
        int idx1 = (index + i + points) % points;

        // Scale float value (0.0 to 1.0) to pixel coordinates
        // Y=0 is top, Y=graphHeight-1 is bottom. Subtract from max Y.
        // Clamp values just in case they go slightly out of 0-1 range
        float val0 = constrain(values[idx0], 0.0, 1.0);
        float val1 = constrain(values[idx1], 0.0, 1.0);

        int y0 = (graphHeight - 1) - (int)(val0 * innerHeight);
        int y1 = (graphHeight - 1) - (int)(val1 * innerHeight);

        // Draw line segment - use graphWidth-1 as max X coordinate
        display.drawLine(i - 1, y0, i, y1, SSD1306_WHITE);
    }
}

void drawBackArrowBox(bool blink) {
    int boxX = SCREEN_WIDTH - BACK_ARROW_BOX_WIDTH + BACK_ARROW_BOX_X_OFFSET;
    if (blink) {
        display.fillRect(boxX, BACK_ARROW_BOX_Y, BACK_ARROW_BOX_WIDTH, BACK_ARROW_BOX_HEIGHT, SSD1306_WHITE);
    } else {
        display.drawRect(boxX, BACK_ARROW_BOX_Y, BACK_ARROW_BOX_WIDTH, BACK_ARROW_BOX_HEIGHT, SSD1306_WHITE);
    }
    display.setTextSize(2);
    display.setTextColor(blink ? SSD1306_BLACK : SSD1306_WHITE);
    display.setCursor(boxX + BACK_ARROW_TEXT_X_OFFSET, BACK_ARROW_BOX_Y + BACK_ARROW_TEXT_Y_OFFSET);
    display.print(""); // Left arrow character ← (Check if your font supports this)
                           // If not, consider drawing lines: display.drawLine(...)
    display.setTextColor(SSD1306_WHITE); // Restore default color
}

void drawInfoPanel(float w, float v, float a, const char* formulaLine1, const char* formulaLine2, uint8_t primaryValueIndex) {
    int infoX = GRAPH_WIDTH + INFO_PANEL_X_OFFSET;

    display.setTextSize(1);
    display.setTextColor(SSD1306_WHITE);

    // Formula Lines
    display.setCursor(infoX, INFO_PANEL_LINE1_Y);
    display.print(formulaLine1);
    display.setCursor(infoX, INFO_PANEL_LINE2_Y);
    display.print(formulaLine2);

    // Calculated Values (W, V, A)
    int valueY = INFO_PANEL_Y + INFO_PANEL_VALUE_Y_OFFSET;

    // Wattage
    display.setCursor(infoX, valueY);
    display.print("W: ");
    display.print(w, 2); // Print with 2 decimal places

    // Voltage
    valueY += INFO_PANEL_VALUE_LINE_SPACING;
    display.setCursor(infoX, valueY);
    display.print("V: ");
    display.print(v, 2);

    // Amperage
    valueY += INFO_PANEL_VALUE_LINE_SPACING;
    display.setCursor(infoX, valueY);
    display.print("A: ");
    display.print(a, 2);
}