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

// --- HARDWARE CONFIGURATION ---
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define OLED_RESET    -1 // Reset pin # (or -1 if sharing Arduino reset pin)
#define SCREEN_ADDRESS 0x3C // See datasheet for Address; 0x3D for 128x64, 0x3C for 128x32

// Pin Definitions for Raspberry Pi Pico (RP2040)
#define PIN_SDA 4
#define PIN_SCL 5

#define ENC_CLK 12
#define ENC_DT  13
#define ENC_SW  14

// NeoPixel Pin (any GPIO is fine, I've chosen 16)
#define PIN_NEOPIXEL 16
#define NUM_PIXELS   64 // For 8x8 WS2812B matrix

// --- TRACKER DATA ---
#define NUM_STEPS 16

struct Step {
  int8_t note; // MIDI Note (0-127), -1 for OFF
};

Step sequence[NUM_STEPS];

// --- STATE ---
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);
Adafruit_NeoPixel pixels(NUM_PIXELS, PIN_NEOPIXEL, NEO_GRB + NEO_KHZ800);

int currentStep = 0;
bool isEditing = false;
int scrollOffset = 0;

// Encoder State
volatile int encoderDelta = 0;
static uint8_t prevNextCode = 0;
static uint16_t store = 0;

// Button State
bool lastButtonState = HIGH;
unsigned long lastDebounceTime = 0;

// --- ENCODER INTERRUPT ---
// Robust Rotary Encoder reading
void readEncoder() {
  static int8_t rot_enc_table[] = {0,1,1,0,1,0,0,1,1,0,0,1,0,1,1,0};

  prevNextCode <<= 2;
  if (digitalRead(ENC_DT)) prevNextCode |= 0x02;
  if (digitalRead(ENC_CLK)) prevNextCode |= 0x01;
  prevNextCode &= 0x0f;

  // If valid state
  if (rot_enc_table[prevNextCode]) {
    store <<= 4;
    store |= prevNextCode;
    if ((store & 0xff) == 0x2b) encoderDelta--;
    if ((store & 0xff) == 0x17) encoderDelta++;
  }
}

void setup() {
  Serial.begin(115200);
  Serial.println(F("Starting."));

  // 1. Setup Encoder
  pinMode(ENC_CLK, INPUT_PULLUP);
  pinMode(ENC_DT, INPUT_PULLUP);
  pinMode(ENC_SW, INPUT_PULLUP);
  
  attachInterrupt(digitalPinToInterrupt(ENC_CLK), readEncoder, CHANGE);
  attachInterrupt(digitalPinToInterrupt(ENC_DT), readEncoder, CHANGE);

  // 2. Setup Display
  Wire.setSDA(PIN_SDA);
  Wire.setSCL(PIN_SCL);
  Wire.begin();

  // SSD1306_SWITCHCAPVCC = generate display voltage from 3.3V internally
  if(!display.begin(SSD1306_SWITCHCAPVCC, SCREEN_ADDRESS)) {
    Serial.println(F("SSD1306 allocation failed"));
    for(;;); // Don't proceed, loop forever
  }

  // 3. Setup NeoPixel Matrix
  pixels.begin();
  pixels.setBrightness(40); // Set brightness to a medium-low value (0-255)
  pixels.clear();
  pixels.show();

  // 4. Init Sequence
  for(int i=0; i<NUM_STEPS; i++) {
    sequence[i].note = -1; // Default to empty
  }
  // Add a simple C-Major scale for testing
  sequence[0].note = 60; // C4
  sequence[2].note = 62; // D4
  sequence[4].note = 64; // E4
  sequence[6].note = 65; // F4

  display.clearDisplay();
  display.display();

  Serial.println(F("Started."));
}

void handleInput() {
  // Handle Encoder Rotation
  int delta = 0;
  noInterrupts();
  delta = encoderDelta;
  encoderDelta = 0;
  interrupts();

  if (delta != 0) {
    if (isEditing) {
      // Change Note
      int newNote = sequence[currentStep].note + delta;
      if (newNote < -1) newNote = -1;
      if (newNote > 127) newNote = 127;
      sequence[currentStep].note = newNote;
    } else {
      // Move Cursor
      currentStep += (delta > 0 ? 1 : -1);
      if (currentStep < 0) currentStep = NUM_STEPS - 1;
      if (currentStep >= NUM_STEPS) currentStep = 0;
      
      // Adjust Scroll to keep cursor in view
      if (currentStep < scrollOffset) scrollOffset = currentStep;
      if (currentStep >= scrollOffset + 6) scrollOffset = currentStep - 5;
    }
  }

  // Handle Button
  int reading = digitalRead(ENC_SW);
  if (reading != lastButtonState) {
    lastDebounceTime = millis();
  }
  
  if ((millis() - lastDebounceTime) > 50) {
    if (reading == LOW) { // Button Pressed
       // Wait for release to toggle mode
       while(digitalRead(ENC_SW) == LOW); 
       isEditing = !isEditing;
    }
  }
  lastButtonState = reading;
}

void drawUI() {
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(0, 0);

  // Header
  display.print(F("TRACKER  "));
  display.print(isEditing ? F("[EDIT]") : F("[NAV]"));
  display.println();
  display.drawLine(0, 8, 128, 8, SSD1306_WHITE);

  // Steps
  int y = 10;
  for (int i = scrollOffset; i < min(scrollOffset + 6, NUM_STEPS); i++) {
    
    // Draw Cursor
    if (i == currentStep) {
      display.fillRect(0, y, 128, 8, SSD1306_WHITE);
      display.setTextColor(SSD1306_BLACK, SSD1306_WHITE); // Invert text
    } else {
      display.setTextColor(SSD1306_WHITE);
    }

    display.setCursor(2, y);
    
    // Step Number
    if (i < 10) display.print(F("0"));
    display.print(i);
    display.print(F(" | "));

    // Note Value
    int n = sequence[i].note;
    if (n == -1) {
      display.print(F("---"));
    } else {
      // Basic Note to String conversion
      const char* noteNames[] = {"C-", "C#", "D-", "D#", "E-", "F-", "F#", "G-", "G#", "A-", "A#", "B-"};
      display.print(noteNames[n % 12]);
      display.print(n / 12 - 1); // Octave
    }

    y += 9;
  }

  display.display();
}

// Helper to convert X,Y to pixel index for an 8x8 matrix.
// Assumes row-major wiring (NOT serpentine).
// If your matrix is wired differently, you'll need to change this function.
int getPixelIndex(int x, int y) {
  return y * 8 + x;
}

void updateLeds() {
  pixels.clear(); // Clear buffer

  for (int s = 0; s < NUM_STEPS; s++) {
    int blockX = (s % 4) * 2;
    int blockY = (s / 4) * 2;

    uint32_t color;

    if (s == currentStep) {
      if (isEditing) {
        color = pixels.Color(50, 0, 0); // Dim Red for editing
      } else {
        color = pixels.Color(40, 40, 40); // Dim White for current step
      }
    } else {
      if (sequence[s].note != -1) {
        color = pixels.Color(0, 0, 50); // Dim Blue for step with note
      } else {
        color = 0; // Off
      }
    }

    // Set the 4 pixels for the 2x2 block
    pixels.setPixelColor(getPixelIndex(blockX, blockY), color);
    pixels.setPixelColor(getPixelIndex(blockX + 1, blockY), color);
    pixels.setPixelColor(getPixelIndex(blockX, blockY + 1), color);
    pixels.setPixelColor(getPixelIndex(blockX + 1, blockY + 1), color);
  }

  pixels.show();
}

void loop() {
  handleInput();
  drawUI();
  updateLeds();
  delay(10); // Small delay to prevent screen tearing/excessive refresh
}