import pygame
import random

# 初始化
pygame.init()

# 常量定义
SCREEN_WIDTH = 1280
SCREEN_HEIGHT = 960
GRID_SIZE = 20  # 网格大小
GRID_WIDTH = SCREEN_WIDTH // GRID_SIZE
GRID_HEIGHT = SCREEN_HEIGHT // GRID_SIZE

# 颜色定义
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
GREEN = (0, 255, 0)
RED = (255, 0, 0)

# 屏幕设置
screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
pygame.display.set_caption("Snake Game")

class Snake(pygame.sprite.Sprite):
    def __init__(self):
        super().__init__()
        self.image = pygame.Surface((GRID_SIZE, GRID_SIZE))
        self.image.fill(GREEN)
        self.rect = self.image.get_rect(center=(SCREEN_WIDTH // 2+10, SCREEN_HEIGHT // 2+10))
        self.direction = "right"
        self.body_segments = []
        self.growth_amount = 1
        self.move_counter = 0
        self.move_interval = 0  # 每5帧移动一次
        self.get_target = 1
        self.target_posx = 0
        self.target_posy = 0

    def update(self): 
        self.move_counter += 1
        if self.move_counter < self.move_interval:
            return

        self.move_counter = 0
        old_head_pos = self.rect.copy()

        # 获取食物位置
        if self.get_target == 1:
            self.target_posx, self.target_posy = locate_food()
            self.get_target = 0

        # 自动路径规划
        current_col = self.rect.x // GRID_SIZE
        current_row = self.rect.y // GRID_SIZE
        target_col = self.target_posx
        target_row = self.target_posy

        # 获取安全方向
        safe_directions = self.get_safe_directions()
        if not safe_directions:
            game_over()
            return

        # 选择最优方向
        best_dir = self.choose_best_direction(safe_directions, current_col, current_row, target_col, target_row)
        self.direction = best_dir

        # 移动头部
        self.move_head()

        # 更新身体段位置
        self.update_body_segments(old_head_pos)

        # 碰撞检测
        self.check_collisions()

    def get_safe_directions(self):
        safe_directions = []
        for direction in ["left", "right", "up", "down"]:
            if self.is_safe(direction):
                safe_directions.append(direction)
        return safe_directions

    def is_safe(self, direction):
        new_x, new_y = self.rect.x, self.rect.y
        if direction == "left":
            new_x -= GRID_SIZE
        elif direction == "right":
            new_x += GRID_SIZE
        elif direction == "up":
            new_y -= GRID_SIZE
        elif direction == "down":
            new_y += GRID_SIZE

        # 边界检测
        if new_x < 0 or new_x >= SCREEN_WIDTH or new_y < 0 or new_y >= SCREEN_HEIGHT:
            return False

        # 身体碰撞检测
        new_rect = pygame.Rect(new_x, new_y, GRID_SIZE, GRID_SIZE)
        for segment in self.body_segments:
            if new_rect.colliderect(segment.rect):
                return False
        return True

    def choose_best_direction(self, safe_directions, current_col, current_row, target_col, target_row):
        best_dir = safe_directions[0]
        min_distance = float('inf')
        
        for direction in safe_directions:
            temp_col, temp_row = current_col, current_row
            if direction == "left":
                temp_col -= 1
            elif direction == "right":
                temp_col += 1
            elif direction == "up":
                temp_row -= 1
            elif direction == "down":
                temp_row += 1

            distance = abs(temp_col - target_col) + abs(temp_row - target_row)
            if distance < min_distance:
                min_distance = distance
                best_dir = direction
        return best_dir

    def move_head(self):
        if self.direction == "left":
            self.rect.x -= GRID_SIZE
        elif self.direction == "right":
            self.rect.x += GRID_SIZE
        elif self.direction == "up":
            self.rect.y -= GRID_SIZE
        elif self.direction == "down":
            self.rect.y += GRID_SIZE

    def update_body_segments(self, old_head_pos):
        if self.body_segments:
            old_positions = [old_head_pos] + [seg.rect.copy() for seg in self.body_segments[:-1]]
            for i, seg in enumerate(self.body_segments):
                seg.rect = old_positions[i]

    def check_collisions(self):
        # 墙体验证
        if (self.rect.left < 0 or self.rect.right > SCREEN_WIDTH or
            self.rect.top < 0 or self.rect.bottom > SCREEN_HEIGHT):
            game_over()

        # 身体碰撞验证
        for segment in self.body_segments:
            if self.rect.colliderect(segment.rect):
                game_over()

    def add_segment(self):
        for _ in range(self.growth_amount):
            new_segment = pygame.sprite.Sprite()
            new_segment.image = pygame.Surface((GRID_SIZE, GRID_SIZE))
            new_segment.image.fill((0, 200, 0))
            
            if self.body_segments:
                base_pos = self.body_segments[-1].rect.copy()
            else:
                base_pos = self.rect.copy()
            
            offset = GRID_SIZE
            if self.direction == "left":
                base_pos.x += offset
            elif self.direction == "right":
                base_pos.x -= offset
            elif self.direction == "up":
                base_pos.y += offset
            elif self.direction == "down":
                base_pos.y -= offset
            
            new_segment.rect = base_pos
            self.body_segments.append(new_segment)
            all_sprites.add(new_segment)

class Food(pygame.sprite.Sprite):
    def __init__(self):
        super().__init__()
        self.image = pygame.Surface((GRID_SIZE, GRID_SIZE))
        self.image.fill(RED)
        self.rect = self.image.get_rect(center=(
            random.randint(0, GRID_WIDTH - 1) * GRID_SIZE + GRID_SIZE // 2,
            random.randint(0, GRID_HEIGHT - 1) * GRID_SIZE + GRID_SIZE // 2
        ))

    def update(self):
        if pygame.sprite.collide_rect(snake, self):
            self.rect.center = (
                random.randint(0, GRID_WIDTH - 1) * GRID_SIZE + GRID_SIZE // 2,
                random.randint(0, GRID_HEIGHT - 1) * GRID_SIZE + GRID_SIZE // 2
            )
            snake.add_segment()
            snake.get_target = 1  # 吃到食物后重新设置目标

def draw_grid():
    for x in range(0, SCREEN_WIDTH, GRID_SIZE):
        pygame.draw.line(screen, WHITE, (x, 0), (x, SCREEN_HEIGHT))
    for y in range(0, SCREEN_HEIGHT, GRID_SIZE):
        pygame.draw.line(screen, WHITE, (0, y), (SCREEN_WIDTH, y))

def game_over():
    global running
    running = False
    print("Game Over! Your score:", len(snake.body_segments))

def locate_food():
    food = random.choice(food_group.sprites())
    food_position = food.rect.center
    food_col = food_position[0] // GRID_SIZE  # 修正行列计算
    food_row = food_position[1] // GRID_SIZE
    return food_col, food_row  # 返回（列，行）

# 创建蛇对象
snake = Snake()

# 创建精灵组
all_sprites = pygame.sprite.Group()
food_group = pygame.sprite.Group()

# 初始化食物
for _ in range(10):
    food = Food()
    all_sprites.add(food)
    food_group.add(food)

all_sprites.add(snake)

running = True
clock = pygame.time.Clock()
while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False

    all_sprites.update()
    screen.fill(BLACK)
    draw_grid()
    all_sprites.draw(screen)
    pygame.display.flip()
    clock.tick(1000)

pygame.quit()