#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[] = {4, 2, 5, 3}; // LEFT, UP, RIGHT, DOWN
const byte buttonPins2[] = {10, 8, 11, 9}; // LEFT, UP, RIGHT, DOWN
typedef enum {
  START,
  RUNNING,
  GAMEOVER
} State;

typedef enum {
    LEFT,
    UP,
    RIGHT,
    DOWN
} Direction;

#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];
int8_t snake2[MAX_SANKE_LENGTH][2];
uint8_t snake_length1,snake_length2;
Direction dir1,dir2;
Direction newDir1,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 i = 0; i < 4; i++) {
    pinMode(buttonPins2[i], INPUT_PULLUP);
  }  

  randomSeed(analogRead(A0));

  setupGame();
}

void setupGame() {
  gameState = START;
  dir1 = RIGHT;
  newDir1 = RIGHT;
  dir2 = RIGHT;
  newDir2 = RIGHT;  
  resetSnake1();
  resetSnake2(); 
  generateFruit();
  display.clearDisplay();
  drawMap();
  drawScore();
  drawPressToStart();
  display.display();
}

void resetSnake1() {
  snake_length1 = STARTING_SNAKE_SIZE;
  for(int i = 0; i < snake_length1; i++) {
    snake1[i][0] = MAP_SIZE_X / 2 - i;
    snake1[i][1] = MAP_SIZE_Y / 4;
  }
}

void resetSnake2() {
  snake_length2 = STARTING_SNAKE_SIZE;
  for(int i = 0; i < snake_length2; i++) {
    snake2[i][0] = MAP_SIZE_X / 2 - i;
    snake2[i][1] = MAP_SIZE_Y / 4 + MAP_SIZE_Y / 2;
  }
}

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

    case RUNNING:
      moveTime++;
      readDirection1();
      readDirection2();
      if(moveTime >= SNAKE_MOVE_DELAY) {
        dir1 = newDir1;
        dir2 = newDir2;
        display.clearDisplay();
        if(moveSnake1() || moveSnake2()) {
          gameState = GAMEOVER;
          drawGameover();
          delay(1000);
        }
        drawMap();
        drawScore();
        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 buttonPin = buttonPins1[i];
    if (digitalRead(buttonPin) == LOW) {
      return true;
    }
  }
  for (byte i = 0; i < 4; i++) {
    byte buttonPin = buttonPins2[i];
    if (digitalRead(buttonPin) == LOW) {
      return true;
    }
  }  
  return false;
}

void readDirection1() {
  for (byte i = 0; i < 4; i++) {
    byte buttonPin = buttonPins1[i];
    if (digitalRead(buttonPin) == LOW && i != ((int)dir1 + 2) % 4) {
      newDir1 = (Direction)i;
      return;
    }
  }
}
void readDirection2() {
  for (byte i = 0; i < 4; i++) {
    byte buttonPin = buttonPins2[i];
    if (digitalRead(buttonPin) == LOW && i != ((int)dir2 + 2) % 4) {
      newDir2 = (Direction)i;
      return;
    }
  }
}

bool moveSnake1() {
  int8_t x = snake1[0][0];
  int8_t y = snake1[0][1];

  switch(dir1) {
    case LEFT:
      x -= 1;
      break;
    case UP:
      y -= 1;
      break;
    case RIGHT:
      x += 1;
      break;
    case DOWN:
      y += 1;
      break;
  }

  if(collisionCheck(x, y))
    return true;

  for(int i = snake_length1 - 1; i > 0; i--) {
    snake1[i][0] = snake1[i - 1][0];
    snake1[i][1] = snake1[i - 1][1];
  }

  snake1[0][0] = x;
  snake1[0][1] = y;
  return false;
}

bool moveSnake2() {
  int8_t x = snake2[0][0];
  int8_t y = snake2[0][1];

  switch(dir2) {
    case LEFT:
      x -= 1;
      break;
    case UP:
      y -= 1;
      break;
    case RIGHT:
      x += 1;
      break;
    case DOWN:
      y += 1;
      break;
  }

  if(collisionCheck(x, y))
    return true;

  for(int i = snake_length2 - 1; i > 0; i--) {
    snake2[i][0] = snake2[i - 1][0];
    snake2[i][1] = snake2[i - 1][1];
  }

  snake2[0][0] = x;
  snake2[0][1] = y;
  return false;
}

void checkFruit() {
  if(fruit[0] == snake1[0][0] && fruit[1] == snake1[0][1])
  {
    if(snake_length1 + 1 <= MAX_SANKE_LENGTH)
      snake_length1++;
    generateFruit();
    return;
  }
  if(fruit[0] == snake2[0][0] && fruit[1] == snake2[0][1])
  {
    if(snake_length2 + 1 <= MAX_SANKE_LENGTH)
      snake_length2++;
    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 < snake_length1; i++) {
      if(fruit[0] == snake1[i][0] && fruit[1] == snake1[i][1]) {
        b = true;
        continue;
      }
    }
    for(int i = 0; i < snake_length2; i++) {
      if(fruit[0] == snake2[i][0] && fruit[1] == snake2[i][1]) {
        b = true;
        continue;
      }
    }    
  } while(b);
}

bool collisionCheck(int8_t x, int8_t y) {
  for(int i = 1; i < snake_length1; i++) {
    if(x == snake1[i][0] && y == snake1[i][1]) return true;
  }
  for(int i = 1; i < snake_length2; i++) {
    if(x == snake2[i][0] && y == snake2[i][1]) return true;
  }  
  if(x < 0 || y < 0 || x >= MAP_SIZE_X || y >= 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 < snake_length1; 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 i = 0; i < snake_length2; i++) {
    display.fillRect(snake2[i][0] * SNAKE_PIECE_SIZE + offsetMapX, snake2[i][1] * SNAKE_PIECE_SIZE + offsetMapY, SNAKE_PIECE_SIZE, SNAKE_PIECE_SIZE, SSD1306_WHITE);
  }  
}

void drawScore() {
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(2,2);
  display.print(F("Score1:"));
  display.println(snake_length1 - STARTING_SNAKE_SIZE);
  display.setCursor(2,10);
  display.print(F("Score2:"));
  display.println(snake_length2 - STARTING_SNAKE_SIZE);  
}

void drawPressToStart() {
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(2,30);
  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"));
}