#include <SPI.h>
#include <Wire.h>
#include "TrackerTypes.h"
#include "MidiDriver.h"
#include "UIManager.h"
#include "config.h"
#include "UIThread.h"
#include "SharedState.h"
#include "SequenceGenerator.h"
#include "Persistence.h"

static Step local_sequence[NUM_TRACKS][NUM_STEPS];

static void handleInput();
static void drawUI();
static void updateLeds();

void saveSequence(bool quiet) {
    Persistence::saveSequence(quiet);
}

bool loadSequence() {
    return Persistence::loadSequence();
}

void factoryReset() {
    Persistence::factoryReset();
}

void generateRandomScale() {
    SequenceGenerator::generateRandomScale();
}

void generateTheme(int themeType) {
  SequenceGenerator::pickRandomScaleType(themeType);
  SequenceGenerator::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]) {
    SequenceGenerator::mutateSequence(target);
}

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[randomizeTrack] += delta;
        if (numSteps[randomizeTrack] < 1) numSteps[randomizeTrack] = 1;
        if (numSteps[randomizeTrack] > NUM_STEPS) numSteps[randomizeTrack] = NUM_STEPS;
        break;
      case UI_EDIT_ROOT:
        globalRoot += delta;
        if (globalRoot < 0) globalRoot = 11;
        if (globalRoot > 11) globalRoot = 0;
        currentRoot = globalRoot; // Reset current to global
        SequenceGenerator::updateScale();
        if (isPlaying) {
             sequenceChangeScheduled = true;
             SequenceGenerator::generateSequenceData((queuedTheme != -1) ? queuedTheme : currentThemeIndex, nextSequence);
        }
        break;
      case UI_EDIT_PROG_ORDER:
        if (delta > 0) progressionOrder = PROG_ORDER_RANDOM;
        else progressionOrder = PROG_ORDER_SEQUENCE;
        break;
      case UI_EDIT_PROG_TEMPLATE:
        templateSelectionIndex += delta;
        if (templateSelectionIndex < 0) templateSelectionIndex = numProgressions - 1;
        if (templateSelectionIndex >= numProgressions) templateSelectionIndex = 0;
        break;
      case UI_EDIT_PROG_LENGTH:
        currentProgression.length += delta;
        if (currentProgression.length < 1) currentProgression.length = 1;
        if (currentProgression.length > MAX_PROG_STEPS) currentProgression.length = MAX_PROG_STEPS;
        progressionStep = 0; // Reset step
        break;
      case UI_EDIT_PROG_STEP_ROOT:
        {
            int stepIdx = stepEditIndex / 3;
            currentProgression.steps[stepIdx].rootOffset += delta;
            if (currentProgression.steps[stepIdx].rootOffset < -12) currentProgression.steps[stepIdx].rootOffset = -12;
            if (currentProgression.steps[stepIdx].rootOffset > 12) currentProgression.steps[stepIdx].rootOffset = 12;
        }
        break;
      case UI_EDIT_PROG_STEP_TYPE:
        {
            int stepIdx = stepEditIndex / 3;
            int type = currentProgression.steps[stepIdx].scaleType;
            type += delta;
            if (type < 0) type = 9;
            if (type > 9) type = 0;
            currentProgression.steps[stepIdx].scaleType = type;
        }
        break;
      case UI_EDIT_PROG_STEP_REPS:
        {
            int stepIdx = stepEditIndex / 3;
            int val = currentProgression.steps[stepIdx].repeats + delta;
            if (val < 1) val = 1;
            if (val > 16) val = 16;
            currentProgression.steps[stepIdx].repeats = val;
        }
        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;
    }
    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);
                randomSeed(melodySeeds[randomizeTrack]);
                for (int i = 0; i < NUM_TRACKS; i++) {
                  numSteps[i] = random(1, NUM_STEPS + 1);
                }
                if (isPlaying) {
                  int theme = (queuedTheme != -1) ? queuedTheme : currentThemeIndex;
                  if (!sequenceChangeScheduled) {
                      memcpy(nextSequence, sequence, sizeof(sequence));
                  }
                  SequenceGenerator::generateTrackData(randomizeTrack, theme, nextSequence);
                  sequenceChangeScheduled = true;
                }
                midi.unlock();
                saveSequence(true);
                break; 
              
              case MENU_ID_ROOT: currentState = UI_EDIT_ROOT; break;
              case MENU_ID_PROG_ORDER: currentState = UI_EDIT_PROG_ORDER; break;
              case MENU_ID_PROG_APPLY_TEMPLATE: currentState = UI_EDIT_PROG_TEMPLATE; templateSelectionIndex = 0; break;
              case MENU_ID_PROG_LENGTH: currentState = UI_EDIT_PROG_LENGTH; break;
              
              case MENU_ID_PROG_STEP_START ... (MENU_ID_PROG_STEP_END - 1):
                stepEditIndex = menuItems[menuSelection].id - MENU_ID_PROG_STEP_START;
                if (stepEditIndex % 3 == 0) currentState = UI_EDIT_PROG_STEP_ROOT;
                else if (stepEditIndex % 3 == 1) currentState = UI_EDIT_PROG_STEP_TYPE;
                else currentState = UI_EDIT_PROG_STEP_REPS;
                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;
                  SequenceGenerator::pickRandomScaleType(queuedTheme);
                  midi.lock();
                  SequenceGenerator::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) Persistence::savePatch(bank, slot);
                  else Persistence::loadPatch(bank, slot, local_sequence);
                  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_ROOT:
              currentState = UI_MENU_MAIN;
              saveSequence(true);
              break;
            case UI_EDIT_PROG_ORDER:
              currentState = UI_MENU_MAIN;
              saveSequence(true);
              break;
            case UI_EDIT_PROG_TEMPLATE:
              // Apply template
              if (templateSelectionIndex > 0) {
                  currentProgression = progressions[templateSelectionIndex];
                  progressionStep = 0;
              } else {
                  // Off/Clear
                  currentProgression.length = 0;
              }
              currentState = UI_MENU_MAIN;
              saveSequence(true);
              break;
            case UI_EDIT_PROG_LENGTH:
            case UI_EDIT_PROG_STEP_ROOT:
            case UI_EDIT_PROG_STEP_TYPE:
            case UI_EDIT_PROG_STEP_REPS:
              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));
                  }
                  SequenceGenerator::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));
                  }
                  SequenceGenerator::generateTrackData(randomizeTrack, theme, nextSequence);
                  sequenceChangeScheduled = true;
                  midi.unlock();
              }
              saveSequence(true);
              break;
            case UI_RANDOMIZE_TRACK_EDIT:
              currentState = UI_MENU_MAIN;
              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_currentTrackNumSteps;
  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;
  local_menuSelection = menuSelection;
  local_midiChannel = midiChannels[local_randomizeTrack];
  local_tempo = tempo;
  local_currentTrackNumSteps = numSteps[local_randomizeTrack];
  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_currentTrackNumSteps,
    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;
  int local_scaleNotes[12], local_numSteps[NUM_TRACKS];
  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;
  for (int i = 0; i < NUM_TRACKS; i++) local_numSteps[i] = numSteps[i];
  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_STEPS || local_currentState == UI_EDIT_FLAVOUR || local_currentState == UI_RANDOMIZE_TRACK_EDIT) {
      // 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
    
    if (currentProgression.length > 0) {
        // Progression Mode
        const ChordProgression& prog = currentProgression;
        
        currentRoot = (globalRoot + prog.steps[progressionStep].rootOffset) % 12;
        if (currentRoot < 0) currentRoot += 12;
        currentScaleType = prog.steps[progressionStep].scaleType;
        SequenceGenerator::updateScale();
        
        queuedProgressionStep = progressionStep;
        repeats = prog.steps[progressionStep].repeats;
        if (progressionOrder == PROG_ORDER_SEQUENCE)
            progressionStep = (progressionStep + 1) % prog.length;
        else
            progressionStep = random(prog.length);
    } else {
        // Random Mode
        SequenceGenerator::pickRandomScaleType(nextTheme);
    }

    SequenceGenerator::generateSequenceData(nextTheme, nextSequence);
    queuedTheme = nextTheme;
    nextSongRepeats = repeats;
    sequenceChangeScheduled = true;
    songModeNeedsNext = false;
  }

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