#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include <Encoder.h>

#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define OLED_RESET -1
#define SCREEN_ADDRESS 0x3C

// Hardware pins
#define ENCODER_PIN_A 2
#define ENCODER_PIN_B 3
#define ENCODER_BTN 4
#define SPEAKER_PIN 8

// Sequencer constants
#define MAX_STEPS 16
#define MIN_BPM 40
#define MAX_BPM 300
#define DEBOUNCE_DELAY 50

// Musical scale (A minor pentatonic)
const int SCALE[] = {220, 262, 294, 330, 392, 440, 523, 587};
const int SCALE_LENGTH = 8;

// View definitions
#define VIEW_SEQUENCE 0
#define VIEW_STEP_EDIT 1
#define VIEW_SETTINGS 2
#define TOTAL_VIEWS 3

Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);
Encoder encoder(ENCODER_PIN_A, ENCODER_PIN_B);

// Enum for Step Edit submenu, sound types, and sequence types
enum StepEditMode {
  EDIT_STEP,
  EDIT_NOTE,
  EDIT_SOUND_TYPE
};

enum SoundType {
  SINE,
  SQUARE,
  NOISE,
  TOTAL_SOUND_TYPES
};

enum SequenceType {
  LINEAR,
  REVERSE,
  PING_PONG,
  RANDOM,
  TOTAL_SEQUENCE_TYPES
};

// State struct with stepEditMode and soundType array
struct {
  int bpm = 120;
  int activeSteps = 8;
  int currentStep = 0;
  byte notes[MAX_STEPS] = {0};  // Initialize with all notes set to 0
  bool isPlaying = true;
  
  int currentView = VIEW_SEQUENCE;
  int editingStep = 0;
  bool isEditing = false;
  StepEditMode stepEditMode = EDIT_STEP;
  SoundType soundTypes[MAX_STEPS] = {SQUARE}; // Default sound type for each step
  SequenceType sequenceType = LINEAR;
  bool pingPongDirection = true; // For ping pong mode
} state;

// Timing variables
unsigned long lastStepTime = 0;
unsigned long lastButtonPress = 0;
bool buttonPressed = false;
long lastEncoderValue = 0;

void setup() {
  pinMode(ENCODER_BTN, INPUT_PULLUP);
  pinMode(SPEAKER_PIN, OUTPUT);

  if(!display.begin(SSD1306_SWITCHCAPVCC, SCREEN_ADDRESS)) {
    for(;;);
  }
  
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
}

void loop() {
  handleInput();
  updateSequencer();
  updateDisplay();
}

void handleInput() {
  long encoderValue = encoder.read() / 4;
  if (encoderValue != lastEncoderValue) {
    handleEncoderChange(encoderValue - lastEncoderValue);
    lastEncoderValue = encoderValue;
  }

  bool currentBtnState = digitalRead(ENCODER_BTN);
  if (!currentBtnState && !buttonPressed && (millis() - lastButtonPress > DEBOUNCE_DELAY)) {
    handleButtonPress();
    buttonPressed = true;
    lastButtonPress = millis();
  } else if (currentBtnState) {
    buttonPressed = false;
  }
}

void handleEncoderChange(int change) {
  switch(state.currentView) {
    case VIEW_SEQUENCE:
      if (state.isEditing) {
        state.editingStep = constrain(state.editingStep + change, 0, state.activeSteps - 1);
      } else {
        state.currentStep = (state.currentStep + change + state.activeSteps) % state.activeSteps;
      }
      break;
      
    case VIEW_STEP_EDIT:
      if (state.isEditing) {
        switch(state.stepEditMode) {
          case EDIT_STEP:
            state.editingStep = constrain(state.editingStep + change, 0, state.activeSteps - 1);
            break;
          case EDIT_NOTE:
            state.notes[state.editingStep] = constrain(state.notes[state.editingStep] + change, 0, SCALE_LENGTH - 1);
            break;
          case EDIT_SOUND_TYPE:
            state.soundTypes[state.editingStep] = static_cast<SoundType>((state.soundTypes[state.editingStep] + change + TOTAL_SOUND_TYPES) % TOTAL_SOUND_TYPES);
            break;
        }
      } else {
        state.stepEditMode = static_cast<StepEditMode>((state.stepEditMode + change + 3) % 3);
      }
      break;
      
    case VIEW_SETTINGS:
      if (state.isEditing) {
        state.bpm = constrain(state.bpm + change * 5, MIN_BPM, MAX_BPM);
      } else {
        if (change > 0) {
          state.activeSteps = constrain(state.activeSteps + change, 1, MAX_STEPS);
        } else {
          state.sequenceType = static_cast<SequenceType>((state.sequenceType + change + TOTAL_SEQUENCE_TYPES) % TOTAL_SEQUENCE_TYPES);
        }
      }
      break;
  }
}

void handleButtonPress() {
  if (state.isEditing) {
    state.isEditing = false;
  } else {
    delay(50);
    if (!digitalRead(ENCODER_BTN)) {
      unsigned long pressStart = millis();
      while (!digitalRead(ENCODER_BTN) && millis() - pressStart < 500);
      
      if (millis() - pressStart >= 500) {
        state.currentView = (state.currentView + 1) % TOTAL_VIEWS;
        state.stepEditMode = EDIT_STEP;
      } else {
        state.isEditing = true;
      }
    }
  }
}

void updateDisplay() {
  display.clearDisplay();
  drawStatusBar();
  
  switch(state.currentView) {
    case VIEW_SEQUENCE:
      drawSequenceView();
      break;
    case VIEW_STEP_EDIT:
      drawStepEditView();
      break;
    case VIEW_SETTINGS:
      drawSettingsView();
      break;
  }
  
  display.display();
}

void drawStatusBar() {
  display.fillRect(0, 0, SCREEN_WIDTH, 10, SSD1306_WHITE);
  display.setTextColor(SSD1306_BLACK);
  display.setCursor(2, 1);
  display.print((state.currentView == VIEW_SEQUENCE) ? "SEQUENCE" : (state.currentView == VIEW_STEP_EDIT) ? "STEP EDIT" : "SETTINGS");
  display.setCursor(SCREEN_WIDTH - 30, 1);
  display.print(state.isPlaying ? ">" : "||");
}

void drawSequenceView() {
  display.setTextColor(SSD1306_WHITE);
  int stepWidth = SCREEN_WIDTH / state.activeSteps;
  int gridY = 15;
  int gridHeight = 35;
  
  for(int i = 0; i < state.activeSteps; i++) {
    int x = i * stepWidth;
    int noteHeight = map(state.notes[i], 0, SCALE_LENGTH - 1, 0, gridHeight);
    display.drawRect(x + 1, gridY, stepWidth - 2, gridHeight, SSD1306_WHITE);
    if (noteHeight > 0) {
      display.fillRect(x + 1, gridY + gridHeight - noteHeight, stepWidth - 2, noteHeight, i == state.currentStep ? SSD1306_WHITE : SSD1306_INVERSE);
    }
    if (i == state.editingStep && state.isEditing) {
      display.drawRect(x, gridY - 1, stepWidth, gridHeight + 2, SSD1306_WHITE);
    }
  }
}

void drawStepEditView() {
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(2, 15);
  display.print("Step: ");
  display.print(state.editingStep + 1);
  
  display.setCursor(2, 30);
  display.print("Note: ");
  display.print(state.notes[state.editingStep]);
  
  display.setCursor(2, 45);
  display.print("Type: ");
  display.print(state.soundTypes[state.editingStep] == SINE ? "Sine" : state.soundTypes[state.editingStep] == SQUARE ? "Square" : "Noise");
  
  int triangleY = (state.stepEditMode == EDIT_STEP) ? 17 : (state.stepEditMode == EDIT_NOTE) ? 32 : 47;
  if (state.isEditing) {
    display.fillTriangle(90, triangleY, 95, triangleY + 4, 90, triangleY + 8, SSD1306_WHITE);
  } else {
    display.drawTriangle(90, triangleY, 95, triangleY + 4, 90, triangleY + 8, SSD1306_WHITE);
  }
}

void drawSettingsView() {
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(2, 15);
  display.print("BPM: ");
  display.print(state.bpm);
  
  display.setCursor(2, 30);
  display.print("Steps: ");
  display.print(state.activeSteps);
  
  display.setCursor(2, 45);
  display.print("Mode: ");
  display.print(state.sequenceType == LINEAR ? "Linear" : state.sequenceType == REVERSE ? "Reverse" : state.sequenceType == PING_PONG ? "Ping Pong" : "Random");
}

void updateSequencer() {
  if (!state.isPlaying) return;

  unsigned long currentMillis = millis();
  unsigned long interval = 60000 / state.bpm;

  if (currentMillis - lastStepTime >= interval) {
    lastStepTime = currentMillis;

    playStep();
    
    if (state.sequenceType == LINEAR) {
      state.currentStep = (state.currentStep + 1) % state.activeSteps;
    } else if (state.sequenceType == REVERSE) {
      state.currentStep = (state.currentStep == 0) ? state.activeSteps - 1 : state.currentStep - 1;
    } else if (state.sequenceType == PING_PONG) {
      if (state.currentStep == 0 || state.currentStep == state.activeSteps - 1) {
        state.pingPongDirection = !state.pingPongDirection;
      }
      state.currentStep += state.pingPongDirection ? 1 : -1;
    } else if (state.sequenceType == RANDOM) {
      state.currentStep = random(state.activeSteps);
    }
  }
}

void playStep() {
  int note = SCALE[state.notes[state.currentStep]];

  if (state.notes[state.currentStep] == 0) {
    noTone(SPEAKER_PIN);  // Mute the step if note is set to 0
  } else {
    switch(state.soundTypes[state.currentStep]) {
      case SINE:
        tone(SPEAKER_PIN, note);  
        break;
      case SQUARE:
        tone(SPEAKER_PIN, note);  
        break;
      case NOISE:
        for (int i = 0; i < 10; i++) {
          tone(SPEAKER_PIN, random(100, 1000), 10);  // Simulate noise with rapid random tones
        }
        break;
    }
  }
}