#include <SPI.h>
#include <Wire.h>
#include <EEPROM.h>
#include "TrackerTypes.h"
#include "MelodyStrategy.h"
#include "LuckyStrategy.h"
#include "ArpStrategy.h"
#include "MidiDriver.h"
#include "UIManager.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

// --- TRACKER DATA ---
#define NUM_STEPS 16

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

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

  midi.lock();
  for (int i=0; i<NUM_STEPS; i++) {
    EEPROM.put(addr, sequence[i]); addr += sizeof(Step);
  }
  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;
  
  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);
  }

  midi.lock();
  for (int i=0; i<NUM_STEPS; i++) {
    EEPROM.get(addr, sequence[i]); addr += sizeof(Step);
  }
  midi.unlock();
  return true;
}

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());
  generateRandomScale();
  melodySeed = random(10000);
  
  EEPROM.begin(512);
  if (!loadSequence()) {
    generateTheme(1);
  }
  isPlaying = false; // Don't start playing on boot

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

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;
  midi.lock();
  generateSequenceData(themeType, sequence);
  midi.unlock();

  clockCount = 0;
  lastClockTime = micros();
  playbackStep = 0;
  midi.sendRealtime(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;
          midi.lock();
          sequence[stepIndex].note = newNote;
          midi.unlock();
        } 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;
                  midi.lock();
                  generateSequenceData(theme, nextSequence);
                  sequenceChangeScheduled = true;
                  midi.unlock();
                }
                saveSequence(true);
              }
              if (menuSelection == 2) {
                melodySeed = random(10000); // Melody
                if (isPlaying) {
                  int theme = (queuedTheme != -1) ? queuedTheme : currentThemeIndex;
                  midi.lock();
                  generateSequenceData(theme, nextSequence);
                  sequenceChangeScheduled = true;
                  midi.unlock();
                }
                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;
                  midi.lock();
                  generateSequenceData(queuedTheme, nextSequence);
                  sequenceChangeScheduled = true;
                  midi.unlock();
                } 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()) ui.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();
        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)

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

    // 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) {
      midi.sendNoteOff(prevNote, shMidiChannel);
    }
    
    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;
      }

      midi.panic(shMidiChannel); // 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;
      midi.sendNoteOn(sequence[playbackStep].note, velocity, shMidiChannel);
    }

    // Note Off for previous step (if tied - delayed Note Off)
    if (isTied && prevNote != -1) {
      midi.sendNoteOff(prevNote, shMidiChannel);
    }

    midi.unlock();
  }
}

void drawUI() {
  midi.lock();
  ui.draw(currentState, menuSelection, navigationSelection, isEditing, midiChannel, tempo, strategies[currentStrategyIndex], queuedTheme, currentThemeIndex, numScaleNotes, scaleNotes, melodySeed, mutationEnabled, songModeEnabled, sequence, scrollOffset, playbackStep, isPlaying, mainMenu, mainMenuCount, randomizeMenu, randomizeMenuCount, setupMenu, setupMenuCount, THEME_1_INDEX);
  midi.unlock();
}

void updateLeds() {
  midi.lock();
  ui.updateLeds(sequence, navigationSelection, playbackStep, isPlaying, currentState, isEditing, songModeEnabled, songRepeatsRemaining, sequenceChangeScheduled);
  midi.unlock();
}

void loop1() {
  if (needsPanic) {
    midi.lock();
    midi.panic(shMidiChannel);
    midi.unlock();
    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
    
    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
}