#include <SPI.h>
#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include <Adafruit_NeoPixel.h>
#include <EEPROM.h>
#include <pico/mutex.h>
#include "TrackerTypes.h"
#include "MelodyStrategy.h"
#include "LuckyStrategy.h"
#include "ArpStrategy.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

// MIDI UART Pins (GP0/GP1) -- OUT only so far
#define PIN_MIDI_TX 0

// 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

Step sequence[NUM_STEPS];
Step nextSequence[NUM_STEPS];
volatile bool sequenceChangeScheduled = false;
volatile bool needsPanic = false;

mutex_t midiMutex;

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

enum UIState {
  UI_TRACKER,
  UI_MENU_MAIN,
  UI_MENU_RANDOMIZE,
  UI_MENU_SETUP,
  UI_SETUP_CHANNEL_EDIT,
  UI_EDIT_TEMPO,
  UI_EDIT_FLAVOUR
};

UIState currentState = UI_MENU_MAIN;

const char* mainMenu[] = { "Tracker", "Randomize", "Setup" };
const int mainMenuCount = sizeof(mainMenu) / sizeof(char*);
const char* randomizeMenu[] = { "Back", "Scale", "Melody", "Flavour", "Tempo", "Mutation", "Song Mode", "Theme 1", "Theme 2", "Theme 3", "Theme 4", "Theme 5", "Theme 6", "Theme 7" };
const int THEME_1_INDEX = 7;
const int randomizeMenuCount = sizeof(randomizeMenu) / sizeof(char*);
const char* setupMenu[] = { "Back", "Channel", "Save", "Load" };
const int setupMenuCount = sizeof(setupMenu) / sizeof(char*);

int menuSelection = 0;
volatile int navigationSelection = 1;
volatile int playbackStep = 0;
int midiChannel = 1;
volatile int shMidiChannel = midiChannel;
int scaleNotes[12];
int numScaleNotes = 0;
int melodySeed = 0;
volatile int queuedTheme = -1;
volatile int currentThemeIndex = 1;
const uint32_t EEPROM_MAGIC = 0x42424246;

MelodyStrategy* strategies[] = { new LuckyStrategy(), new ArpStrategy() };
const int numStrategies = 2;
int currentStrategyIndex = 0;

volatile bool mutationEnabled = false;
volatile bool songModeEnabled = false;
volatile int songRepeatsRemaining = 0;
volatile int nextSongRepeats = 0;
volatile bool songModeNeedsNext = false;
volatile bool isEditing = false;
volatile int scrollOffset = 0;
volatile bool isPlaying = false;
volatile int tempo = 120; // BPM
volatile unsigned long lastClockTime = 0;
volatile int clockCount = 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;
bool buttonActive = false;
bool buttonConsumed = false;
unsigned long buttonPressTime = 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 sortArray(int arr[], int size) {
    for (int i = 0; i < size - 1; i++) {
        for (int j = 0; j < size - i - 1; j++) {
            if (arr[j] > arr[j + 1]) {
                int temp = arr[j];
                arr[j] = arr[j + 1];
                arr[j + 1] = temp;
            }
        }
    }
}

void showMessage(const char* msg) {
  display.clearDisplay();
  display.setCursor(10, 25);
  display.setTextColor(SSD1306_WHITE);
  display.setTextSize(2);
  display.print(msg);
  display.display();
  delay(500);
  display.setTextSize(1);
}

void saveSequence(bool quiet = false) {
  int addr = 0;
  EEPROM.put(addr, EEPROM_MAGIC); addr += sizeof(EEPROM_MAGIC);
  EEPROM.put(addr, midiChannel); addr += sizeof(midiChannel);
  EEPROM.put(addr, melodySeed); addr += sizeof(melodySeed);
  EEPROM.put(addr, currentStrategyIndex); addr += sizeof(currentStrategyIndex);
  EEPROM.put(addr, (int)tempo); addr += sizeof(int);

  EEPROM.put(addr, numScaleNotes); addr += sizeof(numScaleNotes);
  for (int i=0; i<12; i++) {
    EEPROM.put(addr, scaleNotes[i]); addr += sizeof(int);
  }

  mutex_enter_blocking(&midiMutex);
  for (int i=0; i<NUM_STEPS; i++) {
    EEPROM.put(addr, sequence[i]); addr += sizeof(Step);
  }
  mutex_exit(&midiMutex);
  EEPROM.commit();
  if (!quiet) showMessage("SAVED!");
}

bool loadSequence() {
  int addr = 0;
  uint32_t magic;
  EEPROM.get(addr, magic); addr += sizeof(magic);
  if (magic != EEPROM_MAGIC) return false;
  
  EEPROM.get(addr, midiChannel); addr += sizeof(midiChannel);
  shMidiChannel = midiChannel;
  EEPROM.get(addr, melodySeed); addr += sizeof(melodySeed);
  EEPROM.get(addr, currentStrategyIndex); addr += sizeof(currentStrategyIndex);
  int t;
  EEPROM.get(addr, t); addr += sizeof(int);
  tempo = t;

  EEPROM.get(addr, numScaleNotes); addr += sizeof(numScaleNotes);
  for (int i=0; i<12; i++) {
    EEPROM.get(addr, scaleNotes[i]); addr += sizeof(int);
  }

  mutex_enter_blocking(&midiMutex);
  for (int i=0; i<NUM_STEPS; i++) {
    EEPROM.get(addr, sequence[i]); addr += sizeof(Step);
  }
  mutex_exit(&midiMutex);
  return true;
}

void setup() {
  Serial.begin(115200);
  mutex_init(&midiMutex);

  // Use random ADC noise for seed
  delay(5000);
  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
  // Note: Using default I2C pins (SDA=GP4, SCL=GP5) which works on both cores.
  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. Setup MIDI Serial
  Serial1.setTX(PIN_MIDI_TX);
  Serial1.begin(31250);
  Serial.println(F("MIDI Serial initialized on GP0/GP1"));

  // 5. Init Sequence
  randomSeed(micros());
  generateRandomScale();
  melodySeed = random(10000);
  
  EEPROM.begin(512);
  if (!loadSequence()) {
    generateTheme(1);
  }
  isPlaying = false; // Don't start playing on boot

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

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

void sendMidi(uint8_t status, uint8_t note, uint8_t velocity) {
  uint8_t channelStatus = status | (shMidiChannel - 1);
  Serial1.write(channelStatus);
  Serial1.write(note);
  Serial1.write(velocity);
}

void sendMidiRealtime(uint8_t status) {
  mutex_enter_blocking(&midiMutex);
  Serial1.write(status);
  mutex_exit(&midiMutex);
}

void generateRandomScale() {
  numScaleNotes = random(3, 13); // 3 to 12 notes
  for (int i = 0; i < 12; i++) {
    scaleNotes[i] = i; // Fill with all notes
  }
  // Shuffle
  for (int i = 0; i < 12; i++) {
    int j = random(12);
    int temp = scaleNotes[i];
    scaleNotes[i] = scaleNotes[j];
    scaleNotes[j] = temp;
  }
  sortArray(scaleNotes, numScaleNotes);
}

void generateSequenceData(int themeType, Step* target) {
  randomSeed(melodySeed + themeType * 12345); // Deterministic seed for this theme
  strategies[currentStrategyIndex]->generate(target, NUM_STEPS, scaleNotes, numScaleNotes, melodySeed + themeType * 12345);
}

void generateTheme(int themeType) {
  currentThemeIndex = themeType;
  needsPanic = true;
  mutex_enter_blocking(&midiMutex);
  generateSequenceData(themeType, sequence);
  mutex_exit(&midiMutex);

  clockCount = 0;
  lastClockTime = micros();
  playbackStep = 0;
  sendMidiRealtime(0xFA); // MIDI Start
  isPlaying = true;
}

void mutateSequence(Step* target) {
  strategies[currentStrategyIndex]->mutate(target, NUM_STEPS, scaleNotes, numScaleNotes);
}

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

  if (delta != 0) {
    switch(currentState) {
      case UI_TRACKER:
        if (isEditing && navigationSelection > 0) {
          // Change Note
          int stepIndex = navigationSelection - 1;
          int newNote = sequence[stepIndex].note + delta;
          if (newNote < -1) newNote = -1;
          if (newNote > 127) newNote = 127;
          mutex_enter_blocking(&midiMutex);
          sequence[stepIndex].note = newNote;
          mutex_exit(&midiMutex);
        } else {
          // Move Cursor
          navigationSelection += (delta > 0 ? 1 : -1);
          if (navigationSelection < 0) navigationSelection = NUM_STEPS;
          if (navigationSelection > NUM_STEPS) navigationSelection = 0;
          
          // Adjust Scroll to keep cursor in view
          if (navigationSelection < scrollOffset) scrollOffset = navigationSelection;
          if (navigationSelection >= scrollOffset + 6) scrollOffset = navigationSelection - 5;
        }
        break;
      case UI_MENU_MAIN:
        menuSelection += (delta > 0 ? 1 : -1);
        if (menuSelection < 0) menuSelection = mainMenuCount - 1;
        if (menuSelection >= mainMenuCount) menuSelection = 0;
        break;
      case UI_MENU_RANDOMIZE:
        menuSelection += (delta > 0 ? 1 : -1);
        if (menuSelection < 0) menuSelection = randomizeMenuCount - 1;
        if (menuSelection >= randomizeMenuCount) menuSelection = 0;
        break;
      case UI_MENU_SETUP:
        menuSelection += (delta > 0 ? 1 : -1);
        if (menuSelection < 0) menuSelection = setupMenuCount - 1;
        if (menuSelection >= setupMenuCount) menuSelection = 0;
        break;
      case UI_SETUP_CHANNEL_EDIT:
        midiChannel += (delta > 0 ? 1 : -1);
        if (midiChannel < 1) midiChannel = 16;
        if (midiChannel > 16) midiChannel = 1;
        shMidiChannel = midiChannel;
        break;
      case UI_EDIT_TEMPO:
        tempo += delta;
        if (tempo < 40) tempo = 40;
        if (tempo > 240) tempo = 240;
        break;
      case UI_EDIT_FLAVOUR:
        currentStrategyIndex += (delta > 0 ? 1 : -1);
        if (currentStrategyIndex < 0) currentStrategyIndex = numStrategies - 1;
        if (currentStrategyIndex >= numStrategies) currentStrategyIndex = 0;
        break;
    }
  }

  // Handle Button
  int reading = digitalRead(ENC_SW);
  
  if (reading != lastButtonState) {
    lastDebounceTime = millis();
  }
  
  if ((millis() - lastDebounceTime) > 50) {
    if (reading == LOW && !buttonActive) {
       // Button Pressed
       buttonActive = true;
       buttonPressTime = millis();
       buttonConsumed = false;
       Serial.println(F("Button Down"));
    }
    
    if (reading == HIGH && buttonActive) {
       // Button Released
       buttonActive = false;
       if (!buttonConsumed) { // Short press action
          switch(currentState) {
            case UI_TRACKER:
              if (navigationSelection == 0) { // Menu item selected
                currentState = UI_MENU_MAIN;
                menuSelection = 0;
              } else { // A step is selected
                isEditing = !isEditing;
              }
              break;
            case UI_MENU_MAIN:
              if (menuSelection == 0) currentState = UI_TRACKER;
              if (menuSelection == 1) { currentState = UI_MENU_RANDOMIZE; menuSelection = 0; }
              if (menuSelection == 2) { currentState = UI_MENU_SETUP; menuSelection = 0; }
              break;
            case UI_MENU_RANDOMIZE:
              if (menuSelection == 0) { currentState = UI_MENU_MAIN; menuSelection = 1; }
              if (menuSelection == 1) {
                generateRandomScale(); // Scale
                if (isPlaying) {
                  int theme = (queuedTheme != -1) ? queuedTheme : currentThemeIndex;
                  mutex_enter_blocking(&midiMutex);
                  generateSequenceData(theme, nextSequence);
                  sequenceChangeScheduled = true;
                  mutex_exit(&midiMutex);
                }
                saveSequence(true);
              }
              if (menuSelection == 2) {
                melodySeed = random(10000); // Melody
                if (isPlaying) {
                  int theme = (queuedTheme != -1) ? queuedTheme : currentThemeIndex;
                  mutex_enter_blocking(&midiMutex);
                  generateSequenceData(theme, nextSequence);
                  sequenceChangeScheduled = true;
                  mutex_exit(&midiMutex);
                }
                saveSequence(true);
              }
              if (menuSelection == 3) currentState = UI_EDIT_FLAVOUR;
              if (menuSelection == 4) currentState = UI_EDIT_TEMPO;
              if (menuSelection == 5) mutationEnabled = !mutationEnabled;
              if (menuSelection == 6) {
                songModeEnabled = !songModeEnabled;
                if (songModeEnabled) {
                  songModeNeedsNext = true;
                }
              }
              if (menuSelection >= THEME_1_INDEX) { // Themes
                const int selectedTheme = menuSelection - THEME_1_INDEX + 1;
                if (isPlaying) {
                  queuedTheme = selectedTheme;
                  mutex_enter_blocking(&midiMutex);
                  generateSequenceData(queuedTheme, nextSequence);
                  sequenceChangeScheduled = true;
                  mutex_exit(&midiMutex);
                } else {
                  generateTheme(selectedTheme);
                }
              }
              break;
            case UI_MENU_SETUP:
              if (menuSelection == 0) { currentState = UI_MENU_MAIN; menuSelection = 2; }
              if (menuSelection == 1) currentState = UI_SETUP_CHANNEL_EDIT;
              if (menuSelection == 2) { 
                saveSequence(); 
                currentState = UI_MENU_MAIN; 
              }
              if (menuSelection == 3) { 
                if (loadSequence()) showMessage("LOADED!"); 
                currentState = UI_MENU_MAIN; 
              }
              break;
            case UI_SETUP_CHANNEL_EDIT:
              currentState = UI_MENU_SETUP;
              saveSequence(true);
              break;
            case UI_EDIT_TEMPO:
              currentState = UI_MENU_RANDOMIZE;
              saveSequence(true);
              break;
            case UI_EDIT_FLAVOUR:
              currentState = UI_MENU_RANDOMIZE;
              saveSequence(true);
              break;
          }
       }
    }
  }
  
  // Check for Long Press (Start/Stop Playback)
  if (buttonActive && !buttonConsumed && (millis() - buttonPressTime > 600)) {
    // Long press only works from tracker view
    if (currentState == UI_TRACKER) {
      isPlaying = !isPlaying;
      buttonConsumed = true; // Prevent short press action
      Serial.print(F("Playback: ")); Serial.println(isPlaying ? F("ON") : F("OFF"));
      if (isPlaying) {
        playbackStep = 0;
        clockCount = 0;
        lastClockTime = micros();
        sendMidiRealtime(0xFA); // MIDI Start
      } else {
        // Send All Notes Off on stop (CC 123)
        needsPanic = true;
        sendMidiRealtime(0xFC); // MIDI Stop
        queuedTheme = -1;
      }
    }
  }
  
  lastButtonState = reading;
}

void handlePlayback() {
  if (!isPlaying) return;

  unsigned long currentMicros = micros();
  unsigned long clockInterval = 2500000 / tempo; // 60s * 1000000us / (tempo * 24ppqn)

  if (currentMicros - lastClockTime >= clockInterval) {
    lastClockTime += clockInterval;
    
    sendMidiRealtime(0xF8); // MIDI Clock

    clockCount++;
    if (clockCount < 6) return; // 24 ppqn / 4 = 6 pulses per 16th note
    clockCount = 0;

    mutex_enter_blocking(&midiMutex);

    // Determine if we are tying to the next note
    int nextStep = playbackStep + 1;
    if (nextStep >= NUM_STEPS) nextStep = 0;
    
    bool isTied = sequence[playbackStep].tie && (sequence[nextStep].note != -1);
    int prevNote = sequence[playbackStep].note;

    // Note Off for previous step (if NOT tied)
    if (!isTied && prevNote != -1) {
      sendMidi(0x80, prevNote, 0);
    }
    
    playbackStep++;
    if (playbackStep >= NUM_STEPS) {
      playbackStep = 0;

      // Theme change
      if (sequenceChangeScheduled && queuedTheme != -1) {
        currentThemeIndex = queuedTheme;
        queuedTheme = -1;
        // nextSequence is already generated
      }

      // Mutation
      if (mutationEnabled) {
        if (!sequenceChangeScheduled) {
          memcpy(nextSequence, sequence, sizeof(sequence));
        }
        mutateSequence(nextSequence);
        sequenceChangeScheduled = true;
      }

      sendMidi(0xB0, 123, 0); // Panic / All Notes Off
      
      if (sequenceChangeScheduled) {
        memcpy(sequence, nextSequence, sizeof(sequence));
        sequenceChangeScheduled = false;
      }

      // Song Mode? Advance repeats
      if (songModeEnabled) {
        // we just used one repeat
        if (songRepeatsRemaining <= 1) {
          // let's start another round
          songRepeatsRemaining = nextSongRepeats;
        } else {
          // next repeat
          songRepeatsRemaining--;
        }
        // Trigger next song segment generation if we are on the last repeat
        if (songRepeatsRemaining <= 1 && !sequenceChangeScheduled && !songModeNeedsNext) {
          songModeNeedsNext = true;
        }
      }
    }
    
    // Note On for new step
    if (sequence[playbackStep].note != -1) {
      uint8_t velocity = sequence[playbackStep].accent ? 127 : 100;
      sendMidi(0x90, sequence[playbackStep].note, velocity);
    }

    // Note Off for previous step (if tied - delayed Note Off)
    if (isTied && prevNote != -1) {
      sendMidi(0x80, prevNote, 0);
    }

    mutex_exit(&midiMutex);
  }
}

void drawMenu(const char* title, const char* items[], int count, int selection) {
  display.println(title);
  display.drawLine(0, 8, 128, 8, SSD1306_WHITE);

  // Simple scrolling: keep selection visible
  int start = 0;
  if (selection >= 5) {
    start = selection - 4;
  }

  int y = 10;
  for (int i = start; i < count; i++) {
    if (y > 55) break; // Stop if we run out of screen space

    if (i == selection) {
      display.fillRect(0, y, 128, 8, SSD1306_WHITE);
      display.setTextColor(SSD1306_BLACK, SSD1306_WHITE);
    } else {
      display.setTextColor(SSD1306_WHITE);
    }
    display.setCursor(2, y);
    display.print(items[i]);

    // Special case for channel display
    if (currentState == UI_MENU_SETUP && i == 1) {
        display.print(F(": "));
        display.print(midiChannel);
    }

    // Special case for queued theme
    if (currentState == UI_MENU_RANDOMIZE && i >= THEME_1_INDEX && queuedTheme == (i - THEME_1_INDEX + 1)) {
        display.print(F(" [NEXT]"));
    }

    // Special case for active theme
    if (currentState == UI_MENU_RANDOMIZE && i >= THEME_1_INDEX && currentThemeIndex == (i - THEME_1_INDEX + 1)) {
        display.print(F(" *"));
    }

    // Special cases for Randomize Menu values
    if (currentState == UI_MENU_RANDOMIZE) {
      if (i == 1) { // Scale
        display.print(F(": "));
        if (numScaleNotes > 0) {
          const char* noteNames[] = {"C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B"};
          for (int j = 0; j < min(numScaleNotes, 6); j++) {
            display.print(noteNames[scaleNotes[j]]);
            if (j < min(numScaleNotes, 6) - 1) display.print(F(" "));
          }
        }
      } else if (i == 2) { // Melody
        display.print(F(": "));
        display.print(melodySeed);
      } else if (i == 3) { // Flavour
        display.print(F(": "));
        display.print(strategies[currentStrategyIndex]->getName());
      } else if (i == 4) { // Tempo
        display.print(F(": "));
        display.print(tempo);
      } else if (i == 5) { // Mutation
        display.print(F(": "));
        display.print(mutationEnabled ? F("ON") : F("OFF"));
      } else if (i == 6) { // Song Mode
        display.print(F(": "));
        display.print(songModeEnabled ? F("ON") : F("OFF"));
      }
    }
    y += 9;
  }
}

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

  // Header
  display.print(F("SEQ "));
  if (navigationSelection > 0 && isEditing) {
    display.print(F("[EDIT]"));
  } else {
    display.print(F("[NAV] "));
  }
  display.print(F(" CH:"));
  display.print(midiChannel);

  display.println();
  display.drawLine(0, 8, 128, 8, SSD1306_WHITE);

  // Steps
  int y = 10;
  mutex_enter_blocking(&midiMutex);
  for (int i = 0; i < 6; i++) {
    int itemIndex = i + scrollOffset;
    if (itemIndex > NUM_STEPS) break;

    // Draw Cursor
    if (itemIndex == navigationSelection) {
      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);

    if (itemIndex == 0) {
      display.print(F(">> MENU"));
    } else {
      int stepIndex = itemIndex - 1;
      bool isPlayback = isPlaying && (stepIndex == playbackStep);

      if (isPlayback) {
        if (itemIndex == navigationSelection) display.setTextColor(SSD1306_WHITE, SSD1306_BLACK);
        else display.setTextColor(SSD1306_BLACK, SSD1306_WHITE);
      }

      // Step Number
      if (stepIndex < 10) display.print(F("0"));
      display.print(stepIndex);

      if (isPlayback) {
        if (itemIndex == navigationSelection) display.setTextColor(SSD1306_BLACK, SSD1306_WHITE);
        else display.setTextColor(SSD1306_WHITE);
      }

      display.print(F(" | "));

      // Note Value
      int n = sequence[stepIndex].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;
  }
  mutex_exit(&midiMutex);
}

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

  switch(currentState) {
    case UI_TRACKER:
      drawTracker();
      break;
    case UI_MENU_MAIN:
      drawMenu("MAIN MENU", mainMenu, mainMenuCount, menuSelection);
      break;
    case UI_MENU_RANDOMIZE:
      drawMenu("RANDOMIZE", randomizeMenu, randomizeMenuCount, menuSelection);
      break;
    case UI_MENU_SETUP:
      drawMenu("SETUP", setupMenu, setupMenuCount, menuSelection);
      break;
    case UI_SETUP_CHANNEL_EDIT:
      display.println(F("SET MIDI CHANNEL"));
      display.drawLine(0, 8, 128, 8, SSD1306_WHITE);
      display.setCursor(20, 25);
      display.setTextSize(2);
      display.print(F("CH: "));
      if (midiChannel < 10) display.print(F(" "));
      display.print(midiChannel);
      display.setTextSize(1);
      display.setCursor(0, 50);
      display.println(F(" (Press to confirm)"));
      break;
    case UI_EDIT_TEMPO:
      display.println(F("SET TEMPO"));
      display.drawLine(0, 8, 128, 8, SSD1306_WHITE);
      display.setCursor(20, 25);
      display.setTextSize(2);
      display.print(F("BPM: "));
      display.print(tempo);
      display.setTextSize(1);
      display.setCursor(0, 50);
      display.println(F(" (Press to confirm)"));
      break;
    case UI_EDIT_FLAVOUR:
      display.println(F("SET FLAVOUR"));
      display.drawLine(0, 8, 128, 8, SSD1306_WHITE);
      display.setCursor(20, 25);
      display.setTextSize(2);
      display.print(strategies[currentStrategyIndex]->getName());
      display.setTextSize(1);
      display.setCursor(0, 50);
      display.println(F(" (Press to confirm)"));
      break;
  }

  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;
}

uint32_t getNoteColor(int note, bool dim) {
  if (note == -1) return 0;
  // Map note to hue, avoiding Green (approx 21845) which is used for playback cursor.
  // We start from Cyan (~30000) and go up to Orange (~8000 wrapped).
  // Range: 30000 to 65536+8000 = 73536. Width = 43536.
  // Step per semitone: 43536 / 12 = 3628.
  // This ensures notes are distinct colors but never pure Green.
  uint16_t hue = 30000 + (note % 12) * 3628;
  return Adafruit_NeoPixel::ColorHSV(hue, 255, dim ? 10 : 50);
}

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

  mutex_enter_blocking(&midiMutex);
  for (int s = 0; s < NUM_STEPS; s++) {
    int x = s % 8;
    int yBase = (s / 8) * 4;

    uint32_t color = 0;
    uint32_t dimColor = 0;
    if (sequence[s].note != -1) {
      color = getNoteColor(sequence[s].note, sequence[s].tie);
      dimColor = getNoteColor(sequence[s].note, true);
    }

    uint32_t colorP0 = 0;
    uint32_t colorP1 = 0;
    uint32_t colorP2 = 0;
    uint32_t colorP3 = 0;

    if (sequence[s].note != -1) {
      int octave = sequence[s].note / 12;
      
      if (octave > 4) {
        // Octave Up -> Top Row (P0)
        colorP0 = color;
        if (sequence[s].accent) colorP1 = dimColor;
      } else if (octave < 4) {
        // Octave Down -> Bottom Row (P2)
        colorP2 = color;
        if (sequence[s].accent) colorP1 = dimColor;
      } else {
        // Normal -> Middle Row (P1)
        colorP1 = color;
        if (sequence[s].accent) {
          colorP0 = dimColor;
          colorP2 = dimColor;
        }
      }
    }

    int stepNavIndex = navigationSelection - 1;
    uint32_t cursorColor = pixels.Color(0 + (s%2) * 20, 0 + (s%2) * 20, 50 + (s%2) * 20); // Blue for paused
    if (isPlaying) {
      cursorColor = pixels.Color(0, 50, 0); // Green for playback
      if (songModeEnabled) {
        // Song Mode: Show repeats on bottom row
        // Right aligned.
        // If repeats = 1 (last one), blink yellow.
        if (s >= NUM_STEPS/2) { // second half = bottom row
          int col = x;
          int repeatsToDraw = min(songRepeatsRemaining, 8);
          if (col >= (8 - repeatsToDraw)) {
             if (songRepeatsRemaining == 1 && col == 7 && (millis() / 250) % 2) {
               cursorColor = pixels.Color(0, 250, 0); // Max green, change is about to happen
             } else {
               cursorColor = pixels.Color(80, 100, 0); // Yellow-green, remaining repeat
             }
          }
        }
      }
    } else if (currentState == UI_TRACKER) {
      cursorColor = isEditing ? pixels.Color(50, 0, 0) : pixels.Color(40, 40, 40);
    }

    if (cursorColor != 0) {
      bool isCursorHere = (isPlaying && s == playbackStep) || 
                          (!isPlaying && currentState == UI_TRACKER && s == stepNavIndex);
      if (isCursorHere) {
        colorP3 = cursorColor;
      } else {
        // Lightly colored background for cursor row
        uint8_t r = (uint8_t)(cursorColor >> 16);
        uint8_t g = (uint8_t)(cursorColor >> 8);
        uint8_t b = (uint8_t)cursorColor;
        colorP3 = pixels.Color(r/5, g/5, b/5);
      }
    }

    pixels.setPixelColor(getPixelIndex(x, yBase), colorP0);
    pixels.setPixelColor(getPixelIndex(x, yBase + 1), colorP1);
    pixels.setPixelColor(getPixelIndex(x, yBase + 2), colorP2);
    pixels.setPixelColor(getPixelIndex(x, yBase + 3), colorP3);
  }

  if (sequenceChangeScheduled && (millis() / 125) % 2) {
    pixels.setPixelColor(NUM_PIXELS - 1, pixels.Color(127, 50, 0));
  }
  mutex_exit(&midiMutex);

  pixels.show();
}

void loop1() {
  if (needsPanic) {
    mutex_enter_blocking(&midiMutex);
    sendMidi(0xB0, 123, 0);
    mutex_exit(&midiMutex);
    needsPanic = false;
  }
  handlePlayback();
}

void loop() {
  // Handle Song Mode Generation in UI Thread
  if (songModeNeedsNext) {
    int nextTheme = random(1, 8); // Themes 1-7
    int repeats = random(1, 9);   // 1-8 repeats
    
    mutex_enter_blocking(&midiMutex);
    generateSequenceData(nextTheme, nextSequence);
    queuedTheme = nextTheme;
    nextSongRepeats = repeats;
    sequenceChangeScheduled = true;
    mutex_exit(&midiMutex);
    
    songModeNeedsNext = false;
  }

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