#include <FastLED.h>

// LED matrix and joystick setup
#define LED_TYPE    WS2812B
#define NUM_LEDS    64
#define DATA_PIN    3
#define VERT_PIN    A0
#define HORZ_PIN    A1
#define SEL_PIN     2

CRGB leds[NUM_LEDS];

void controlLED(int y, int x, CRGB color) {

  if (x < 0 || x > 7 || y < 0 || y > 7) return;

  int ledIndex = y * 8 + x;

  leds[ledIndex] = color;
}


class Head {
  public:
    int x = 3;
    int y = 3;

    void move(int dx, int dy) {
      x += dx;
      y += dy;
    }

    void draw() {
      controlLED(x, y, CRGB::Green);
    }
};

class Body {
  private:
    struct Segment { int x, y; };
    Segment segments[64];
    int length = 1;

  public:
    void addSegment(int x, int y, int count) {
      // Add segments
      for (int c = 0; c < count; c++) {
        for (int i = length; i > 0; i--) segments[i] = segments[i - 1];
        segments[0] = {x, y};
        length++;
        if (length >= 64) break; 
      }
    }

    void removeSegment(int count) {
  
      length -= count;
      if (length < 1) length = 1;
    }

    void move(int newX, int newY) {
      for (int i = length - 1; i > 0; i--) segments[i] = segments[i - 1];
      segments[0] = {newX, newY};
    }

    void draw() {
      for (int i = 0; i < length; i++) controlLED(segments[i].y, segments[i].x, CRGB::Green);
    }

    bool checkCollision(int x, int y) {
      for (int i = 0; i < length; i++) if (segments[i].x == x && segments[i].y == y) return true;
      return false;
    }

    int getLength() { return length; }
};

class Apple {
  public:
    int x = random(0, 8);
    int y = random(0, 8);

    void spawn() {
      x = random(0, 8);
      y = random(0, 8);
    }

    void draw() {
      controlLED(y, x, CRGB::Red);
    }
};

class Game {
  private:
    Head head;
    Body body;
    Apple apple;
    int dx = 1, dy = 0;

  public:
  void readJoystick() {
  int vert = analogRead(VERT_PIN);
  int horz = analogRead(HORZ_PIN);

  if (vert < 400) { dx = -1; dy = 0; } // Down
  if (vert > 600) { dx = 1; dy = 0; }  // Up
  if (horz < 400) { dx = 0; dy = 1; } // Left
  if (horz > 600) { dx = 0; dy = -1; }  // Right
}


void update() {
      int newX = head.x + dx, newY = head.y + dy;
      if (newX < 0 || newX > 7 || newY < 0 || newY > 7 || body.checkCollision(newX, newY)) { gameOver(); return; }

      if (newX == apple.x && newY == apple.y) {
        int randomChange = random(-2, 3); 
        if (randomChange > 0) {
          body.addSegment(head.x, head.y, randomChange); 
        } else if (randomChange < 0) {
          body.removeSegment(-randomChange); 
        }
        apple.spawn(); 
      } else {
        body.move(head.x, head.y);
      }

      head.x = newX; head.y = newY;
    }

    void draw() {
      FastLED.clear(); // Clear all LEDs
      controlLED(apple.y, apple.x, CRGB::Red);
      body.draw();
      controlLED(head.y, head.x, CRGB::Green);
      FastLED.show();
    }

    void reset(){
      head=Head();
      body=Body();
      apple=Apple();
      dx=1;
      dy=0;
      apple.spawn();
    }

    void gameOver() {
      
      for (int i = 0; i < 3; i++) {
        FastLED.showColor(CRGB::Red);
        delay(500);
        FastLED.clear();
        delay(500);
      }
      reset();
    }

    void run() {
      readJoystick();
      update();
      draw();
    }
};

Game snakeGame;

void setup() {
  FastLED.addLeds<LED_TYPE, DATA_PIN, GRB>(leds, NUM_LEDS);
  pinMode(VERT_PIN, INPUT);
  pinMode(HORZ_PIN, INPUT);
  pinMode(SEL_PIN, INPUT_PULLUP);
  Serial.begin(9600);
}

void loop() {
  snakeGame.run();
  delay(200); // Game speed
}