#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 "MidiDriver.h"
#include "UIManager.h"
#include "config.h"
#include "UIThread.h"
#include "SharedState.h"

static Step local_sequence[NUM_TRACKS][NUM_STEPS];

static void handleInput();
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]);

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

  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) 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];
  EEPROM.get(addr, melodySeeds); addr += sizeof(melodySeeds);
  EEPROM.get(addr, currentStrategyIndices); addr += sizeof(currentStrategyIndices);
  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;

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

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

void factoryReset() {
  ui.showMessage("RESETTING...");
  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, NUM_STEPS, scaleNotes, numScaleNotes, melodySeeds[track] + themeType * 12345);
}

void generateRandomScale() {
  // All tracks share the same scale for now
  strategies[currentStrategyIndices[0]]->generateScale(scaleNotes, numScaleNotes);
}

static void generateSequenceData(int themeType, Step (*target)[NUM_STEPS]) {
  for(int i=0; i<NUM_TRACKS; i++) {
    generateTrackData(i, themeType, target);
  }
}

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

  midi.lock();
  memcpy(sequence, local_sequence, sizeof(local_sequence));
  needsPanic = true;
  midi.unlock();

  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, NUM_STEPS, scaleNotes, numScaleNotes);
}

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

  if (delta != 0) {
    switch(currentState) {
      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 = randomizeMenuPolyCount - 1;
          if (menuSelection >= randomizeMenuPolyCount) 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:
        {
          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_FLAVOUR:
        {
          currentStrategyIndices[randomizeTrack] += (delta > 0 ? 1 : -1);
          if (currentStrategyIndices[randomizeTrack] < 0) currentStrategyIndices[randomizeTrack] = numStrategies - 1;
          if (currentStrategyIndices[randomizeTrack] >= numStrategies) currentStrategyIndices[randomizeTrack] = 0;
        }
        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 (menuSelection == 0) { currentState = UI_MENU_RANDOMIZE; menuSelection = 0; break;  }
              if (menuSelection == 1) { currentState = UI_MENU_SETUP; menuSelection = 0; break;  }
              break;
            case UI_MENU_RANDOMIZE:
              {
              if (menuSelection == 0) { currentState = UI_MENU_SETUP; menuSelection = 0; break;  }

              if (menuSelection == 1) { // 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; 
              }
              if (menuSelection == 2) { // Scale
                generateRandomScale();
                if (isPlaying) {
                  int theme = (queuedTheme != -1) ? queuedTheme : currentThemeIndex;
                  midi.lock();
                  // Regenerate all tracks with new scale
                  generateSequenceData(theme, nextSequence);
                  sequenceChangeScheduled = true;
                  midi.unlock();
                }
                saveSequence(true);
                break; 
              }
              if (menuSelection == 3) { currentState = UI_EDIT_TEMPO; break; } // Tempo
              if (menuSelection == 4) { // Song Mode
                songModeEnabled = !songModeEnabled;
                if (songModeEnabled) {
                  songModeNeedsNext = true;
                }
                break; 
              }

              if (menuSelection == 5) { currentState = UI_RANDOMIZE_TRACK_EDIT; break; } // Track
              if (menuSelection == 6) { trackMute[randomizeTrack] = !trackMute[randomizeTrack]; break; } // Mute
              if (menuSelection == 7) { currentState = UI_EDIT_FLAVOUR;  break; } // Flavour
              if (menuSelection == 8) { mutationEnabled = !mutationEnabled; break; } // Mutation

              if (menuSelection >= THEME_1_INDEX_POLY) { // Themes
                const int selectedTheme = menuSelection - THEME_1_INDEX_POLY + 1;
                if (isPlaying) {
                  queuedTheme = selectedTheme;
                  midi.lock();
                  generateSequenceData(queuedTheme, nextSequence);
                  sequenceChangeScheduled = true;
                  midi.unlock();
                } else {
                  generateTheme(selectedTheme);
                }
                break; 
              }
              }
              break;
            case UI_MENU_SETUP:
              if (menuSelection == 0) { currentState = UI_MENU_RANDOMIZE; menuSelection = 0; break; }
              if (menuSelection == 1) { currentState = UI_SETUP_CHANNEL_EDIT; break; }
              if (menuSelection == 2) { factoryReset(); break; }
              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;
              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_RANDOMIZE;
              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;
  bool local_mutationEnabled, local_songModeEnabled, local_isPlaying;
  bool local_trackMute[NUM_TRACKS];
  int local_midiChannel;
  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_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));
  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_mutationEnabled, local_songModeEnabled, (const Step (*)[NUM_STEPS])local_sequence, local_playbackStep, local_isPlaying,
    mainMenu, mainMenuCount, randomizeMenuPoly, randomizeMenuPolyCount, setupMenu, setupMenuCount,
    THEME_1_INDEX_POLY, MODE_POLY, local_randomizeTrack, (const bool*)local_trackMute);
}

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];
  bool local_trackMute[NUM_TRACKS];

  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;
  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_RANDOMIZE) {
      // Main section items: Setup(0), Melody(1), Scale(2), Tempo(3), Song Mode(4)
      bool isMainSection = (local_menuSelection <= 4);
      if (!isMainSection) {
          // 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) {
      // 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_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
}