#include <SPI.h>
#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include <Adafruit_NeoPixel.h>
#include <EEPROM.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

// MIDI UART Pins (GP0/GP1) -- OUT only so far
#define PIN_MIDI_TX 0

// NeoPixel Pin (any GPIO is fine, I've chosen 16)
#define PIN_NEOPIXEL 16
#define NUM_PIXELS   64 // For 8x8 WS2812B matrix

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

struct Step {
  int8_t note; // MIDI Note (0-127), -1 for OFF
  bool accent;
  bool tie;
};

Step sequence[NUM_STEPS];

// --- STATE ---
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);
Adafruit_NeoPixel pixels(NUM_PIXELS, PIN_NEOPIXEL, NEO_GRB + NEO_KHZ800);

enum UIState {
  UI_TRACKER,
  UI_MENU_MAIN,
  UI_MENU_RANDOMIZE,
  UI_MENU_SETUP,
  UI_SETUP_CHANNEL_EDIT
};

UIState currentState = UI_TRACKER;

const char* mainMenu[] = { "Tracker", "Randomize", "Setup" };
const int mainMenuCount = sizeof(mainMenu) / sizeof(char*);
const char* randomizeMenu[] = { "Back", "Scale", "Melody", "Theme 1", "Theme 2", "Theme 3", "Theme 4", "Theme 5", "Theme 6", "Theme 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;
int navigationSelection = 1;
int playbackStep = 0;
int midiChannel = 1;
int scaleNotes[12];
int numScaleNotes = 0;
int melodySeed = 0;
int queuedTheme = -1;
const uint32_t EEPROM_MAGIC = 0x42424244;

bool isEditing = false;
int scrollOffset = 0;
bool isPlaying = false;
unsigned long lastStepTime = 0;
int tempo = 120; // BPM


// 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 showMessage(const char* msg) {
  display.clearDisplay();
  display.setCursor(10, 25);
  display.setTextColor(SSD1306_WHITE);
  display.setTextSize(2);
  display.print(msg);
  display.display();
  delay(500);
  display.setTextSize(1);
}

void saveSequence() {
  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, numScaleNotes); addr += sizeof(numScaleNotes);
  for (int i=0; i<12; i++) {
    EEPROM.put(addr, scaleNotes[i]); addr += sizeof(int);
  }

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

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

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

  // 2. Setup Display
  // Note: Using default I2C pins (SDA=GP4, SCL=GP5) which works on both cores.
  Wire.begin();

  // SSD1306_SWITCHCAPVCC = generate display voltage from 3.3V internally
  if(!display.begin(SSD1306_SWITCHCAPVCC, SCREEN_ADDRESS)) {
    Serial.println(F("SSD1306 allocation failed"));
    for(;;); // Don't proceed, loop forever
  }

  // 3. Setup NeoPixel Matrix
  pixels.begin();
  pixels.setBrightness(40); // Set brightness to a medium-low value (0-255)
  pixels.clear();
  pixels.show();

  // 4. Setup MIDI Serial
  Serial1.setTX(PIN_MIDI_TX);
  Serial1.begin(31250);
  Serial.println(F("MIDI Serial initialized on GP0/GP1"));

  // 5. Init Sequence
  randomSeed(micros());
  generateRandomScale();
  melodySeed = random(10000);
  
  EEPROM.begin(512);
  if (!loadSequence()) {
    generateTheme(1);
  }
  isPlaying = false; // Don't start playing on boot

  display.clearDisplay();
  display.display();

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

void sendMidi(uint8_t status, uint8_t note, uint8_t velocity) {
  uint8_t channelStatus = status | (midiChannel - 1);
  Serial1.write(channelStatus);
  Serial1.write(note);
  Serial1.write(velocity);
}

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 generateTheme(int themeType) {
  sendMidi(0xB0, 123, 0); // Panic / All Notes Off
  randomSeed(melodySeed + themeType * 12345); // Deterministic seed for this theme
  if (numScaleNotes == 0) generateRandomScale();

  for (int i = 0; i < NUM_STEPS; i++) {
    int octave = random(3) + 3; // 3, 4, 5 (Base is 4)
    sequence[i].note = (random(100) < 50) ? (12 * octave + scaleNotes[random(numScaleNotes)]) : -1;
    sequence[i].accent = (random(100) < 30);
    sequence[i].tie = (random(100) < 20);
  }
  randomSeed(micros()); // Restore randomness
  isPlaying = true;
}

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;
          sequence[stepIndex].note = newNote;
        } 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;
        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 (menuSelection == 2) melodySeed = random(10000); // Melody
              if (menuSelection >= 3) { // Themes
                if (isPlaying) {
                  queuedTheme = menuSelection - 2;
                } else {
                  generateTheme(menuSelection - 2);
                }
              }
              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()) showMessage("LOADED!"); 
                currentState = UI_MENU_MAIN; 
              }
              break;
            case UI_SETUP_CHANNEL_EDIT:
              currentState = UI_MENU_SETUP;
              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 = navigationSelection > 0 ? navigationSelection - 1 : 0;
        lastStepTime = millis(); // Reset timer to start immediately
      } else {
        // Send All Notes Off on stop (CC 123)
        sendMidi(0xB0, 123, 0); 
        queuedTheme = -1;
      }
    }
  }
  
  lastButtonState = reading;
}

void handlePlayback() {
  if (!isPlaying) return;

  unsigned long interval = 15000 / tempo; // 16th notes (60000 / tempo / 4)
  if (millis() - lastStepTime > interval) {
    lastStepTime = millis();
    
    // 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) {
      sendMidi(0x80, prevNote, 0);
    }
    
    playbackStep++;
    if (playbackStep >= NUM_STEPS) {
      playbackStep = 0;
      if (queuedTheme != -1) {
        generateTheme(queuedTheme);
        queuedTheme = -1;
      }
    }
    
    // Note On for new step
    if (sequence[playbackStep].note != -1) {
      uint8_t velocity = sequence[playbackStep].accent ? 127 : 100;
      sendMidi(0x90, sequence[playbackStep].note, velocity);
    }

    // Note Off for previous step (if tied - delayed Note Off)
    if (isTied && prevNote != -1) {
      sendMidi(0x80, prevNote, 0);
    }
    
    // Auto-scroll navigation cursor if not editing
    if (!isEditing) {
      navigationSelection = playbackStep + 1; // +1 because 0 is menu
      if (navigationSelection < scrollOffset) scrollOffset = navigationSelection;
      if (navigationSelection >= scrollOffset + 6) scrollOffset = navigationSelection - 5;
    }
  }
}

void drawMenu(const char* title, const char* items[], int count, int selection) {
  display.println(title);
  display.drawLine(0, 8, 128, 8, SSD1306_WHITE);

  // Simple scrolling: keep selection visible
  int start = 0;
  if (selection >= 5) {
    start = selection - 4;
  }

  int y = 10;
  for (int i = start; i < count; i++) {
    if (y > 55) break; // Stop if we run out of screen space

    if (i == selection) {
      display.fillRect(0, y, 128, 8, SSD1306_WHITE);
      display.setTextColor(SSD1306_BLACK, SSD1306_WHITE);
    } else {
      display.setTextColor(SSD1306_WHITE);
    }
    display.setCursor(2, y);
    display.print(items[i]);

    // Special case for channel display
    if (currentState == UI_MENU_SETUP && i == 1) {
        display.print(F(": "));
        display.print(midiChannel);
    }

    // Special case for queued theme
    if (currentState == UI_MENU_RANDOMIZE && i >= 3 && queuedTheme == (i - 2)) {
        display.print(F(" [NEXT]"));
    }

    // Special cases for Randomize Menu values
    if (currentState == UI_MENU_RANDOMIZE) {
      if (i == 1) { // Scale
        display.print(F(": "));
        if (numScaleNotes > 0) {
          const char* noteNames[] = {"C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B"};
          for (int j = 0; j < min(numScaleNotes, 6); j++) {
            display.print(noteNames[scaleNotes[j]]);
            if (j < min(numScaleNotes, 6) - 1) display.print(F(" "));
          }
        }
      } else if (i == 2) { // Melody
        display.print(F(": "));
        display.print(melodySeed);
      }
    }
    y += 9;
  }
}

void drawTracker() {
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(0, 0);

  // Header
  display.print(F("SEQ "));
  if (navigationSelection > 0 && isEditing) {
    display.print(F("[EDIT]"));
  } else {
    display.print(F("[NAV] "));
  }
  display.print(F(" CH:"));
  display.print(midiChannel);

  display.println();
  display.drawLine(0, 8, 128, 8, SSD1306_WHITE);

  // Steps
  int y = 10;
  for (int i = 0; i < 6; i++) {
    int itemIndex = i + scrollOffset;
    if (itemIndex > NUM_STEPS) break;

    // Draw Cursor
    if (itemIndex == navigationSelection) {
      display.fillRect(0, y, 128, 8, SSD1306_WHITE);
      display.setTextColor(SSD1306_BLACK, SSD1306_WHITE); // Invert text
    } else {
      display.setTextColor(SSD1306_WHITE);
    }
    display.setCursor(2, y);

    if (itemIndex == 0) {
      display.print(F(">> MENU"));
    } else {
      int stepIndex = itemIndex - 1;
      // Step Number
      if (stepIndex < 10) display.print(F("0"));
      display.print(stepIndex);
      display.print(F(" | "));

      // Note Value
      int n = sequence[stepIndex].note;
      if (n == -1) {
        display.print(F("---"));
      } else {
        // Basic Note to String conversion
        const char* noteNames[] = {"C-", "C#", "D-", "D#", "E-", "F-", "F#", "G-", "G#", "A-", "A#", "B-"};
        display.print(noteNames[n % 12]);
        display.print(n / 12 - 1); // Octave
      }
    }

    y += 9;
  }
}

void drawUI() {
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(0, 0);

  switch(currentState) {
    case UI_TRACKER:
      drawTracker();
      break;
    case UI_MENU_MAIN:
      drawMenu("MAIN MENU", mainMenu, mainMenuCount, menuSelection);
      break;
    case UI_MENU_RANDOMIZE:
      drawMenu("RANDOMIZE", randomizeMenu, randomizeMenuCount, menuSelection);
      break;
    case UI_MENU_SETUP:
      drawMenu("SETUP", setupMenu, setupMenuCount, menuSelection);
      break;
    case UI_SETUP_CHANNEL_EDIT:
      display.println(F("SET MIDI CHANNEL"));
      display.drawLine(0, 8, 128, 8, SSD1306_WHITE);
      display.setCursor(20, 25);
      display.setTextSize(2);
      display.print(F("CH: "));
      if (midiChannel < 10) display.print(F(" "));
      display.print(midiChannel);
      display.setTextSize(1);
      display.setCursor(0, 50);
      display.println(F(" (Press to confirm)"));
      break;
  }

  display.display();
}

// Helper to convert X,Y to pixel index for an 8x8 matrix.
// Assumes row-major wiring (NOT serpentine).
// If your matrix is wired differently, you'll need to change this function.
int getPixelIndex(int x, int y) {
  return y * 8 + x;
}

uint32_t getNoteColor(int note) {
  if (note == -1) return 0;
  // Map note to hue, avoiding Green (approx 21845) which is used for playback cursor.
  // We start from Cyan (~30000) and go up to Orange (~8000 wrapped).
  // Range: 30000 to 65536+8000 = 73536. Width = 43536.
  // Step per semitone: 43536 / 12 = 3628.
  // This ensures notes are distinct colors but never pure Green.
  uint16_t hue = 30000 + (note % 12) * 3628;
  return Adafruit_NeoPixel::ColorHSV(hue, 255, 50);
}

void updateLeds() {
  pixels.clear(); // Clear buffer

  for (int s = 0; s < NUM_STEPS; s++) {
    int blockX = (s % 4) * 2;
    int blockY = (s / 4) * 2;

    // --- Top Row: Attributes ---
    uint32_t colorTL = 0; // Octave
    uint32_t colorTR = 0; // Accent/Tie

    if (sequence[s].note != -1) {
      int octave = sequence[s].note / 12;
      // Base octave 4 (MIDI 48-59).
      if (octave > 4) colorTL = pixels.Color(0, 50, 0); // Up (Green)
      else if (octave < 4) colorTL = pixels.Color(50, 0, 0); // Down (Red)
      
      if (sequence[s].accent && sequence[s].tie) colorTR = pixels.Color(50, 0, 50); // Both (Purple)
      else if (sequence[s].accent) colorTR = pixels.Color(50, 50, 50); // Accent (White)
      else if (sequence[s].tie) colorTR = pixels.Color(50, 50, 0); // Tie (Yellow)
    }

    // --- Bottom Row: Note / Cursor ---
    uint32_t colorBL = 0;
    uint32_t colorBR = 0;

    if (sequence[s].note != -1) {
      uint32_t noteColor = getNoteColor(sequence[s].note);
      colorBL = noteColor;
      colorBR = noteColor;
    }

    int stepNavIndex = navigationSelection - 1;
    if (isPlaying && s == playbackStep) {
      colorBL = colorBR = pixels.Color(0, 50, 0); // Green for playback
    } else if (currentState == UI_TRACKER && s == stepNavIndex) {
      colorBL = colorBR = isEditing ? pixels.Color(50, 0, 0) : pixels.Color(40, 40, 40);
    }

    pixels.setPixelColor(getPixelIndex(blockX, blockY), colorTL);
    pixels.setPixelColor(getPixelIndex(blockX + 1, blockY), colorTR);
    pixels.setPixelColor(getPixelIndex(blockX, blockY + 1), colorBL);
    pixels.setPixelColor(getPixelIndex(blockX + 1, blockY + 1), colorBR);
  }

  pixels.show();
}

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