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

#define SCREEN_WIDTH        128 // OLED display width, in pixels
#define SCREEN_HEIGHT        64 // OLED display height, in pixels

// Declaration for an SSD1306 display connected to I2C (SDA, SCL pins)
#define OLED_RESET            4 // Reset pin # (or -1 if sharing Arduino reset pin)
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);
const byte buttonPins1[] = {9, 8, 10, 11};
const byte buttonPins2[] = {4, 2, 5, 3};


typedef enum {
  START,
  RUNNING,
  GAMEOVER
} State;

typedef enum {
    LEFT1,
    UP1,
    RIGHT1,
    DOWN1
} Direction1;

typedef enum {
    LEFT2,
    UP2,
    RIGHT2,
    DOWN2
} Direction2;

#define SNAKE_PIECE_SIZE     3
#define MAX_SANKE_LENGTH    65
#define MAP_SIZE_X          20
#define MAP_SIZE_Y          20
#define STARTING_SNAKE_SIZE  5
#define SNAKE_MOVE_DELAY    30

State gameState;

int8_t snake1[MAX_SANKE_LENGTH][2];
uint8_t snake1_length;
int8_t snake2[MAX_SANKE_LENGTH][2];
uint8_t snake2_length;
Direction1 dir1;
Direction1 newDir1;
Direction2 dir2;
Direction2 newDir2;

int8_t fruit[2];

void setup() {
  Serial.begin(9600);
  if(!display.begin(SSD1306_SWITCHCAPVCC, 0x3D)) {
    Serial.println(F("SSD1306 allocation failed"));
    for(;;);
  }

  for (byte i = 0; i < 4; i++) {
    pinMode(buttonPins1[i], INPUT_PULLUP);
  }
  for (byte j = 0; j < 4; j++) {
    pinMode(buttonPins2[j], INPUT_PULLUP);
  }

  randomSeed(analogRead(A0));

  setupGame();
}

void setupGame() {
  gameState = START;
  dir1 = RIGHT1;
  newDir1 = RIGHT1;
  dir2 = RIGHT2;
  newDir2 = RIGHT2;
  resetSnake();
  generateFruit();
  display.clearDisplay();
  drawMap();
  drawScore1();
  drawScore2();
  drawPressToStart();
  display.display();
}

void resetSnake() {
  snake1_length = STARTING_SNAKE_SIZE;
  for(int i = 0; i < snake1_length; i++) {
    snake1[i][0] = MAP_SIZE_X / 2 - i;
    snake1[i][1] = MAP_SIZE_Y / 2;
  }
  snake2_length = STARTING_SNAKE_SIZE;
  for(int j = 0; j < snake2_length; j++) {
    snake2[j][0] = MAP_SIZE_X / 3 - j;
    snake2[j][1] = MAP_SIZE_Y / 3;
  }
}

int moveTime = 0;
void loop() {
  switch(gameState) {
    case START:
      if(buttonPress()) gameState = RUNNING;
      break;
    

    case RUNNING:
      moveTime++;
      readDirection();
      if(moveTime >= SNAKE_MOVE_DELAY) {
        dir1 = newDir1;
        dir2 = newDir2;
        display.clearDisplay();
        if(moveSnake()) { 
          gameState = GAMEOVER;
          drawGameover();
          delay(1000);
        }
        drawMap();
        drawScore1();
        drawScore2();
        display.display();
        checkFruit();
        moveTime = 0;
      }
      break;
    
    case GAMEOVER:
      if(buttonPress()) { 
        delay(500);
        setupGame();
        gameState = START;
      }
      break;
  }
  
  delay(10);
}

bool buttonPress() {
  for (byte i = 0; i < 4; i++) {
    byte buttonPin1 = buttonPins1[i];
    if (digitalRead(buttonPin1) == LOW) {
      return true;
    }
  }
  for (byte j = 0; j < 4; j++) {
    byte buttonPin2 = buttonPins2[j];
    if (digitalRead(buttonPin2) == LOW) {
      //drawTest();
      return true;
    }
  }
  return false;
}

void readDirection() {
  for (byte i = 0; i < 4; i++) {
    byte buttonPin1 = buttonPins1[i];
    if (digitalRead(buttonPin1) == LOW && i != ((int)dir1 + 2) %4) {
      newDir1 = (Direction1)i;
      return;
    }
  }
   for (byte j = 0; j < 4; j++) {
    byte buttonPin2 = buttonPins2[j];
    if (digitalRead(buttonPin2) == LOW && j != ((int)dir2 +2) %4) {
      newDir2 = (Direction2)j;
      return;
    }
  }
}

bool moveSnake() {
  int8_t x1 = snake1[0][0];
  int8_t y1 = snake1[0][1];
  int8_t x2 = snake2[0][0];
  int8_t y2 = snake2[0][1];

  switch(dir1) {
    case LEFT1:
      x1 -= 1;
      break;
    case UP1:
      y1 -= 1;
      break;
    case RIGHT1:
      x1 += 1;
      break;
    case DOWN1:
      y1 += 1;
      break;
  }
  switch(dir2) {
    case LEFT2:
      x2 -= 1;
      break;
    case UP2:
      y2 -= 1;
      break;
    case RIGHT2:
      x2 += 1;
      break;
    case DOWN2:
      y2 += 1;
      break;
  }

  if(collisionCheck(x1, y1, x2, y2))
    return true;

  for(int i = snake1_length - 1; i > 0; i--) {
    snake1[i][0] = snake1[i - 1][0];
    snake1[i][1] = snake1[i - 1][1];
  }
  for(int j = snake2_length - 1; j > 0; j--) {
    snake2[j][0] = snake2[j - 1][0];
    snake2[j][1] = snake2[j - 1][1];
  }
  for(int i = 1; i < snake1_length; i++) {
    if((snake2[0][0] == snake1[i][0] && snake2[0][1] == snake1[i][1])) { 
      gameState = GAMEOVER;
      drawGameover();
      delay(500);
    }
  }
  for(int i = 1; i < snake2_length; i++) {
    if(snake1[0][0] == snake2[i][0] && snake1[0][1] == snake2[i][1]) { 
      gameState = GAMEOVER;
      drawGameover();
      delay(500);
    }
  }

  snake1[0][0] = x1;
  snake1[0][1] = y1;
  snake2[0][0] = x2;
  snake2[0][1] = y2;
  return false;
}

void checkFruit() {
  if(fruit[0] == snake1[0][0] && fruit[1] == snake1[0][1])
  {
    if(snake1_length + 1 <= MAX_SANKE_LENGTH)
      snake1_length++;
    generateFruit();
  }
  if(fruit[0] == snake2[0][0] && fruit[1] == snake2[0][1])
  {
    if(snake2_length + 1 <= MAX_SANKE_LENGTH)
      snake2_length++;
    generateFruit();
  }
}

void generateFruit() {
  bool b = false;
  do {
    b = false;
    fruit[0] = random(0, MAP_SIZE_X);
    fruit[1] = random(0, MAP_SIZE_Y);
    for(int i = 0; i < snake1_length; i++) {
      if(fruit[0] == snake1[i][0] && fruit[1] == snake1[i][1]) {
        b = true;
        continue;
      }
    }
    for(int j = 0; j < snake2_length; j++) {
      if(fruit[0] == snake2[j][0] && fruit[1] == snake2[j][1]) {
        b = true;
        continue;
      }
    }
  } while(b);
}

bool collisionCheck(int8_t x1, int8_t y1, int8_t x2, int8_t y2) {
  for(int i = 1; i < snake1_length; i++) {
    if(x1 == snake1[i][0] && y1 == snake1[i][1]) return true;
  }
  for(int j = 1; j < snake2_length; j++) {
    if(x2 == snake2[j][0] && y2 == snake2[j][1]) return true;
  }
  if(x1 < 0 || y1 < 0 || x1 >= MAP_SIZE_X || y1 >= MAP_SIZE_Y || x2 < 0 || y2 < 0 || x2 >= MAP_SIZE_X || y2 >= MAP_SIZE_Y) return true;
  return false;
}

void drawMap() {
  int offsetMapX = SCREEN_WIDTH - SNAKE_PIECE_SIZE * MAP_SIZE_X - 2;
  int offsetMapY = 2;

  display.drawRect(fruit[0] * SNAKE_PIECE_SIZE + offsetMapX, fruit[1] * SNAKE_PIECE_SIZE + offsetMapY, SNAKE_PIECE_SIZE, SNAKE_PIECE_SIZE, SSD1306_INVERSE);
  display.drawRect(offsetMapX - 2, 0, SNAKE_PIECE_SIZE * MAP_SIZE_X + 4, SNAKE_PIECE_SIZE * MAP_SIZE_Y + 4, SSD1306_WHITE);
  for(int i = 0; i < snake1_length; i++) {
    display.fillRect(snake1[i][0] * SNAKE_PIECE_SIZE + offsetMapX, snake1[i][1] * SNAKE_PIECE_SIZE + offsetMapY, SNAKE_PIECE_SIZE, SNAKE_PIECE_SIZE, SSD1306_WHITE);
  }
   for(int j = 0; j < snake2_length; j++) {
    display.fillRect(snake2[j][0] * SNAKE_PIECE_SIZE + offsetMapX, snake2[j][1] * SNAKE_PIECE_SIZE + offsetMapY, SNAKE_PIECE_SIZE, SNAKE_PIECE_SIZE, SSD1306_WHITE);
  }
}

void drawScore1() {
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(1,1);
  display.print(F("Score1:"));
  display.println(snake1_length - STARTING_SNAKE_SIZE);
}

void drawScore2() {
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(1,10);
  display.print(F("Score2:"));
  display.println(snake2_length - STARTING_SNAKE_SIZE);
}

void drawPressToStart() {
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(2,20);
  display.print(F("Press a\n button to\n start the\n game!"));
}

void drawGameover() {
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(2,50);
  display.println(F("GAMEOVER"));
}

void drawTest() {
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(2,70);
  display.println(F("Test"));
}