#include <SPI.h>
#include <Wire.h>
#include <EEPROM.h>
#include "TrackerTypes.h"
#include "MelodyStrategy.h"
#include "LuckyStrategy.h"
#include "ArpStrategy.h"
#include "EuclideanStrategy.h"
#include "MarkovStrategy.h"
#include "CellularAutomataStrategy.h"
#include "LSystemStrategy.h"
#include "DroneStrategy.h"
#include "MidiDriver.h"
#include "UIManager.h"
#include "config.h"
#include "UIThread.h"
#include "SharedState.h"

extern volatile int trackIntensity[NUM_TRACKS];

static Step local_sequence[NUM_TRACKS][NUM_STEPS];

static void handleInput();
static int scaleEditSelection = 0;
static int scaleEditNoteIndex = 0;
static void drawUI();
static void updateLeds();
static void generateTrackData(int track, int themeType, Step (*target)[NUM_STEPS]);
static void generateSequenceData(int themeType, Step (*target)[NUM_STEPS]);
static void savePatch(int bank, int slot);
static void loadPatch(int bank, int slot);

void saveSequence(bool quiet) {
  midi.lock();
  int addr = 0;
  EEPROM.put(addr, EEPROM_MAGIC); addr += sizeof(EEPROM_MAGIC);
  int channels[NUM_TRACKS];
  for(int i=0; i<NUM_TRACKS; i++) channels[i] = midiChannels[i];
  EEPROM.put(addr, channels); addr += sizeof(channels);
  EEPROM.put(addr, melodySeeds); addr += sizeof(melodySeeds);
  EEPROM.put(addr, currentStrategyIndices); addr += sizeof(currentStrategyIndices);
  bool mutes[NUM_TRACKS];
  for(int i=0; i<NUM_TRACKS; i++) mutes[i] = trackMute[i];
  EEPROM.put(addr, mutes); addr += sizeof(mutes);
  EEPROM.put(addr, (int)tempo); addr += sizeof(int);
  int intensities[NUM_TRACKS];
  for(int i=0; i<NUM_TRACKS; i++) intensities[i] = trackIntensity[i];
  EEPROM.put(addr, intensities); addr += sizeof(intensities);
  EEPROM.put(addr, (int)numSteps); 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);
  }

  EEPROM.put(addr, sequence); addr += sizeof(sequence);
  midi.unlock();
  EEPROM.commit();
  if (!quiet) ui.showMessage("SAVED!");
}

bool loadSequence() {
  midi.lock();
  int addr = 0;
  uint32_t magic;
  EEPROM.get(addr, magic); addr += sizeof(magic);
  if (magic != EEPROM_MAGIC) {
    midi.unlock();
    return false;
  }
  
  int channels[NUM_TRACKS];
  EEPROM.get(addr, channels); addr += sizeof(channels);
  for(int i=0; i<NUM_TRACKS; i++) {
    midiChannels[i] = channels[i];
    if (midiChannels[i] < 1) midiChannels[i] = 1;
    if (midiChannels[i] > 16) midiChannels[i] = 16;
  }
  EEPROM.get(addr, melodySeeds); addr += sizeof(melodySeeds);
  EEPROM.get(addr, currentStrategyIndices); addr += sizeof(currentStrategyIndices);
  for(int i=0; i<NUM_TRACKS; i++) {
    if (currentStrategyIndices[i] < 0 || currentStrategyIndices[i] >= numStrategies) currentStrategyIndices[i] = 0;
  }
  bool mutes[NUM_TRACKS];
  EEPROM.get(addr, mutes); addr += sizeof(mutes);
  for(int i=0; i<NUM_TRACKS; i++) trackMute[i] = mutes[i];
  int t;
  EEPROM.get(addr, t); addr += sizeof(int);
  tempo = t;
  if (tempo < 40) tempo = 40;
  if (tempo > 240) tempo = 240;
  int intensities[NUM_TRACKS];
  EEPROM.get(addr, intensities); addr += sizeof(intensities);
  for(int i=0; i<NUM_TRACKS; i++) {
    trackIntensity[i] = intensities[i];
    if (trackIntensity[i] < 1) trackIntensity[i] = 1;
    if (trackIntensity[i] > 10) trackIntensity[i] = 10;
  }
  EEPROM.get(addr, t); addr += sizeof(int);
  numSteps = t;
  if (numSteps <= 0 || numSteps >= NUM_STEPS) numSteps = NUM_STEPS;

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

  EEPROM.get(addr, sequence); addr += sizeof(sequence);
  midi.unlock();
  return true;
}

static void savePatch(int bank, int slot) {
    int patchIndex = bank * 4 + slot;
    int addr = 512 + patchIndex * 256; // Start after main save, 256 bytes per patch
    
    midi.lock();
    EEPROM.put(addr, numScaleNotes); addr += sizeof(numScaleNotes);
    for (int i = 0; i < 12; i++) {
        EEPROM.put(addr, scaleNotes[i]); addr += sizeof(int);
    }
    EEPROM.put(addr, currentStrategyIndices); addr += sizeof(currentStrategyIndices);
    EEPROM.put(addr, melodySeeds); addr += sizeof(melodySeeds);
    EEPROM.put(addr, (int)numSteps); addr += sizeof(int);
    bool mutes[NUM_TRACKS];
    for(int i=0; i<NUM_TRACKS; i++) mutes[i] = trackMute[i];
    EEPROM.put(addr, mutes); addr += sizeof(mutes);
    int intensities[NUM_TRACKS];
    for(int i=0; i<NUM_TRACKS; i++) intensities[i] = trackIntensity[i];
    EEPROM.put(addr, intensities); addr += sizeof(intensities);
    midi.unlock();
    EEPROM.commit();
    ui.showMessage("SAVED!");
}

static void loadPatch(int bank, int slot) {
    int patchIndex = bank * 4 + slot;
    int addr = 512 + patchIndex * 256;

    midi.lock();
    EEPROM.get(addr, numScaleNotes); addr += sizeof(numScaleNotes);
    if (numScaleNotes < 0 || numScaleNotes > 12) numScaleNotes = 0;
    for (int i = 0; i < 12; i++) {
        EEPROM.get(addr, scaleNotes[i]); addr += sizeof(int);
        if (scaleNotes[i] < 0) scaleNotes[i] = 0;
        if (scaleNotes[i] > 11) scaleNotes[i] = 11;
    }
    EEPROM.get(addr, currentStrategyIndices); addr += sizeof(currentStrategyIndices);
    for(int i=0; i<NUM_TRACKS; i++) {
      if (currentStrategyIndices[i] < 0 || currentStrategyIndices[i] >= numStrategies) currentStrategyIndices[i] = 0;
    }
    EEPROM.get(addr, melodySeeds); addr += sizeof(melodySeeds);
    int t;
    EEPROM.get(addr, t); addr += sizeof(int);
    numSteps = t;
    if (numSteps <= 0 || numSteps >= NUM_STEPS) numSteps = NUM_STEPS;
    
    bool mutes[NUM_TRACKS];
    EEPROM.get(addr, mutes); addr += sizeof(mutes);
    for(int i=0; i<NUM_TRACKS; i++) trackMute[i] = mutes[i];

    int intensities[NUM_TRACKS];
    EEPROM.get(addr, intensities); addr += sizeof(intensities);
    for(int i=0; i<NUM_TRACKS; i++) {
      trackIntensity[i] = intensities[i];
      if (trackIntensity[i] < 1) trackIntensity[i] = 1;
      if (trackIntensity[i] > 10) trackIntensity[i] = 10;
    }
    
    if (isPlaying) {
        generateSequenceData(currentThemeIndex, nextSequence);
        sequenceChangeScheduled = true;
    } else {
        generateSequenceData(currentThemeIndex, local_sequence);
        memcpy(sequence, local_sequence, sizeof(local_sequence));
    }
    midi.unlock();
    ui.showMessage("LOADED!");
}

void factoryReset() {
  ui.showMessage("RESETTING...");
  for(int i=0; i<NUM_TRACKS; i++) {
    trackIntensity[i] = 10;
  }
  uint32_t magic = 0;
  EEPROM.put(0, magic);
  EEPROM.commit();
  delay(500);
  rp2040.reboot();
}

static void generateTrackData(int track, int themeType, Step (*target)[NUM_STEPS]) {
    randomSeed(melodySeeds[track] + themeType * 12345);
    strategies[currentStrategyIndices[track]]->generate(target, track, numSteps, scaleNotes, numScaleNotes, melodySeeds[track] + themeType * 12345, trackIntensity[track]);
}

void generateRandomScale() {
  Serial.println(F("Generating new scale."));
  // All tracks share the same scale for now
  strategies[currentStrategyIndices[0]]->generateScale(scaleNotes, numScaleNotes);
}

static void generateSequenceData(int themeType, Step (*target)[NUM_STEPS]) {
  Serial.println(F("Generating sequence."));
  for(int i=0; i<NUM_TRACKS; i++) {
    generateTrackData(i, themeType, target);
  }
}

void generateTheme(int themeType) {
  generateSequenceData(themeType, local_sequence);

  Serial.println(F("Generating theme."));
  midi.lock();
  memcpy(sequence, local_sequence, sizeof(local_sequence));
  needsPanic = true;
  midi.unlock();

  Serial.println(F("Theme ready."));
  currentThemeIndex = themeType;
  clockCount = 0;
  lastClockTime = micros();
  playbackStep = 0;
  isPlaying = true;
}

void mutateSequence(Step (*target)[NUM_STEPS]) {
  for(int i=0; i<NUM_TRACKS; i++) strategies[currentStrategyIndices[i]]->mutate(target, i, numSteps, scaleNotes, numScaleNotes, trackIntensity[i]);
}

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

  if (delta != 0) {
    switch(currentState) {
      case UI_MENU_MAIN:
        {
            int next = menuSelection;
            int count = 0;
            do {
                next += (delta > 0 ? 1 : -1);
                if (next < 0) next = menuItemsCount - 1;
                if (next >= menuItemsCount) next = 0;
                count++;
            } while (!isItemVisible(next) && count < menuItemsCount);
            menuSelection = next;
        }
        break;
      case UI_SETUP_CHANNEL_EDIT:
        {
          midiChannels[randomizeTrack] += (delta > 0 ? 1 : -1);
          if (midiChannels[randomizeTrack] < 1) midiChannels[randomizeTrack] = 16;
          if (midiChannels[randomizeTrack] > 16) midiChannels[randomizeTrack] = 1;
        }
        break;
      case UI_EDIT_TEMPO:
        tempo += delta;
        if (tempo < 40) tempo = 40;
        if (tempo > 240) tempo = 240;
        break;
      case UI_EDIT_STEPS:
        numSteps += delta;
        if (numSteps < 1) numSteps = 1;
        if (numSteps > NUM_STEPS) numSteps = NUM_STEPS;
        break;
      case UI_EDIT_FLAVOUR:
        {
          currentStrategyIndices[randomizeTrack] += (delta > 0 ? 1 : -1);
          if (currentStrategyIndices[randomizeTrack] < 0) currentStrategyIndices[randomizeTrack] = numStrategies - 1;
          if (currentStrategyIndices[randomizeTrack] >= numStrategies) currentStrategyIndices[randomizeTrack] = 0;
        }
        break;
      case UI_EDIT_INTENSITY:
        {
          int current = trackIntensity[randomizeTrack];
          current += delta;
          if (current < 1) current = 1;
          if (current > 10) current = 10;
          trackIntensity[randomizeTrack] = current;
        }
        break;
      case UI_SCALE_EDIT:
        scaleEditSelection += (delta > 0 ? 1 : -1);
        if (scaleEditSelection < 0) scaleEditSelection = numScaleNotes + 4;
        if (scaleEditSelection > numScaleNotes + 4) scaleEditSelection = 0;
        break;
      case UI_SCALE_NOTE_EDIT:
        scaleNotes[scaleEditNoteIndex] += (delta > 0 ? 1 : -1);
        if (scaleNotes[scaleEditNoteIndex] < 0) scaleNotes[scaleEditNoteIndex] = 11;
        if (scaleNotes[scaleEditNoteIndex] > 11) scaleNotes[scaleEditNoteIndex] = 0;
        midi.lock();
        midi.sendNoteOn(60 + scaleNotes[scaleEditNoteIndex], 100, midiChannels[randomizeTrack]);
        midi.unlock();
        delay(50);
        midi.lock();
        midi.sendNoteOff(60 + scaleNotes[scaleEditNoteIndex], midiChannels[randomizeTrack]);
        midi.unlock();
        break;
      case UI_SCALE_TRANSPOSE:
        if (delta != 0) {
            int shift = delta % 12;
            if (shift < 0) shift += 12;
            for(int i=0; i<numScaleNotes; i++) scaleNotes[i] = (scaleNotes[i] + shift) % 12;
            sortArray(scaleNotes, numScaleNotes);
            if (isPlaying) {
                int theme = (queuedTheme != -1) ? queuedTheme : currentThemeIndex;
                midi.lock();
                generateSequenceData(theme, nextSequence);
                sequenceChangeScheduled = true;
                midi.unlock();
            }
        }
        break;
    }
    if (currentState == UI_RANDOMIZE_TRACK_EDIT) {
        randomizeTrack += (delta > 0 ? 1 : -1);
        if (randomizeTrack < 0) randomizeTrack = NUM_TRACKS - 1;
        if (randomizeTrack >= NUM_TRACKS) randomizeTrack = 0;
    }

  }

  // 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_MENU_MAIN:
              if (menuItems[menuSelection].isGroup) {
                  menuItems[menuSelection].expanded = !menuItems[menuSelection].expanded;
                  break;
              }
              
              switch(menuItems[menuSelection].id) {
              case MENU_ID_PLAYBACK:
                isPlaying = !isPlaying;
                if (isPlaying) {
                  playbackStep = 0;
                  clockCount = 0;
                  lastClockTime = micros();
                } else {
                  queuedTheme = -1;
                }
                break;
              case MENU_ID_MELODY:
                midi.lock();
                melodySeeds[randomizeTrack] = random(10000);
                if (isPlaying) {
                  int theme = (queuedTheme != -1) ? queuedTheme : currentThemeIndex;
                  if (!sequenceChangeScheduled) {
                      memcpy(nextSequence, sequence, sizeof(sequence));
                  }
                  generateTrackData(randomizeTrack, theme, nextSequence);
                  sequenceChangeScheduled = true;
                }
                midi.unlock();
                saveSequence(true);
                break; 
              
              case MENU_ID_SCALE:
                currentState = UI_SCALE_EDIT;
                scaleEditSelection = 0;
                break; 
              
              case MENU_ID_TEMPO: currentState = UI_EDIT_TEMPO; break;
              case MENU_ID_STEPS: currentState = UI_EDIT_STEPS; break;
              
              case MENU_ID_SONG_MODE:
                songModeEnabled = !songModeEnabled;
                if (songModeEnabled) {
                  songModeNeedsNext = true;
                }
                break; 

              case MENU_ID_TRACK_SELECT: currentState = UI_RANDOMIZE_TRACK_EDIT; break;
              case MENU_ID_MUTE: trackMute[randomizeTrack] = !trackMute[randomizeTrack]; break;
              case MENU_ID_FLAVOUR: currentState = UI_EDIT_FLAVOUR;  break;
              case MENU_ID_INTENSITY: currentState = UI_EDIT_INTENSITY; break;
              case MENU_ID_MUTATION: mutationEnabled = !mutationEnabled; break;
              
              case MENU_ID_CHANNEL: currentState = UI_SETUP_CHANNEL_EDIT; break;
              case MENU_ID_PROTECTED_MODE: protectedMode = !protectedMode; break;
              case MENU_ID_RESET: factoryReset(); break;

              default:
                if (menuItems[menuSelection].id >= MENU_ID_THEME_1 && menuItems[menuSelection].id <= MENU_ID_THEME_7) {
                const int selectedTheme = menuItems[menuSelection].id - MENU_ID_THEME_1 + 1;
                if (isPlaying) {
                  queuedTheme = selectedTheme;
                  midi.lock();
                  generateSequenceData(queuedTheme, nextSequence);
                  sequenceChangeScheduled = true;
                  midi.unlock();
                } else {
                  generateTheme(selectedTheme);
                }
                break; 
              }
              if (menuItems[menuSelection].id >= MENU_ID_PATCH_ACTIONS_START) {
                  int offset = menuItems[menuSelection].id - MENU_ID_PATCH_ACTIONS_START;
                  int bank = offset / 8;
                  int sub = offset % 8;
                  bool isSave = sub >= 4;
                  int slot = sub % 4;
                  if (isSave) savePatch(bank, slot);
                  else loadPatch(bank, slot);
                  break;
              }
              break;
              }
              break;
            case UI_SETUP_CHANNEL_EDIT:
              currentState = UI_MENU_MAIN;
              saveSequence(true);
              break;
            case UI_EDIT_TEMPO:
              currentState = UI_MENU_MAIN;
              saveSequence(true);
              break;
            case UI_EDIT_STEPS:
              currentState = UI_MENU_MAIN;
              saveSequence(true);
              break;
            case UI_EDIT_FLAVOUR:
              currentState = UI_MENU_MAIN;
              if (isPlaying) {
                  int theme = (queuedTheme != -1) ? queuedTheme : currentThemeIndex;
                  midi.lock();
                  if (!sequenceChangeScheduled) {
                      memcpy(nextSequence, sequence, sizeof(sequence));
                  }
                  generateTrackData(randomizeTrack, theme, nextSequence);
                  sequenceChangeScheduled = true;
                  midi.unlock();
              }
              saveSequence(true);
              break;
            case UI_EDIT_INTENSITY:
              currentState = UI_MENU_MAIN;
              if (isPlaying) {
                  int theme = (queuedTheme != -1) ? queuedTheme : currentThemeIndex;
                  midi.lock();
                  if (!sequenceChangeScheduled) {
                      memcpy(nextSequence, sequence, sizeof(sequence));
                  }
                  generateTrackData(randomizeTrack, theme, nextSequence);
                  sequenceChangeScheduled = true;
                  midi.unlock();
              }
              saveSequence(true);
              break;
            case UI_RANDOMIZE_TRACK_EDIT:
              currentState = UI_MENU_MAIN;
              saveSequence(true);
              break;
            case UI_SCALE_EDIT:
              if (scaleEditSelection == 0) {
                  currentState = UI_MENU_MAIN;
                  saveSequence(true);
              } else if (scaleEditSelection == 1) {
                  generateRandomScale();
                  if (isPlaying) {
                      int theme = (queuedTheme != -1) ? queuedTheme : currentThemeIndex;
                      midi.lock();
                      generateSequenceData(theme, nextSequence);
                      sequenceChangeScheduled = true;
                      midi.unlock();
                  }
                  saveSequence(true);
              } else if (scaleEditSelection <= numScaleNotes + 1) {
                  scaleEditNoteIndex = scaleEditSelection - 2;
                  currentState = UI_SCALE_NOTE_EDIT;
              } else if (scaleEditSelection == numScaleNotes + 2) {
                  currentState = UI_SCALE_TRANSPOSE;
              } else if (scaleEditSelection == numScaleNotes + 3) {
                  if (numScaleNotes < 12) {
                      int next = (numScaleNotes > 0) ? (scaleNotes[numScaleNotes-1] + 1) % 12 : 0;
                      scaleNotes[numScaleNotes] = next;
                      numScaleNotes++;
                      scaleEditSelection--; // Move cursor to the new note
                      if (isPlaying) {
                          int theme = (queuedTheme != -1) ? queuedTheme : currentThemeIndex;
                          midi.lock();
                          generateSequenceData(theme, nextSequence);
                          sequenceChangeScheduled = true;
                          midi.unlock();
                      }
                      saveSequence(true);
                  }
              } else {
                  if (numScaleNotes > 1) {
                      numScaleNotes--;
                      scaleEditSelection--;
                      if (isPlaying) {
                          int theme = (queuedTheme != -1) ? queuedTheme : currentThemeIndex;
                          midi.lock();
                          generateSequenceData(theme, nextSequence);
                          sequenceChangeScheduled = true;
                          midi.unlock();
                      }
                      saveSequence(true);
                  }
              }
              break;
            case UI_SCALE_NOTE_EDIT:
              sortArray(scaleNotes, numScaleNotes);
              if (isPlaying) {
                  int theme = (queuedTheme != -1) ? queuedTheme : currentThemeIndex;
                  midi.lock();
                  generateSequenceData(theme, nextSequence);
                  sequenceChangeScheduled = true;
                  midi.unlock();
              }
              currentState = UI_SCALE_EDIT;
              saveSequence(true);
              break;
            case UI_SCALE_TRANSPOSE:
              currentState = UI_SCALE_EDIT;
              saveSequence(true);
              break;
          }
       }
    }
  }
  
  // Check for Long Press (Start/Stop Playback)
  if (buttonActive && !buttonConsumed && (millis() - buttonPressTime > 600)) {
    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();
    } else {
      queuedTheme = -1;
    }
  }
  
  lastButtonState = reading;
}

static void drawUI() {
  // Make local copies of shared data inside a critical section
  // to avoid holding the lock during slow display operations.
  UIState local_currentState;
  int local_menuSelection, local_randomizeTrack, local_tempo, local_currentThemeIndex, local_queuedTheme, local_numScaleNotes;
  int local_melodySeed, local_numSteps;
  bool local_mutationEnabled, local_songModeEnabled, local_isPlaying;
  bool local_trackMute[NUM_TRACKS];
  int local_midiChannel;
  int local_trackIntensities[NUM_TRACKS];
  MelodyStrategy* local_strategy;
  int local_playbackStep;
  int local_scaleNotes[12];

  midi.lock();
  local_randomizeTrack = randomizeTrack;

  local_currentState = currentState;
  if (local_currentState == UI_SCALE_EDIT || local_currentState == UI_SCALE_NOTE_EDIT || local_currentState == UI_SCALE_TRANSPOSE) {
      local_menuSelection = scaleEditSelection;
  } else {
      local_menuSelection = menuSelection;
  }
  local_midiChannel = midiChannels[local_randomizeTrack];
  local_tempo = tempo;
  local_numSteps = numSteps;
  local_strategy = strategies[currentStrategyIndices[local_randomizeTrack]];
  local_queuedTheme = queuedTheme;
  local_currentThemeIndex = currentThemeIndex;
  local_numScaleNotes = numScaleNotes;
  memcpy(local_scaleNotes, scaleNotes, sizeof(local_scaleNotes));
  local_melodySeed = melodySeeds[local_randomizeTrack];
  local_mutationEnabled = mutationEnabled;
  local_songModeEnabled = songModeEnabled;
  memcpy(local_sequence, sequence, sizeof(local_sequence));
  local_playbackStep = playbackStep;
  local_isPlaying = isPlaying;
  memcpy(local_trackMute, (const void*)trackMute, sizeof(local_trackMute));
  memcpy(local_trackIntensities, (const void*)trackIntensity, sizeof(local_trackIntensities));
  midi.unlock();

  ui.draw(local_currentState, local_menuSelection, 
    local_midiChannel, local_tempo, local_strategy, 
    local_queuedTheme, local_currentThemeIndex, local_numScaleNotes, local_scaleNotes, local_melodySeed, local_numSteps,
    local_mutationEnabled, local_songModeEnabled, (const Step (*)[NUM_STEPS])local_sequence, local_playbackStep, local_isPlaying, local_randomizeTrack, (const bool*)local_trackMute, (const int*)local_trackIntensities);
}

static void updateLeds() {
  // Make local copies of shared data inside a critical section
  // to avoid holding the lock during slow LED update operations.
  int local_playbackStep;
  bool local_isPlaying;
  UIState local_currentState;
  int local_menuSelection;
  bool local_songModeEnabled;
  int local_songRepeatsRemaining;
  bool local_sequenceChangeScheduled;
  PlayMode local_playMode;
  int local_numScaleNotes, local_numSteps;
  int local_scaleNotes[12];
  bool local_trackMute[NUM_TRACKS];
  int local_randomizeTrack;

  midi.lock();
  memcpy(local_sequence, sequence, sizeof(local_sequence));
  local_playbackStep = playbackStep;
  local_isPlaying = isPlaying;
  local_currentState = currentState;
  local_menuSelection = menuSelection;
  local_songModeEnabled = songModeEnabled;
  local_songRepeatsRemaining = songRepeatsRemaining;
  local_sequenceChangeScheduled = sequenceChangeScheduled;
  local_playMode = playMode;
  local_numScaleNotes = numScaleNotes;
  local_numSteps = numSteps;
  local_randomizeTrack = randomizeTrack;
  memcpy(local_scaleNotes, scaleNotes, sizeof(local_scaleNotes));
  memcpy(local_trackMute, (const void*)trackMute, sizeof(local_trackMute));
  midi.unlock();

  PlayMode ledDisplayMode = MODE_POLY; // Default to POLY (MAIN section view)

  if (local_currentState == UI_MENU_MAIN) {
      MenuItemID id = menuItems[local_menuSelection].id;
      // Check if we are in the Track group (IDs between TRACK_SELECT and THEME_7)
      if (id >= MENU_ID_TRACK_SELECT && id <= MENU_ID_THEME_7) {
          // It's a TRACK section item (Track, Mute, Flavour, Mutation, Themes)
          ledDisplayMode = MODE_MONO;
      }
  } else if (local_currentState == UI_EDIT_FLAVOUR || local_currentState == UI_RANDOMIZE_TRACK_EDIT || local_currentState == UI_SCALE_EDIT || local_currentState == UI_SCALE_NOTE_EDIT || local_currentState == UI_SCALE_TRANSPOSE) {
      // These are entered from TRACK section items
      ledDisplayMode = MODE_MONO;
  }

  ui.updateLeds((const Step (*)[NUM_STEPS])local_sequence, local_playbackStep, local_isPlaying, 
                local_currentState, local_songModeEnabled, local_songRepeatsRemaining, 
                local_sequenceChangeScheduled, ledDisplayMode, local_randomizeTrack, local_numSteps, local_numScaleNotes, 
                local_scaleNotes, (const bool*)local_trackMute);
}

void loopUI() {
  // 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
    
    generateSequenceData(nextTheme, nextSequence);
    queuedTheme = nextTheme;
    nextSongRepeats = repeats;
    sequenceChangeScheduled = true;
    songModeNeedsNext = false;
  }

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