#include <Wire.h>
#include <Keypad.h>
#include <LiquidCrystal_I2C.h>

#define LCD_ADDR 0x27
#define LCD_COLS 16
#define LCD_ROWS 2
LiquidCrystal_I2C lcd(LCD_ADDR, LCD_COLS, LCD_ROWS);

const byte ROWS = 4;
const byte COLS = 4;
char keys[ROWS][COLS] = {
  {'1','2','3','A'},
  {'4','5','6','B'},
  {'7','8','9','C'},
  {'*','0','#','D'}
};
byte rowPins[ROWS] = {12, 11, 10, 9};
byte colPins[COLS] = {8, 7, 6, 5};
Keypad keypad = Keypad(makeKeymap(keys), rowPins, colPins, ROWS, COLS);

#define BUZZER_PIN 2

// Пользовательские символы
byte LCD_CHAR_FLAG[8] = {B11111,B11110,B11111,B10000,B10000,B10000,B10000,B10000};
byte LCD_CHAR_HEART_0[8] = {B00000,B01010,B10101,B10001,B01110,B00100,B00000,B00000};
byte LCD_CHAR_HEART_1[8] = {B00000,B01010,B11111,B11111,B01110,B00100,B00000,B00000};
byte LCD_CHAR_CLOCK[8] = {B00000,B11111,B11111,B01110,B00100,B01010,B10001,B11111};

// Настройки по умолчанию
#define DEFAULT_TIMER_MAX 0
#define DEFAULT_SCORE_MAX 100
#define DEFAULT_LIFE_MAX 5
#define DEFAULT_N_MAX 10

enum GameState {
  START,
  GAME_A_1, GAME_A_2,
  GAME_B_1, GAME_B_2,
  GAME_C_1, GAME_C_2,
  GAME_D_1, GAME_D_2,
  SET_A, SET_B, SET_C, SET_D,
  STAT
};

struct GameSettings {
  int timerMax;
  int scoreMax;
  int lifeMax;
  int nMax;
};

GameState currentState = START;
GameSettings settings = {DEFAULT_TIMER_MAX, DEFAULT_SCORE_MAX, DEFAULT_LIFE_MAX, DEFAULT_N_MAX};

int score = 0;
int lives = DEFAULT_LIFE_MAX;
int currentTimerValue = 0;
unsigned long lastTimerUpdate = 0;
int correctStreak = 0;

char currentQuestion[16];
int correctAnswer;
int wrongAnswers[3];
char inputBuffer[5] = "";
byte inputPos = 0;
bool inputMode = false;

void setup() {
  lcd.init();
  lcd.backlight();
  
  lcd.createChar(0, LCD_CHAR_FLAG);
  lcd.createChar(1, LCD_CHAR_HEART_0);
  lcd.createChar(2, LCD_CHAR_HEART_1);
  lcd.createChar(3, LCD_CHAR_CLOCK);
  
  playStartGameSound();
  updateScreen();
}

void loop() {
  char key = keypad.getKey();
  if (key) {
    playCoinSound();
    handleInput(key);
  }

  updateTimer();
  delay(100);
}

// Добавленная функция для обработки настроек
void handleSettingsInput(char key) {
  if (inputMode) {
    if (key == '#') {
      if (inputPos > 0) {
        int value = atoi(inputBuffer);
        switch (currentState) {
          case SET_A: settings.timerMax = min(value, 60); break;
          case SET_B: settings.scoreMax = constrain(value, 10, 1000); break;
          case SET_C: settings.lifeMax = max(value, 1); break;
          case SET_D: settings.nMax = constrain(value, 3, 100); break;
        }
      }
      inputMode = false;
      clearInputBuffer();
    }
    else if (isdigit(key) && inputPos < 4) {
      inputBuffer[inputPos++] = key;
      inputBuffer[inputPos] = '\0';
    }
    else if (key == '*') {
      inputMode = false;
      clearInputBuffer();
    }
  } else {
    if (key == '#') {
      inputMode = true;
    }
    else if (key == '*') {
      currentState = START;
    }
    else if (key >= 'A' && key <= 'D') {
      switch (key) {
        case 'A': currentState = SET_A; break;
        case 'B': currentState = SET_B; break;
        case 'C': currentState = SET_C; break;
        case 'D': currentState = SET_D; break;
      }
    }
  }
}

void handleInput(char key) {
  switch (currentState) {
    case START:
      if (key == 'A') startGame(GAME_A_1);
      else if (key == 'B') startGame(GAME_B_1);
      else if (key == 'C') startGame(GAME_C_1);
      else if (key == 'D') startGame(GAME_D_1);
      else if (key == '*') currentState = SET_A;
      break;
      
    case GAME_A_1:
    case GAME_C_1:
      if (key == '*') currentState = START;
      else if (key == '#') currentState = STAT;
      else if (key >= '1' && key <= '4') checkAnswer(key - '1');
      break;
      
    case GAME_B_1:
    case GAME_D_1:
      if (key == '*') {
        if (inputPos > 0) inputBuffer[--inputPos] = '\0';
        else currentState = START;
      }
      else if (key == '#') {
        if (inputPos > 0) checkManualAnswer();
        else currentState = STAT;
      }
      else if (isdigit(key) && inputPos < 4) {
        inputBuffer[inputPos++] = key;
        inputBuffer[inputPos] = '\0';
      }
      break;
      
    case GAME_A_2:
    case GAME_B_2:
    case GAME_C_2:
    case GAME_D_2:
      if (key == '#') {
        if (lives > 0 && score < settings.scoreMax) {
          generateQuestion();
          resetTimer();
          // Возвращаемся к соответствующему экрану вопроса
          switch(currentState) {
            case GAME_A_2: currentState = GAME_A_1; break;
            case GAME_B_2: currentState = GAME_B_1; break;
            case GAME_C_2: currentState = GAME_C_1; break;
            case GAME_D_2: currentState = GAME_D_1; break;
            default: break;
          }
        } else {
          currentState = STAT;
        }
      }
      break;
      
    case STAT:
      if (key == '#') {
        if (lives > 0 && score < settings.scoreMax) {
          currentState = (currentState == GAME_A_2) ? GAME_A_1 :
                         (currentState == GAME_B_2) ? GAME_B_1 :
                         (currentState == GAME_C_2) ? GAME_C_1 : GAME_D_1;
        }
        else currentState = START;
      }
      else currentState = START;
      break;
  }
  updateScreen();
}

// Остальные функции остаются без изменений
void checkAnswer(byte index) {
  int answer = (index == 3) ? correctAnswer : wrongAnswers[index];
  processAnswer(answer == correctAnswer);
}

void checkManualAnswer() {
  int answer = atoi(inputBuffer);
  processAnswer(answer == correctAnswer);
  clearInputBuffer();
}

void processAnswer(bool isCorrect) {
  if (isCorrect) {
    score++;
    correctStreak++;
    if (correctStreak >= 3) {
      playBonusSound();
      correctStreak = 0;
    } else {
      playCorrectSound();
    }
  } else {
    playWrongSound();
    lives--;
    correctStreak = 0;
  }
  
  // Переходим на экран результата
  currentState = (GameState)((int)currentState + 1);
  
  if (lives <= 0 || score >= settings.scoreMax) {
    currentState = STAT;
  }
  updateScreen(); // Обновляем экран сразу после ответа
}

void startGame(GameState gameType) {
  currentState = gameType;
  score = 0;
  lives = settings.lifeMax;
  correctStreak = 0;
  generateQuestion();
  resetTimer();
  playStartGameSound();
}

void generateQuestion() {
  int a, b;
  char op;
  
  switch (currentState) {
    case GAME_A_1:
    case GAME_B_1:
      a = random(1, settings.nMax);
      b = random(1, settings.nMax);
      op = random(0, 2) ? '+' : '-';
      correctAnswer = (op == '+') ? a + b : a - b;
      sprintf(currentQuestion, "%d%c%d=", a, op, b);
      break;
      
    case GAME_C_1:
      a = random(1, settings.nMax);
      b = random(1, settings.nMax);
      correctAnswer = a * b;
      sprintf(currentQuestion, "%d*%d=", a, b);
      break;
      
    case GAME_D_1:
      b = random(1, settings.nMax);
      correctAnswer = random(1, settings.nMax);
      a = b * correctAnswer;
      sprintf(currentQuestion, "%d/%d=", a, b);
      break;
  }
  
  if (currentState == GAME_A_1 || currentState == GAME_C_1) {
    do {
      wrongAnswers[0] = correctAnswer + random(1, 3);
      wrongAnswers[1] = correctAnswer - random(1, 3);
      wrongAnswers[2] = correctAnswer + random(3, 5);
    } while (wrongAnswers[0] == wrongAnswers[1] || 
             wrongAnswers[0] == wrongAnswers[2] || 
             wrongAnswers[1] == wrongAnswers[2] ||
             wrongAnswers[0] == correctAnswer ||
             wrongAnswers[1] == correctAnswer ||
             wrongAnswers[2] == correctAnswer);
  }
}

void resetTimer() {
  currentTimerValue = settings.timerMax;
  lastTimerUpdate = millis();
}

void updateTimer() {
  if ((currentState == GAME_A_1 || currentState == GAME_B_1 || 
       currentState == GAME_C_1 || currentState == GAME_D_1) && 
      settings.timerMax > 0) {
    if (millis() - lastTimerUpdate >= 1000) {
      lastTimerUpdate = millis();
      if (--currentTimerValue <= 0) {
        playTimeOutSound();
        processAnswer(false);
        updateScreen();
      }
    }
  }
}

void clearInputBuffer() {
  inputPos = 0;
  memset(inputBuffer, 0, sizeof(inputBuffer));
}

void updateScreen() {
  lcd.clear();
  
  switch (currentState) {
    case START:
      lcd.print("   ARITHMETIK   ");
      lcd.setCursor(0, 1);
      lcd.print("*SET A B C D #");
      break;
      
    case GAME_A_1:
    case GAME_C_1:
      lcd.print(currentQuestion);
      lcd.print("***");
      displayTimer();
      
      lcd.setCursor(0, 1);
      lcd.print(wrongAnswers[0]);
      lcd.print(" ");
      lcd.print(wrongAnswers[1]);
      lcd.print(" ");
      lcd.print(wrongAnswers[2]);
      lcd.print(" ");
      lcd.print(correctAnswer);
      break;
      
    case GAME_B_1:
    case GAME_D_1:
      lcd.print(currentQuestion);
      lcd.print("***");
      displayTimer();
      
      lcd.setCursor(0, 1);
      lcd.print("*     ");
      lcd.print(inputBuffer);
      lcd.print("    #");
      break;
      
    case GAME_A_2:
    case GAME_B_2:
    case GAME_C_2:
    case GAME_D_2:
      lcd.print(currentQuestion);
      lcd.print(correctAnswer);
      
      lcd.setCursor(0, 1);
      displayScoreAndLives();
      lcd.print(" #-next");
      break;
      
    case SET_A:
    case SET_B:
    case SET_C:
    case SET_D:
      lcd.print("SET:");
      lcd.print(currentState == SET_A ? " >T" : "  T");
      lcd.print(currentState == SET_B ? " >F" : "  F");
      lcd.print(currentState == SET_C ? " >H" : "  H");
      lcd.print(currentState == SET_D ? " >N" : "  N");
      
      lcd.setCursor(0, 1);
      lcd.print("*      ");
      if (inputMode) lcd.print(inputBuffer);
      else {
        switch (currentState) {
          case SET_A: lcd.print(settings.timerMax); break;
          case SET_B: lcd.print(settings.scoreMax); break;
          case SET_C: lcd.print(settings.lifeMax); break;
          case SET_D: lcd.print(settings.nMax); break;
        }
      }
      lcd.print("      #");
      break;
      
    case STAT:
      if (lives <= 0) {
        lcd.print("     GAME OVER    ");
        playGameOverSound();
      } 
      else if (score >= settings.scoreMax) {
        lcd.print("     YOU WIN!     ");
        playMarioSound();
      } 
      else {
        lcd.print("      PAUSE       ");
      }
      
      lcd.setCursor(0, 1);
      displayScoreAndLives();
      break;
  }
}

void displayTimer() {
  lcd.setCursor(13, 0);
  if (settings.timerMax > 0) {
    if (currentTimerValue >= 10) lcd.print(currentTimerValue);
    else {
      lcd.print(" ");
      lcd.print(currentTimerValue);
    }
    lcd.write(3);
  } else {
    lcd.print("--");
    lcd.write(3);
  }
}

void displayScoreAndLives() {
  lcd.print(score);
  lcd.write(0);
  lcd.print(" ");
  for (int i = 0; i < settings.lifeMax; i++) {
    lcd.write(i < lives ? 2 : 1);
  }
}

// Звуковые функции остаются без изменений
void playCoinSound() {
  tone(BUZZER_PIN, 784, 100);
  delay(100);
  tone(BUZZER_PIN, 1046, 100);
  delay(100);
  noTone(BUZZER_PIN);
}

void playCorrectSound() {
  tone(BUZZER_PIN, 1046, 100);
  delay(100);
  tone(BUZZER_PIN, 1175, 100);
  delay(100);
  tone(BUZZER_PIN, 1318, 200);
  delay(200);
  noTone(BUZZER_PIN);
}

void playWrongSound() {
  for (int i = 0; i < 3; i++) {
    tone(BUZZER_PIN, 200, 200);
    delay(200);
    noTone(BUZZER_PIN);
    delay(100);
  }
}

void playStartGameSound() {
  tone(BUZZER_PIN, 523, 100);
  delay(100);
  tone(BUZZER_PIN, 659, 100);
  delay(100);
  tone(BUZZER_PIN, 784, 100);
  delay(100);
  noTone(BUZZER_PIN);
}

void playBonusSound() {
  tone(BUZZER_PIN, 1318, 150);
  delay(150);
  tone(BUZZER_PIN, 1568, 150);
  delay(150);
  tone(BUZZER_PIN, 1760, 150);
  delay(150);
  tone(BUZZER_PIN, 1975, 300);
  delay(300);
  noTone(BUZZER_PIN);
}

void playMarioSound() {
  const int m[] = {523,392,330,440,494,466,440,392,659,784,880,698,784,659,523,587,494,523,392,330,440,494,466,440,392,659,784,880,698,784,659,523,587,494,784,740,698,622,659,415,440,523,440,523,587,784,740,698,622,659,1047,1047,1047,1047,1047,1047};
  const byte d[] = {200,100,200,200,200,100,200,125,125,125,200,100,200,200,100,200,200,200,100,200,200,200,100,200,125,125,125,200,100,200,200,100,200,200,100,100,100,200,100,100,100,100,100,100,100,100,100,100,200,100,200,200,200,100,400};
  for (int i = 0; i < 56; i++) {
    tone(BUZZER_PIN, m[i], d[i]);
    delay(d[i] * 1.3);
    noTone(BUZZER_PIN);
  }
}

void playGameOverSound() {
  tone(BUZZER_PIN, 220, 400);
  delay(400);
  tone(BUZZER_PIN, 196, 400);
  delay(400);
  tone(BUZZER_PIN, 175, 400);
  delay(400);
  tone(BUZZER_PIN, 165, 800);
  delay(800);
  noTone(BUZZER_PIN);
}

void playTimeOutSound() {
  tone(BUZZER_PIN, 100, 500);
  delay(500);
  tone(BUZZER_PIN, 80, 500);
  delay(500);
  noTone(BUZZER_PIN);
}