from typing import List
import pygame
from random import randint


class Pos:
    def __init__(self, x, y):
        self.x, self.y = int(x), int(y)

    @classmethod
    def random(cls):
        return cls(randint(0, Game.MAP.x - 1), randint(0, Game.MAP.y - 1))

    def copy(self):
        return Pos(self.x, self.y)

    def __eq__(self, other):
        return self.x == other.x and self.y == other.y


class Colours:
    BACKGROUND = (0, 0, 0)
    FOREGROUND = (255, 0, 0)
    SNAKE = (0, 255, 0)
    FOOD = (0, 0, 255)


class Game:
    MAP = Pos(20, 20)
    CUBE = 20
    FPS = 60
    SPEED = 120
    DEBUG = False

    def __init__(self):
        start_pos = Pos(self.MAP.x / 2, self.MAP.y / 2)
        self.snake = Snake(start_pos)
        self.dead = False
        self.running = True
        self.score = 0

        self.screen = pygame.display.set_mode((self.MAP.y * self.CUBE, self.MAP.x * self.CUBE))
        self.clock = pygame.time.Clock()
        self.timer = 0.0
        self.font = pygame.font.SysFont('Arial', 20, True, False)

    def proc_events(self, e):
        if e.type == pygame.QUIT:
            self.running = False
        if e.type != pygame.KEYDOWN:
            return

        if e.key == pygame.K_SPACE and self.dead:
            start_pos = Pos(self.MAP.x / 2, self.MAP.y / 2)
            self.snake = Snake(start_pos)
            self.dead = False

        if e.key == pygame.K_UP:
            if self.snake.dir != 1:
                self.snake.dir_next = 0

        elif e.key == pygame.K_DOWN:
            if self.snake.dir != 0:
                self.snake.dir_next = 1

        elif e.key == pygame.K_LEFT:
            if self.snake.dir != 3:
                self.snake.dir_next = 2

        elif e.key == pygame.K_RIGHT:
            if self.snake.dir != 2:
                self.snake.dir_next = 3

        elif e.key == pygame.K_d:
            self.DEBUG = not self.DEBUG

        elif e.key == pygame.K_ESCAPE:
            self.running = False

    def loop(self):
        for event in pygame.event.get():
            self.proc_events(event)

        if self.dead:
            text = self.font.render(f"SCORE: {self.score}", True, Colours.FOREGROUND)
            self.screen.blit(text, (self.MAP.x * self.CUBE / 2, self.MAP.y * self.CUBE / 2))
            pygame.display.flip()
            self.clock.tick(30)
            return

        if 0 > self.snake.head.x or self.snake.head.x >= self.MAP.x or \
            0 > self.snake.head.y or self.snake.head.y >= self.MAP.y:
            self.dead = True

        self.screen.fill(Colours.BACKGROUND)
        pygame.draw.rect(
            self.screen, Colours.FOOD,
            (self.snake.food.y * self.CUBE, self.snake.food.x * self.CUBE, self.CUBE - 1, self.CUBE - 1), 0)

        for c in self.snake.body:
            pygame.draw.rect(
                self.screen, Colours.SNAKE,
                (c.y * self.CUBE, c.x * self.CUBE, self.CUBE - 1, self.CUBE - 1), 0)

        self.timer += self.clock.get_time()
        if self.timer >= self.SPEED:
            self.timer = 0.0
            if self.snake.dir_next == 0:
                self.snake.head.x -= 1
            elif self.snake.dir_next == 1:
                self.snake.head.x += 1
            elif self.snake.dir_next == 2:
                self.snake.head.y -= 1
            else:
                self.snake.head.y += 1
            self.snake.dir = self.snake.dir_next

            if self.snake.collide(self.snake.head):
                self.dead = True

            if self.snake.head == self.snake.food:
                self.snake.add_food()
                self.snake.size += 1
                self.score += 1

            self.snake.move()

        if self.DEBUG:
            text1 = self.font.render(
                f"X: {self.snake.head.x:02d} Y: {self.snake.head.y:02d} "
                f"S: {self.snake.size} FPS: {self.clock.get_fps():.1f}",
                True, Colours.FOREGROUND)
            text2 = self.font.render(
                f"F X: {self.snake.food.x:02d} Y: {self.snake.food.y:02d} S: {self.score}",
                True, Colours.FOREGROUND)
            self.screen.blit(text1, (10, 10))
            self.screen.blit(text2, (10, 32))

        pygame.display.flip()
        self.clock.tick(60)


class Snake:
    def __init__(self, start: Pos, size=3, dir=0):
        self.size: int = size
        self.head: Pos = start
        self.body: List[Pos] = [start.copy()]
        self.dir: int = dir  # Direction 0=N, 1=S, 2=W, 3=E
        self.dir_next: int = dir  # Next direction
        self.food: Pos = None
        self.add_food()

    def add_food(self):
        self.food = Pos.random()
        while self.collide(self.food):
            self.food = Pos.random()

    def collide(self, pos: Pos) -> bool:
        for l in self.body:
            if pos == l:
                return True
        return False

    def move(self):
        self.body.append(self.head.copy())
        if len(self.body) > self.size:
            self.body.pop(0)


if __name__ == '__main__':
    pygame.init()
    pygame.display.set_caption("Snake")
    game = Game()
    game.DEBUG = True
    while game.running:
        game.loop()
    pygame.quit()