#include <Arduino.h>
#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include <ESP8266WiFi.h>
#include <espnow.h>

extern "C" {
  #include <user_interface.h>
}

#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, -1);

// === 1. 硬件引脚 ===
const int POT_PIN = A0;
const int BTN_SELECT_PIN = 0; // D3
const int BTN_EXIT_PIN = 14;  // D5
const int BUZZER_PIN = 12;    // D6

// === 2. ESP-NOW 无线通信 ===
uint8_t teacherAddress[] = {0x48, 0x3F, 0xDA, 0x60, 0x7F, 0x0C};

typedef struct struct_message {
  int mode; int score; int errorBPM; bool isActive;
} struct_message;

struct_message sendData;

void sendToTeacher(int m, int s, int err, bool active) {
  sendData.mode = m;
  sendData.score = s;
  sendData.errorBPM = err;
  sendData.isActive = active;
  esp_now_send(teacherAddress, (uint8_t *) &sendData, sizeof(sendData));
}

void OnDataSent(uint8_t *mac_addr, uint8_t sendStatus) {
  Serial.print("Send Status: "); 
  Serial.println(sendStatus == 0 ? "SUCCESS" : "FAIL");
}

// === 3. 状态机 ===
enum State { 
  MAIN_MENU, 
  PITCH_SUB, PRACTICE_PITCH, TEST_PITCH_PLAY, TEST_PITCH_ANSWER, TEST_PITCH_RESULT,
  RHYTHM_SUB, PRACTICE_RHYTHM, TEST_RHYTHM_PLAY, TEST_RHYTHM_ANSWER,
  INTEGRATED_SUB, TEST_INT_PLAY, TEST_INT_PITCH, TEST_INT_RHYTHM, TEST_INT_RESULT 
};
State currentState = MAIN_MENU;

// === 4. 全局变量 ===
int menuIdx = 0; int subMenuIdx = 0;
unsigned long lastBeatTime = 0;
int freqs[] = {262, 294, 330, 349, 392}; 
const char* noteNames[] = {"C4 (Do)", "D4 (Re)", "E4 (Mi)", "F4 (Fa)", "G4 (Sol)"};
int targetNote, userNote, p_total = 0, p_correct = 0;
int targetBPM, userBPM;
int tapCount = 0; unsigned long tapTime = 0; unsigned long userIntervalSum = 0; 
int targetInterval = 0; unsigned long pitchSelectStartTime = 0; 

// === 5. 初始化 ===
void setup() {
  Serial.begin(115200);
  pinMode(BTN_SELECT_PIN, INPUT_PULLUP);
  pinMode(BTN_EXIT_PIN, INPUT_PULLUP);
  pinMode(BUZZER_PIN, OUTPUT);
  
  display.begin(SSD1306_SWITCHCAPVCC, 0x3C);
  display.setTextColor(WHITE);
  randomSeed(analogRead(A0));

  WiFi.mode(WIFI_STA);
  WiFi.disconnect(); 
  
  wifi_set_channel(11);

  if (esp_now_init() != 0) return;
  esp_now_set_self_role(ESP_NOW_ROLE_CONTROLLER);
  esp_now_register_send_cb(OnDataSent); 
  
  // 【关键修改】Peer 绑定也必须是信道 11
  esp_now_add_peer(teacherAddress, ESP_NOW_ROLE_SLAVE, 11, NULL, 0);
  
  Serial.println("Sender Ready on Channel 11");
}

int checkButton() {
  if (digitalRead(BTN_SELECT_PIN) == LOW) {
    delay(30); while (digitalRead(BTN_SELECT_PIN) == LOW); return 1; 
  }
  if (digitalRead(BTN_EXIT_PIN) == LOW) {
    delay(30); while (digitalRead(BTN_EXIT_PIN) == LOW); return 2; 
  }
  return 0;
}

// === 6. 主循环 ===
void loop() {
  int potValue = analogRead(POT_PIN);
  int btn = checkButton();

  // 随时按 D5 强制返回主页面
  if (btn == 2) {
    currentState = MAIN_MENU;
    return;
  }

  switch (currentState) {
    case MAIN_MENU:
      menuIdx = map(potValue, 0, 1023, 0, 2);
      display.clearDisplay(); display.setCursor(0,0);
      display.setTextSize(1);
      display.println("--- MAIN MENU ---");
      display.setCursor(0, 25);
      display.print(menuIdx == 0 ? "> " : "  "); display.println("1. Pitch Training");
      display.print(menuIdx == 1 ? "> " : "  "); display.println("2. Rhythm Training");
      display.print(menuIdx == 2 ? "> " : "  "); display.println("3. Integrated Task");
      display.display();
      if (btn == 1) {
        if (menuIdx == 0) currentState = PITCH_SUB;
        else if (menuIdx == 1) currentState = RHYTHM_SUB;
        else currentState = INTEGRATED_SUB; 
      }
      break;

    case PITCH_SUB:
      subMenuIdx = map(potValue, 0, 1023, 0, 1);
      display.clearDisplay(); display.println("PITCH: SELECT");
      display.setCursor(0, 25);
      display.print(subMenuIdx == 0 ? "> " : "  "); display.println("Practice Mode");
      display.print(subMenuIdx == 1 ? "> " : "  "); display.println("Actual Test");
      display.display();
      if (btn == 1) currentState = (subMenuIdx == 0) ? PRACTICE_PITCH : TEST_PITCH_PLAY;
      break;

    case PRACTICE_PITCH:
      userNote = map(potValue, 0, 1023, 0, 4);
      display.clearDisplay(); display.println("PITCH PRACTICE");
      display.setCursor(30, 30); display.setTextSize(2);
      display.println(noteNames[userNote]); display.setTextSize(1);
      display.display();
      if (btn == 1) tone(BUZZER_PIN, freqs[userNote], 300);
      break;

    case TEST_PITCH_PLAY:
      targetNote = random(0, 5);
      display.clearDisplay(); display.println("PITCH TEST: Listen"); display.display();
      delay(800);
      tone(BUZZER_PIN, freqs[targetNote], 500);
      sendToTeacher(1, 0, 0, true);
      currentState = TEST_PITCH_ANSWER;
      break;

    case TEST_PITCH_ANSWER:
      userNote = map(potValue, 0, 1023, 0, 4);
      display.clearDisplay(); display.println("YOUR ANSWER?");
      display.setCursor(30, 30); display.setTextSize(2);
      display.println(noteNames[userNote]); display.display(); display.setTextSize(1);
      if (btn == 1) {
        p_total++;
        if (userNote == targetNote) p_correct++;
        currentState = TEST_PITCH_RESULT;
      }
      break;

    case TEST_PITCH_RESULT:
      { 
        int thisResult = (userNote == targetNote) ? 100 : 0;
        display.clearDisplay(); display.setCursor(0,0);
        display.println(userNote == targetNote ? "CORRECT!" : "WRONG!");
        display.print("Accuracy: "); display.print(thisResult); display.println("%");
        display.display();
        sendToTeacher(1, thisResult, 0, false);
        delay(3000); 
        currentState = MAIN_MENU; // 自动回主页面
      }
      break;

    case RHYTHM_SUB:
      subMenuIdx = map(potValue, 0, 1023, 0, 1);
      display.clearDisplay(); display.println("RHYTHM: SELECT");
      display.setCursor(0, 25);
      display.print(subMenuIdx == 0 ? "> " : "  "); display.println("Practice BPM");
      display.print(subMenuIdx == 1 ? "> " : "  "); display.println("Tap Test"); 
      display.display();
      if (btn == 1) currentState = (subMenuIdx == 0) ? PRACTICE_RHYTHM : TEST_RHYTHM_PLAY;
      break;

    case PRACTICE_RHYTHM:
      userBPM = map(potValue, 0, 1023, 40, 200);
      display.clearDisplay(); display.println("RHYTHM PRACTICE");
      display.setCursor(30, 30); display.setTextSize(2);
      display.print(userBPM); display.println(" BPM"); display.setTextSize(1);
      display.display();
      if (millis() - lastBeatTime >= (60000 / userBPM)) {
        tone(BUZZER_PIN, 1000, 50); lastBeatTime = millis();
      }
      break;

    case TEST_RHYTHM_PLAY:
      targetBPM = random(60, 160); targetInterval = 60000 / targetBPM; 
      display.clearDisplay(); display.println("RHYTHM TEST: Listen");
      display.setCursor(0, 30); display.print("Target: "); display.print(targetBPM); display.println(" BPM");
      display.display();
      delay(1000);
      for(int i=0; i<4; i++) { tone(BUZZER_PIN, 1200, 100); delay(targetInterval); }
      sendToTeacher(2, 0, 0, true);
      tapCount = 0; userIntervalSum = 0;
      currentState = TEST_RHYTHM_ANSWER;
      break;

    case TEST_RHYTHM_ANSWER:
      display.clearDisplay(); display.println("TAP 4 TIMES (D3)");
      display.setCursor(40, 30); display.setTextSize(2);
      display.print(tapCount); display.println("/4"); display.display(); display.setTextSize(1);
      if (digitalRead(BTN_SELECT_PIN) == LOW) {
        unsigned long currentTap = millis(); tone(BUZZER_PIN, 800, 50); 
        while(digitalRead(BTN_SELECT_PIN) == LOW); delay(30); 
        if (tapCount > 0) userIntervalSum += (currentTap - tapTime);
        tapTime = currentTap; tapCount++;
        if (tapCount >= 4) {
          userBPM = 60000 / (userIntervalSum / 3);
          int errorBPM = userBPM - targetBPM; 
          int rhythmAcc = (abs(errorBPM) <= 15) ? 100 : 0;
          display.clearDisplay(); display.setCursor(0,0);
          display.println("--- RESULT ---");
          display.print("Error: "); display.print(errorBPM); display.println(" BPM"); // 显示误差
          display.display();
          sendToTeacher(2, rhythmAcc, errorBPM, false);
          delay(4000); 
          currentState = MAIN_MENU; // 自动回主页面
        }
      }
      break;

    case INTEGRATED_SUB:
      display.clearDisplay(); display.println("INTEGRATED TEST");
      display.setCursor(0, 20); display.println("1. Listen\n2. Pitch\n3. Rhythm");
      display.display();
      if (btn == 1) currentState = TEST_INT_PLAY;
      break;

    case TEST_INT_PLAY:
      targetNote = random(0, 5); targetBPM = random(60, 150); targetInterval = 60000 / targetBPM;
      display.clearDisplay(); display.setTextSize(2); display.setCursor(20, 20);
      display.println("LISTEN!"); display.display(); display.setTextSize(1);
      delay(1000);
      for(int i=0; i<4; i++) { tone(BUZZER_PIN, freqs[targetNote], 150); delay(targetInterval); }
      sendToTeacher(3, 0, 0, true);
      pitchSelectStartTime = millis(); currentState = TEST_INT_PITCH;
      break;

    case TEST_INT_PITCH:
      { 
        userNote = map(potValue, 0, 1023, 0, 4);
        int timeLeft = 5 - ((millis() - pitchSelectStartTime) / 1000);
        display.clearDisplay(); display.println("SELECT PITCH:");
        display.setCursor(30, 20); display.setTextSize(2); display.println(noteNames[userNote]);
        display.setTextSize(1); display.setCursor(0, 50);
        display.print("Time: "); display.print(timeLeft); display.display();
        if (timeLeft <= 0 || digitalRead(BTN_SELECT_PIN) == LOW) {
          if (digitalRead(BTN_SELECT_PIN) == LOW) delay(300); 
          tapCount = 0; userIntervalSum = 0;
          currentState = TEST_INT_RHYTHM;
        }
      }
      break;

    case TEST_INT_RHYTHM:
      display.clearDisplay(); display.println("TAP RHYTHM (D3)");
      display.setCursor(40, 30); display.setTextSize(2);
      display.print(tapCount); display.println("/4"); display.display(); display.setTextSize(1);
      if (digitalRead(BTN_SELECT_PIN) == LOW) {
        unsigned long currentTap = millis(); tone(BUZZER_PIN, freqs[userNote], 50); 
        while(digitalRead(BTN_SELECT_PIN) == LOW); delay(30);
        if (tapCount > 0) userIntervalSum += (currentTap - tapTime);
        tapTime = currentTap; tapCount++;
        if (tapCount >= 4) {
          userBPM = 60000 / (userIntervalSum / 3);
          currentState = TEST_INT_RESULT;
        }
      }
      break;

    case TEST_INT_RESULT:
      { 
        int errorBPM = userBPM - targetBPM;
        int acc = 0;
        if (userNote == targetNote) acc += 50;
        if (abs(errorBPM) <= 15) acc += 50;
        display.clearDisplay(); display.setCursor(0,0);
        display.println("--- RESULT ---");
        display.print("Acc: "); display.print(acc); display.println("%");
        display.print("Err: "); display.print(errorBPM); display.println(" BPM"); // 显示误差
        display.display();
        sendToTeacher(3, acc, errorBPM, false);
        delay(4000); 
        currentState = MAIN_MENU; // 自动回主页面
      }
      break;
  }
}