#include <Wire.h>                   // Must include Wire for I2C
#include <Adafruit_GFX.h>           // Base graphics library
#include <Adafruit_SSD1306.h>       // OLED library
#include "icons.h"                  // Your custom icons

// Hardware Pins mapped for Wokwi simulation
#define BTN_WHITE 2   // Pin 2: White Button
#define BTN_GREEN 3   // Pin 3: Green Button
#define OLED_SDA 4    // Pin 4: OLED SDA
#define OLED_SCL 5    // Pin 5: OLED SCK
#define ENC_SW 14     // Pin 14: Encoder Switch
#define ENC_DT 15     // Pin 15: Encoder DT
#define ENC_CLK 26    // Pin 26: Encoder CLK

// Simulation-friendly pin remappings for Wokwi testing:
#define CHARGE_PIN 28   // Slide switch for charge status (HIGH/LOW)
#define BAT_ADC_PIN 27  // Potentiometer connected to GP27 for smooth voltage testing

const unsigned char *bluetoothIcons[] = { epd_bitmap_BluetoothOFF, epd_bitmap_BluetoothConnected };
const unsigned char *irIcons[] = { epd_bitmap_IRreceive, epd_bitmap_IRTransmit };
const unsigned char *keyboardIcons[] = { epd_bitmap_Keyboard, epd_bitmap_KeyboardInverted };
const unsigned char *modeIcons[] = { epd_bitmap_VolumeMute, epd_bitmap_VolumeDown, epd_bitmap_VolumeUP, epd_bitmap_Scroll };
const unsigned char *batteryIcons[] = { 
    epd_bitmap_BatteryEmptyNPI, epd_bitmap_BatteryEmptyl, 
    epd_bitmap_Battery1barNPI, epd_bitmap_Battery1bar, 
    epd_bitmap_Battery2barNPI, epd_bitmap_Battery2bar, 
    epd_bitmap_Battery3barNPI, epd_bitmap_Battery3bar, 
    epd_bitmap_BatteryFullNPI, epd_bitmap_BatteryFull 
};

// Initialize the display object globally
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 32
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, -1);

enum IconSlot { SLOT_BLUETOOTH, SLOT_IR, SLOT_KEYBOARD, SLOT_MODE, SLOT_BATTERY, SLOT_COUNT };

struct Slot {
    int x, y;
    const unsigned char **icons;
    int count, current;
    void next() { current++; if (current >= count) current = 0; }
    void draw() { display.drawBitmap(x, y, icons[current], 16, 16, SSD1306_WHITE); }
};

// Define the slot array
Slot slots[SLOT_COUNT] = {
    {0, 0, bluetoothIcons, 2, 0},
    {28, 0, irIcons, 2, 0},
    {56, 0, keyboardIcons, 2, 0},
    {84, 0, modeIcons, 4, 0},
    {112, 0, batteryIcons, 10, 0}
};

void drawSlots() {
    display.clearDisplay();
    for(int i = 0; i < SLOT_COUNT; i++) slots[i].draw();
    display.display();
}

void updateBatteryStatus() {
    int rawCharge = digitalRead(CHARGE_PIN); 
    int rawBat = analogRead(BAT_ADC_PIN);    
    
    bool isPluggedIn = (rawCharge == HIGH); 
    
    // Scale potentiometer input (0 - 4095) directly to a simulated 0.0V - 4.2V scale
    float voltage = (rawBat / 4095.0) * 4.2; 

    // Debug output to Serial1 terminal
    Serial1.print("Raw Bat: "); Serial1.print(rawBat);
    Serial1.print(" | Calc Voltage: "); Serial1.println(voltage);

    int newIconIndex = 0;
    if (voltage > 4.1) newIconIndex = isPluggedIn ? 9 : 8;      
    else if (voltage > 3.9) newIconIndex = isPluggedIn ? 7 : 6; 
    else if (voltage > 3.8) newIconIndex = isPluggedIn ? 5 : 4; 
    else if (voltage > 3.7) newIconIndex = isPluggedIn ? 3 : 2; 
    else newIconIndex = isPluggedIn ? 1 : 0;                    

    // Only update the display if the battery icon actually changes to prevent flicker
    if (slots[SLOT_BATTERY].current != newIconIndex) {
        slots[SLOT_BATTERY].current = newIconIndex;
        drawSlots();
    }
}

void setup() {
  Serial1.begin(115200);
  delay(1000);
  while (!Serial1 && millis() < 5000) { delay(100); } 
  Serial1.println("--- Starting Setup ---");
  
  analogReadResolution(12); // Ensure RP2040 reads a full 12-bit range (0 - 4095)
  
  pinMode(BTN_WHITE, INPUT_PULLUP);
  pinMode(BTN_GREEN, INPUT_PULLUP);
  pinMode(ENC_SW, INPUT_PULLUP);
  pinMode(ENC_CLK, INPUT_PULLUP);
  pinMode(ENC_DT, INPUT_PULLUP);
  pinMode(CHARGE_PIN, INPUT_PULLDOWN); 
  
  Wire.setSDA(OLED_SDA);
  Wire.setSCL(OLED_SCL);
  Wire.begin();
  
  if(!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) { 
    Serial1.println("DISPLAY FAILED!"); 
    while(1); 
  }
  
  display.clearDisplay();
  display.display(); 
  drawSlots();       
  Serial1.println("Setup Complete & Display Rendered!");
}

void loop() {
  // 1. Receive Display/Icon Commands
  if (Serial.available() || Serial1.available()) {
    String cmd = Serial.available() ? Serial.readStringUntil('\n') : Serial1.readStringUntil('\n');
    cmd.trim();
    
    if (cmd.startsWith("D:")) {
      int comma = cmd.indexOf(',');
      int slot = cmd.substring(2, comma).toInt();
      int icon = cmd.substring(comma + 1).toInt();
      if (slot >= 0 && slot < SLOT_COUNT) {
        slots[slot].current = icon;
        drawSlots();
        Serial1.println("Updated Slot " + String(slot) + " to icon " + String(icon));
      }
    }
  }

  // 2. Report Buttons
  if (digitalRead(BTN_WHITE) == LOW) {
    Serial1.println("B:WHITE,1");
    delay(200);
  }
  
  if (digitalRead(BTN_GREEN) == LOW) {
    Serial1.println("B:GREEN,1");
    delay(200);
  }

  // 3. Update Battery Status Periodically (every 1 second for smooth potentiometer testing)
  static unsigned long lastBatteryUpdate = 0;
  if (millis() - lastBatteryUpdate > 1000) {
      updateBatteryStatus();
      lastBatteryUpdate = millis();
  }

  // 4. Rotary Encoder Rotation Logic
  static int lastClk = HIGH;
  int currentClk = digitalRead(ENC_CLK);
  if (currentClk != lastClk && currentClk == LOW) {
      if (digitalRead(ENC_DT) != currentClk) {
          Serial1.println("E:CW");
      } else {
          Serial1.println("E:CCW");
      }
  }
  lastClk = currentClk;

  // 5. Rotary Encoder Push Button Logic
  static bool lastEncBtnState = HIGH;
  bool currentEncBtnState = digitalRead(ENC_SW);
  if (currentEncBtnState != lastEncBtnState) {
      if (currentEncBtnState == LOW) {
          Serial1.println("B:ENCODER,1");
      }
      delay(50); 
      lastEncBtnState = currentEncBtnState;
  }
}