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

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

UIState currentState = UI_MENU_RANDOMIZE; // Let's start in the Play menu

const char* mainMenu[] = { "Randomize", "Setup" };
const int mainMenuCount = sizeof(mainMenu) / sizeof(char*);

const char* randomizeMenuMono[] = { "Setup", "Melody", "Flavour", "Scale", "Tempo", "Mutation", "Song Mode", "Theme 1", "Theme 2", "Theme 3", "Theme 4", "Theme 5", "Theme 6", "Theme 7" };
const int randomizeMenuMonoCount = sizeof(randomizeMenuMono) / sizeof(char*);
const int THEME_1_INDEX_MONO = 7;

const char* randomizeMenuPoly[] = { "Setup", "Track", "Mute", "Melody", "Flavour", "Scale", "Tempo", "Mutation", "Song Mode", "Theme 1", "Theme 2", "Theme 3", "Theme 4", "Theme 5", "Theme 6", "Theme 7" };
const int randomizeMenuPolyCount = sizeof(randomizeMenuPoly) / sizeof(char*);
const int THEME_1_INDEX_POLY = 9;

const char* setupMenu[] = { "Back", "Play Mode", "Channel", "Factory Reset" };
const int setupMenuCount = sizeof(setupMenu) / sizeof(char*);

int menuSelection = 0;
volatile bool trackMute[NUM_TRACKS];
int randomizeTrack = 0;
volatile int playbackStep = 0;
volatile int midiChannels[NUM_TRACKS];
int scaleNotes[12];
int numScaleNotes = 0;
int melodySeeds[NUM_TRACKS];
volatile int queuedTheme = -1;
volatile int currentThemeIndex = 1;
const uint32_t EEPROM_MAGIC = 0x4242424B;

MelodyStrategy* strategies[] = { new LuckyStrategy(), new ArpStrategy(), new EuclideanStrategy() };
const int numStrategies = 3;
int currentStrategyIndices[NUM_TRACKS];

volatile PlayMode playMode = MODE_MONO;
volatile bool mutationEnabled = false;
volatile bool songModeEnabled = false;
volatile int songRepeatsRemaining = 0;
volatile int nextSongRepeats = 0;
volatile bool songModeNeedsNext = false;
volatile bool isPlaying = false;
volatile int tempo = 120; // BPM
volatile unsigned long lastClockTime = 0;
volatile int clockCount = 0;


// Watchdog
volatile unsigned long lastLoop0Time = 0;
volatile unsigned long lastLoop1Time = 0;
volatile bool watchdogActive = false;

// 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 saveSequence(bool quiet = false) {
  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, (int)playMode); 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);
  }

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

bool loadSequence() {
  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, t); addr += sizeof(int);
  playMode = (PlayMode)t;

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

  midi.lock();
  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();
}

void setup() {
  Serial.begin(115200);

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

  ui.begin();
  midi.begin();

  // 5. Init Sequence
  randomSeed(micros());
  
  EEPROM.begin(512);
  if (!loadSequence()) {
    generateRandomScale();
    for(int i=0; i<NUM_TRACKS; i++) {
      midiChannels[i] = i + 1;
      melodySeeds[i] = random(10000);
      currentStrategyIndices[i] = 0;
      trackMute[i] = false;
    }
    generateTheme(1);
  }
  isPlaying = false; // Don't start playing on boot

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

  // Enable Watchdog
  lastLoop0Time = millis();
  lastLoop1Time = millis();
  watchdogActive = true;
}

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

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

void generateTheme(int themeType) {
  currentThemeIndex = themeType;
  needsPanic = true;
  midi.lock();
  generateSequenceData(themeType, sequence);
  midi.unlock();

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

void mutateSequence(Step (*target)[NUM_STEPS]) {
  if (playMode == MODE_POLY) {
    for(int i=0; i<NUM_TRACKS; i++) strategies[currentStrategyIndices[i]]->mutate(target, i, NUM_STEPS, scaleNotes, numScaleNotes);
  } else {
    strategies[currentStrategyIndices[0]]->mutate(target, 0, 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_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);
          int count = (playMode == MODE_POLY) ? randomizeMenuPolyCount : randomizeMenuMonoCount;
          if (menuSelection < 0) menuSelection = count - 1;
          if (menuSelection >= count) 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:
        {
          int trackToEdit = (playMode == MODE_POLY) ? randomizeTrack : 0;
          midiChannels[trackToEdit] += (delta > 0 ? 1 : -1);
          if (midiChannels[trackToEdit] < 1) midiChannels[trackToEdit] = 16;
          if (midiChannels[trackToEdit] > 16) midiChannels[trackToEdit] = 1;
        }
        break;
      case UI_EDIT_TEMPO:
        tempo += delta;
        if (tempo < 40) tempo = 40;
        if (tempo > 240) tempo = 240;
        break;
      case UI_EDIT_FLAVOUR:
        {
          int trackToEdit = playMode == MODE_POLY ? randomizeTrack : 0;
          currentStrategyIndices[trackToEdit] += (delta > 0 ? 1 : -1);
          if (currentStrategyIndices[trackToEdit] < 0) currentStrategyIndices[trackToEdit] = numStrategies - 1;
          if (currentStrategyIndices[trackToEdit] >= numStrategies) currentStrategyIndices[trackToEdit] = 0;
        }
        break;
      case UI_SETUP_PLAYMODE_EDIT:
        playMode = (playMode == MODE_MONO) ? MODE_POLY : MODE_MONO;
        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:
              {
              int track_offset = (playMode == MODE_POLY) ? 2 : 0;
              int theme_1_index = (playMode == MODE_POLY) ? THEME_1_INDEX_POLY : THEME_1_INDEX_MONO;
              if (menuSelection == 0) { currentState = UI_MENU_SETUP; menuSelection = 0; break;  }
              if (playMode == MODE_POLY) {
                  if (menuSelection == 1) { currentState = UI_RANDOMIZE_TRACK_EDIT; break; }
                  if (menuSelection == 2) { trackMute[randomizeTrack] = !trackMute[randomizeTrack]; break; }
              }

              if (menuSelection == 1 + track_offset) { // Melody
                int track = playMode == MODE_POLY ? randomizeTrack : 0;
                midi.lock();
                melodySeeds[track] = random(10000);
                if (isPlaying) {
                  int theme = (queuedTheme != -1) ? queuedTheme : currentThemeIndex;
                  if (!sequenceChangeScheduled) {
                      memcpy(nextSequence, sequence, sizeof(sequence));
                  }
                  generateTrackData(track, theme, nextSequence);
                  sequenceChangeScheduled = true;
                }
                midi.unlock();
                saveSequence(true);
                break; 
              }
              if (menuSelection == 2 + track_offset) { currentState = UI_EDIT_FLAVOUR;  break; } // Flavour
              if (menuSelection == 3 + track_offset) { // 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 == 4 + track_offset) { currentState = UI_EDIT_TEMPO; break; }
              if (menuSelection == 5 + track_offset) { mutationEnabled = !mutationEnabled; break; }
              if (menuSelection == 6 + track_offset) {
                songModeEnabled = !songModeEnabled;
                if (songModeEnabled) {
                  songModeNeedsNext = true;
                }
                break; 
              }
              if (menuSelection >= theme_1_index) { // Themes
                const int selectedTheme = menuSelection - theme_1_index + 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_PLAYMODE_EDIT; break; }
              if (menuSelection == 2) { currentState = UI_SETUP_CHANNEL_EDIT; break; }
              if (menuSelection == 3) { 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;
                  int track = playMode == MODE_POLY ? randomizeTrack : 0;
                  midi.lock();
                  if (!sequenceChangeScheduled) {
                      memcpy(nextSequence, sequence, sizeof(sequence));
                  }
                  generateTrackData(track, theme, nextSequence);
                  sequenceChangeScheduled = true;
                  midi.unlock();
              }
              saveSequence(true);
              break;
            case UI_SETUP_PLAYMODE_EDIT:
              currentState = UI_MENU_SETUP;
              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();
      midi.sendRealtime(0xFA); // MIDI Start
    } else {
      // Send All Notes Off on stop (CC 123)
      needsPanic = true;
      midi.sendRealtime(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)
  int tracksToPlay = (playMode == MODE_POLY) ? NUM_TRACKS : 1;

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

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

    midi.lock();

    

    for(int t=0; t<tracksToPlay; t++) {
      int trackChannel = playMode == MODE_POLY ? midiChannels[t] : midiChannels[0];
      int nextStep = playbackStep + 1;
      if (nextStep >= NUM_STEPS) nextStep = 0;
      
      // Determine if we are tying to the next note
      bool isTied = sequence[t][playbackStep].tie && (sequence[t][nextStep].note != -1);
      int prevNote = sequence[t][playbackStep].note;

      // Note Off for previous step (if NOT tied)
      if (!isTied && prevNote != -1) {
        midi.sendNoteOff(prevNote, trackChannel);
      }
    }
    
    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;
      }

      for (int i=0; i<tracksToPlay; i++) midi.panic(midiChannels[i]);
      
      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
    for(int t=0; t<tracksToPlay; t++) {
      int trackChannel = playMode == MODE_POLY ? midiChannels[t] : midiChannels[0];
      if (!trackMute[t] && sequence[t][playbackStep].note != -1) {
        uint8_t velocity = sequence[t][playbackStep].accent ? 127 : 100;
        midi.sendNoteOn(sequence[t][playbackStep].note, velocity, trackChannel);
      }

      int prevStep = (playbackStep == 0) ? NUM_STEPS - 1 : playbackStep - 1;
      bool wasTied = sequence[t][prevStep].tie && (sequence[t][playbackStep].note != -1);
      int prevNote = sequence[t][prevStep].note;
      // Note Off for previous step (if tied - delayed Note Off)
      if (wasTied && prevNote != -1) {
        midi.sendNoteOff(prevNote, trackChannel);
      }
    }

    midi.unlock();
  }
}

void drawUI() {
  midi.lock();
  const char **randMenu;
  int randMenuCount;
  int themeIndex;
  if (playMode == MODE_POLY) {
    randMenu = randomizeMenuPoly;
    randMenuCount = randomizeMenuPolyCount;
    themeIndex = THEME_1_INDEX_POLY;
  } else {
    randMenu = randomizeMenuMono;
    randMenuCount = randomizeMenuMonoCount;
    themeIndex = THEME_1_INDEX_MONO;
  }
  
  int ui_track = 0;
  if (playMode == MODE_POLY) {
    ui_track = randomizeTrack;
  }

  ui.draw(currentState, menuSelection, 
    midiChannels[ui_track], tempo, strategies[currentStrategyIndices[ui_track]], 
    queuedTheme, currentThemeIndex, numScaleNotes, scaleNotes, melodySeeds[ui_track], 
    mutationEnabled, songModeEnabled, sequence, playbackStep, isPlaying,
    mainMenu, mainMenuCount, randMenu, randMenuCount, setupMenu, setupMenuCount,
    themeIndex, playMode, randomizeTrack, (const bool*)trackMute);
  midi.unlock();
}

void updateLeds() {
  midi.lock();
  ui.updateLeds(sequence, playbackStep, isPlaying, currentState, songModeEnabled, songRepeatsRemaining, sequenceChangeScheduled, playMode, numScaleNotes, scaleNotes, (const bool*)trackMute);
  midi.unlock();
}

void loop1() {
  unsigned long now = millis();
  lastLoop1Time = now;
  if (watchdogActive && (now - lastLoop0Time > 1000)) {
    Serial.println("Core 0 Freeze detected");
    rp2040.reboot();
  }

  if (needsPanic) {
    midi.lock();
    if (playMode == MODE_POLY) {
      for (int i=0; i<NUM_TRACKS; i++) midi.panic(midiChannels[i]);
    } else {
      midi.panic(midiChannels[0]);
    }
    midi.unlock();
    needsPanic = false;
  }
  handlePlayback();
}

void loop() {
  unsigned long now = millis();
  lastLoop0Time = now;
  if (watchdogActive && (now - lastLoop1Time > 1000)) {
    Serial.println("Core 1 Freeze detected");
    rp2040.reboot();
  }

  // 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
    
    midi.lock();
    generateSequenceData(nextTheme, nextSequence);
    queuedTheme = nextTheme;
    nextSongRepeats = repeats;
    sequenceChangeScheduled = true;
    midi.unlock();
    
    songModeNeedsNext = false;
  }

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