import pygame
import random
import sys

# 初始化pygame
pygame.init()

# 游戏常量
WIDTH, HEIGHT = 600, 600
GRID_SIZE = 20
GRID_WIDTH = WIDTH // GRID_SIZE
GRID_HEIGHT = HEIGHT // GRID_SIZE
FPS = 10

# 颜色定义
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
GREEN = (0, 255, 0)
RED = (255, 0, 0)
BLUE = (0, 120, 255)
DARK_GREEN = (0, 180, 0)
GRAY = (40, 40, 40)

# 方向常量
UP = (0, -1)
DOWN = (0, 1)
LEFT = (-1, 0)
RIGHT = (1, 0)


class Snake:
    def __init__(self):
        self.reset()

    def reset(self):
        # 蛇初始位置在屏幕中央
        self.length = 3
        self.positions = [(GRID_WIDTH // 2, GRID_HEIGHT // 2)]
        self.direction = RIGHT
        self.next_direction = RIGHT
        self.score = 0
        self.grow_pending = 2  # 初始长度为3，需要再生长2节

    def get_head_position(self):
        return self.positions[0]

    def turn(self, direction):
        # 防止直接反向移动
        if (direction[0] * -1, direction[1] * -1) != self.direction:
            self.next_direction = direction

    def move(self):
        self.direction = self.next_direction
        head = self.get_head_position()
        new_x = (head[0] + self.direction[0]) % GRID_WIDTH
        new_y = (head[1] + self.direction[1]) % GRID_HEIGHT
        new_position = (new_x, new_y)

        # 检查是否撞到自己
        if new_position in self.positions[1:]:
            return False  # 游戏结束

        # 添加新的头部
        self.positions.insert(0, new_position)

        # 如果有生长待处理，不删除尾部
        if self.grow_pending > 0:
            self.grow_pending -= 1
            self.length += 1
        else:
            # 否则删除尾部，保持长度不变
            self.positions.pop()

        return True  # 游戏继续

    def grow(self):
        self.grow_pending += 1
        self.score += 10

    def draw(self, surface):
        for i, p in enumerate(self.positions):
            # 绘制蛇身矩形
            rect = pygame.Rect(p[0] * GRID_SIZE, p[1] *
                               GRID_SIZE, GRID_SIZE, GRID_SIZE)

            # 蛇头用不同颜色
            if i == 0:
                pygame.draw.rect(surface, GREEN, rect)
                pygame.draw.rect(surface, DARK_GREEN, rect, 1)

                # 绘制蛇眼睛
                eye_size = GRID_SIZE // 5
                # 根据方向确定眼睛位置
                if self.direction == RIGHT:
                    left_eye = (rect.right - eye_size, rect.top + eye_size*2)
                    right_eye = (rect.right - eye_size,
                                 rect.bottom - eye_size*2)
                elif self.direction == LEFT:
                    left_eye = (rect.left + eye_size, rect.top + eye_size*2)
                    right_eye = (rect.left + eye_size,
                                 rect.bottom - eye_size*2)
                elif self.direction == UP:
                    left_eye = (rect.left + eye_size*2, rect.top + eye_size)
                    right_eye = (rect.right - eye_size*2, rect.top + eye_size)
                else:  # DOWN
                    left_eye = (rect.left + eye_size*2, rect.bottom - eye_size)
                    right_eye = (rect.right - eye_size*2,
                                 rect.bottom - eye_size)

                pygame.draw.circle(surface, BLACK, left_eye, eye_size)
                pygame.draw.circle(surface, BLACK, right_eye, eye_size)
            else:
                # 蛇身
                pygame.draw.rect(surface, DARK_GREEN, rect)
                pygame.draw.rect(surface, GREEN, rect, 1)


class Food:
    def __init__(self):
        self.position = (0, 0)
        self.randomize_position()

    def randomize_position(self):
        self.position = (random.randint(0, GRID_WIDTH - 1),
                         random.randint(0, GRID_HEIGHT - 1))

    def draw(self, surface):
        rect = pygame.Rect(self.position[0] * GRID_SIZE,
                           self.position[1] * GRID_SIZE,
                           GRID_SIZE, GRID_SIZE)
        pygame.draw.rect(surface, RED, rect)
        pygame.draw.rect(surface, (255, 200, 200), rect, 2)

        # 绘制食物内部细节
        inner_rect = pygame.Rect(
            self.position[0] * GRID_SIZE + GRID_SIZE//4,
            self.position[1] * GRID_SIZE + GRID_SIZE//4,
            GRID_SIZE//2, GRID_SIZE//2
        )
        pygame.draw.rect(surface, (255, 150, 150), inner_rect, 1)


class Game:
    def __init__(self):
        self.screen = pygame.display.set_mode((WIDTH, HEIGHT))
        pygame.display.set_caption("贪吃蛇游戏")
        self.clock = pygame.time.Clock()
        self.font = pygame.font.SysFont(None, 36)
        self.small_font = pygame.font.SysFont(None, 24)

        self.snake = Snake()
        self.food = Food()
        self.game_over = False
        self.game_started = False

    def handle_keys(self):
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
            elif event.type == pygame.KEYDOWN:
                if self.game_over:
                    if event.key == pygame.K_SPACE:
                        self.snake.reset()
                        self.food.randomize_position()
                        self.game_over = False
                elif not self.game_started:
                    if event.key == pygame.K_SPACE:
                        self.game_started = True
                else:
                    if event.key == pygame.K_UP:
                        self.snake.turn(UP)
                    elif event.key == pygame.K_DOWN:
                        self.snake.turn(DOWN)
                    elif event.key == pygame.K_LEFT:
                        self.snake.turn(LEFT)
                    elif event.key == pygame.K_RIGHT:
                        self.snake.turn(RIGHT)

    def update(self):
        if not self.game_started or self.game_over:
            return

        # 移动蛇
        if not self.snake.move():
            self.game_over = True
            return

        # 检查是否吃到食物
        if self.snake.get_head_position() == self.food.position:
            self.snake.grow()
            self.food.randomize_position()

            # 确保食物不出现在蛇身上
            while self.food.position in self.snake.positions:
                self.food.randomize_position()

    def draw_grid(self):
        for x in range(0, WIDTH, GRID_SIZE):
            pygame.draw.line(self.screen, GRAY, (x, 0), (x, HEIGHT), 1)
        for y in range(0, HEIGHT, GRID_SIZE):
            pygame.draw.line(self.screen, GRAY, (0, y), (WIDTH, y), 1)

    def draw_score(self):
        score_text = self.font.render(f"得分: {self.snake.score}", True, WHITE)
        self.screen.blit(score_text, (10, 10))

        length_text = self.small_font.render(
            f"长度: {self.snake.length}", True, WHITE)
        self.screen.blit(length_text, (10, 50))

    def draw_instructions(self):
        if not self.game_started:
            title = self.font.render("贪吃蛇游戏", True, GREEN)
            start_text = self.font.render("按空格键开始游戏", True, WHITE)

            instructions = [
                "使用方向键控制蛇的移动",
                "吃到红色食物可以增加长度和得分",
                "不要撞到自己的身体",
                "按空格键重新开始游戏"
            ]

            self.screen.blit(
                title, (WIDTH//2 - title.get_width()//2, HEIGHT//2 - 100))
            self.screen.blit(
                start_text, (WIDTH//2 - start_text.get_width()//2, HEIGHT//2 - 40))

            for i, line in enumerate(instructions):
                text = self.small_font.render(line, True, WHITE)
                self.screen.blit(
                    text, (WIDTH//2 - text.get_width()//2, HEIGHT//2 + 20 + i*30))

        elif self.game_over:
            game_over_text = self.font.render("游戏结束!", True, RED)
            score_text = self.font.render(
                f"最终得分: {self.snake.score}", True, WHITE)
            restart_text = self.font.render("按空格键重新开始", True, GREEN)

            self.screen.blit(game_over_text, (WIDTH//2 -
                             game_over_text.get_width()//2, HEIGHT//2 - 50))
            self.screen.blit(
                score_text, (WIDTH//2 - score_text.get_width()//2, HEIGHT//2))
            self.screen.blit(
                restart_text, (WIDTH//2 - restart_text.get_width()//2, HEIGHT//2 + 50))

    def run(self):
        while True:
            self.handle_keys()
            self.update()

            # 绘制背景
            self.screen.fill(BLACK)

            # 绘制网格
            self.draw_grid()

            # 绘制游戏元素
            self.snake.draw(self.screen)
            self.food.draw(self.screen)

            # 绘制分数
            self.draw_score()

            # 绘制游戏状态提示
            self.draw_instructions()

            # 显示控制提示
            if self.game_started and not self.game_over:
                controls = self.small_font.render("方向键控制移动", True, WHITE)
                self.screen.blit(
                    controls, (WIDTH - controls.get_width() - 10, 10))

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


if __name__ == "__main__":
    game = Game()
    game.run()