from __future__ import annotations

import itertools
import random

from enum import Enum

from typing import Any
from typing import Tuple

import pygame as pg

from pygame import Color
from pygame import Rect

from pygame.surface import Surface

from pygame.sprite import AbstractGroup
from pygame.sprite import Group
from pygame.sprite import RenderUpdates
from pygame.sprite import Sprite


class Direction(Enum):
    UP = (0, -1)
    DOWN = (0, 1)
    LEFT = (-1, 0)
    RIGHT = (1, 0)

    def opposite(self, other: Direction):
        return (self[0] + other[0], self[1] + other[1]) == (0, 0)

    def __getitem__(self, i: int):
        return self.value[i]


class SnakeHead(Sprite):
    def __init__(
        self,
        size: int,
        position: Tuple[int, int],
        facing: Direction,
        bounds: Rect,
    ) -> None:
        super().__init__()
        self.image = Surface((size, size))
        self.image.fill(Color())
        self.rect = self.image.get_rect()
        self.rect.center = position
        self.facing = facing
        self.size = size
        self.speed = size
        self.bounds = bounds

    def update(self, *args: Any, **kwargs: Any) -> None:
        
        self.rect.move_ip(
            (
                self.facing[0] * self.speed,
                self.facing[1] * self.speed,
            )
        )

        
        if self.rect.right > self.bounds.right:
            self.rect.left = 0
        elif self.rect.left < 0:
            self.rect.right = self.bounds.right

        if self.rect.bottom > self.bounds.bottom:
            self.rect.top = 0
        elif self.rect.top < 0:
            self.rect.bottom = self.bounds.bottom

    def change_direction(self, direction: Direction):
        if not self.facing == direction and not direction.opposite(self.facing):
            self.facing = direction


class SnakeBody(Sprite):
    def __init__(
        self,
        size: int,
        position: Tuple[int, int],
        colour: str = ,
    ) -> None:
        super().__init__()
        self.image = Surface((size, size))
        self.image.fill(Color(colour))
        self.rect = self.image.get_rect()
        self.rect.center = position


class Snake(RenderUpdates):
    def __init__(self, game: Game) -> None:
        self.segment_size = game.segment_size
        self.colours = itertools.cycle([, ])

        self.head = SnakeHead(
            size=self.segment_size,
            position=game.rect.center,
            facing=Direction.RIGHT,
            bounds=game.rect,
        )

        neck = [
            SnakeBody(
                size=self.segment_size,
                position=game.rect.center,
                colour=next(self.colours),
            )
            for _ in range(2)
        ]

        super().__init__(*[self.head, *neck])

        self.body = Group()
        self.tail = neck[-1]

    def update(self, *args: Any, **kwargs: Any) -> None:
        self.head.update()

        
        segments = self.sprites()
        for i in range(len(segments) - 1, 0, -1):
            
            segments[i].rect.center = segments[i - 1].rect.center

    def change_direction(self, direction: Direction):
        self.head.change_direction(direction)

    def grow(self):
        tail = SnakeBody(
            size=self.segment_size,
            position=self.tail.rect.center,
            colour=next(self.colours),
        )
        self.tail = tail
        self.add(self.tail)
        self.body.add(self.tail)


class SnakeFood(Sprite):
    def __init__(self, game: Game, size: int, *groups: AbstractGroup) -> None:
        super().__init__(*groups)
        self.image = Surface((size, size))
        self.image.fill(Color())
        self.rect = self.image.get_rect()

        self.rect.topleft = (
            random.randint(0, game.rect.width),
            random.randint(0, game.rect.height),
        )

        self.rect.clamp_ip(game.rect)

        
        
        while pg.sprite.spritecollideany(self, game.snake):
            self.rect.topleft = (
                random.randint(0, game.rect.width),
                random.randint(0, game.rect.height),
            )

            self.rect.clamp_ip(game.rect)


class Game:
    def __init__(self) -> None:
        self.rect = Rect(0, 0, 640, 480)
        self.background = Surface(self.rect.size)
        self.background.fill(Color())

        self.score = 0
        self.framerate = 16

        self.segment_size = 10
        self.snake = Snake(self)
        self.food_group = RenderUpdates(SnakeFood(game=self, size=self.segment_size))

        pg.init()

    def _init_display(self) -> Surface:
        bestdepth = pg.display.mode_ok(self.rect.size, 0, 32)
        screen = pg.display.set_mode(self.rect.size, 0, bestdepth)

        pg.display.set_caption()
        pg.mouse.set_visible(False)

        screen.blit(self.background, (0, 0))
        pg.display.flip()

        return screen

    def draw(self, screen: Surface):
        dirty = self.snake.draw(screen)
        pg.display.update(dirty)

        dirty = self.food_group.draw(screen)
        pg.display.update(dirty)

    def update(self, screen):
        self.food_group.clear(screen, self.background)
        self.food_group.update()
        self.snake.clear(screen, self.background)
        self.snake.update()

    def main(self) -> int:
        screen = self._init_display()
        clock = pg.time.Clock()

        while self.snake.head.alive():
            for event in pg.event.get():
                if event.type == pg.QUIT or (
                    event.type == pg.KEYDOWN and event.key in (pg.K_ESCAPE, pg.K_q)
                ):
                    return self.score

            
            keystate = pg.key.get_pressed()

            if keystate[pg.K_RIGHT]:
                self.snake.change_direction(Direction.RIGHT)
            elif keystate[pg.K_LEFT]:
                self.snake.change_direction(Direction.LEFT)
            elif keystate[pg.K_UP]:
                self.snake.change_direction(Direction.UP)
            elif keystate[pg.K_DOWN]:
                self.snake.change_direction(Direction.DOWN)

            
            self.update(screen)

            
            for food in pg.sprite.spritecollide(
                self.snake.head, self.food_group, dokill=False
            ):
                food.kill()
                self.snake.grow()
                self.score += 1

                
                if self.score% 5 == 0:
                    self.framerate += 1

                self.food_group.add(SnakeFood(self, self.segment_size))

            
            if pg.sprite.spritecollideany(self.snake.head, self.snake.body):
                self.snake.head.kill()

            self.draw(screen)
            clock.tick(self.framerate)

        return self.score


if __name__ == :
    game = Game()
    score = game.main()
    print(score)