import pygame
import sys
import random
import math
import os

# Инициализация Pygame
pygame.init()

# Размеры экрана
WIDTH, HEIGHT = 800, 600
SCREEN = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Snake Game")

# Цвета
BACKGROUND_COLOR = (20, 21, 31)
SNAKE_COLOR = (50, 205, 50)
FOOD_COLOR = (255, 69, 0)
PARTICLE_COLOR = (255, 69, 0)

# Цвета для свечения
SNAKE_GLOW_COLOR = (80, 220, 80, 76)  # RGBA для прозрачности
FOOD_GLOW_COLOR = (255, 100, 50, 76)

# Шрифты
FONT = pygame.font.SysFont('Arial', 20)
BIG_FONT = pygame.font.SysFont('Arial', 50)

# FPS
FPS = 60
CLOCK = pygame.time.Clock()

# Путь к файлу рекорда
HIGH_SCORE_FILE = "high_score.txt"

# Класс Змейки
class Snake:
    def __init__(self, x, y):
        self.angle = -90  # Направление вверх
        self.turn_direction = 0  # -1 для поворота влево, 1 для поворота вправо, 0 для прямого движения
        self.speed = 5  # Увеличена скорость для соответствия spacing
        self.turn_speed = 3  # градусов за кадр
        self.radius = 10
        self.length = 5  # Начальная длина змейки
        self.spacing = 20  # Расстояние между сегментами увеличено до 20

        # История пути (path tracking)
        self.path = []  # Список позиций головы
        self.path_length = self.spacing * self.length  # Общее расстояние, на которое сегменты должны следовать

        # Инициализируем начальную позицию головы
        self.head_pos = {'x': x, 'y': y}
        self.path.append((self.head_pos['x'], self.head_pos['y']))
        print(f"[DEBUG] Snake initialized with head at ({x}, {y})")

    def turn(self, direction):
        self.turn_direction = direction
        print(f"[DEBUG] Snake turn direction set to {direction}")

    def grow(self):
        self.length += 1
        self.path_length = self.spacing * self.length
        print(f"[DEBUG] Snake grew to length {self.length}")

    def update(self):
        # Обновляем угол направления
        self.angle += self.turn_direction * self.turn_speed
        self.angle %= 360  # Обеспечиваем, что угол остается в диапазоне [0, 360)
        radians = math.radians(self.angle)
        dx = math.cos(radians) * self.speed
        dy = math.sin(radians) * self.speed

        # Обновляем позицию головы
        self.head_pos['x'] += dx
        self.head_pos['y'] += dy
        self.path.append((self.head_pos['x'], self.head_pos['y']))
        print(f"[DEBUG] Snake head moved to ({self.head_pos['x']:.2f}, {self.head_pos['y']:.2f})")

        # Проверка выхода за границы экрана
        if not (0 <= self.head_pos['x'] <= WIDTH and 0 <= self.head_pos['y'] <= HEIGHT):
            print("[DEBUG] Snake collided with wall")
            return True  # Возвращаем True для обработки столкновения с границей

        # Очищаем историю пути, чтобы она не была слишком длинной
        while self.calculate_path_length() > self.path_length:
            removed_position = self.path.pop(0)
            print(f"[DEBUG] Removed position from path: {removed_position}")

        return False  # Нет столкновения с границей

    def calculate_path_length(self):
        length = 0
        for i in range(1, len(self.path)):
            prev = self.path[i-1]
            curr = self.path[i]
            length += math.hypot(curr[0] - prev[0], curr[1] - prev[1])
        return length

    def get_segment_position(self, index):
        target_length = self.spacing * index
        accumulated_length = 0
        for i in range(len(self.path)-1, 0, -1):
            prev = self.path[i-1]
            curr = self.path[i]
            segment = math.hypot(curr[0] - prev[0], curr[1] - prev[1])
            if accumulated_length + segment >= target_length:
                ratio = (target_length - accumulated_length) / segment if segment != 0 else 0
                x = curr[0] + (prev[0] - curr[0]) * ratio
                y = curr[1] + (prev[1] - curr[1]) * ratio
                return (x, y)
            accumulated_length += segment
        return self.path[0]

    def draw(self, surface):
        # Рисуем сегменты змейки
        for i in range(1, self.length + 1):
            pos = self.get_segment_position(i)
            x, y = pos

            # Эффект свечения
            alpha = 0.3 * (1 - i / self.length)
            glow_surface = pygame.Surface((self.radius * 2 + 8, self.radius * 2 + 8), pygame.SRCALPHA)
            pygame.draw.circle(glow_surface, (80, 220, 80, int(255 * alpha)), (self.radius + 4, self.radius + 4), self.radius + 4)
            surface.blit(glow_surface, (x - self.radius - 4, y - self.radius - 4))
            
            # Рисуем тело змейки
            pygame.draw.circle(surface, SNAKE_COLOR, (int(x), int(y)), self.radius)
            pygame.draw.circle(surface, (64, 255, 64), (int(x), int(y)), self.radius, 2)

    def check_self_collision(self, start_delay_frames):
        if start_delay_frames > 0:
            # Пропускаем проверку в течение начальной задержки
            return False

        # Голова змейки
        head_x, head_y = self.head_pos['x'], self.head_pos['y']
        for i in range(2, self.length + 1):
            segment_x, segment_y = self.get_segment_position(i)
            distance = math.hypot(head_x - segment_x, head_y - segment_y)
            if distance < self.radius * 2:
                print("[DEBUG] Snake collided with itself")
                return True
        return False

# Класс Еды
class Food:
    def __init__(self):
        self.radius = 6
        self.respawn()
        print(f"[DEBUG] Food initialized at ({self.x}, {self.y})")

    def respawn(self):
        self.x = random.randint(self.radius, WIDTH - self.radius)
        self.y = random.randint(self.radius, HEIGHT - self.radius)
        print(f"[DEBUG] Food respawned at ({self.x}, {self.y})")

    def draw(self, surface):
        # Эффект свечения
        glow_surface = pygame.Surface((self.radius * 2 + 8, self.radius * 2 + 8), pygame.SRCALPHA)
        pygame.draw.circle(glow_surface, FOOD_GLOW_COLOR, (self.radius + 4, self.radius + 4), self.radius + 4)
        surface.blit(glow_surface, (self.x - self.radius - 4, self.y - self.radius - 4))
        
        # Рисуем еду
        pygame.draw.circle(surface, FOOD_COLOR, (int(self.x), int(self.y)), self.radius)
        pygame.draw.circle(surface, (255, 99, 71), (int(self.x), int(self.y)), self.radius, 2)

# Класс Частиц
class ParticleSystem:
    def __init__(self):
        self.particles = []
        print("[DEBUG] ParticleSystem initialized")

    def emit(self, x, y):
        for _ in range(20):
            angle = random.uniform(0, 2 * math.pi)
            speed = random.uniform(1, 4)
            self.particles.append({
                'x': x,
                'y': y,
                'dx': math.cos(angle) * speed,
                'dy': math.sin(angle) * speed,
                'life': 1.0,
                'color': PARTICLE_COLOR
            })
        print(f"[DEBUG] Emitted 20 particles at ({x}, {y})")

    def update(self):
        original_count = len(self.particles)
        self.particles = [p for p in self.particles if p['life'] > 0]
        removed_particles = original_count - len(self.particles)
        if removed_particles > 0:
            print(f"[DEBUG] Removed {removed_particles} particles")
        for p in self.particles:
            p['x'] += p['dx']
            p['y'] += p['dy']
            p['life'] -= 0.02

    def draw(self, surface):
        for p in self.particles:
            if p['life'] <= 0:
                continue
            alpha = max(int(255 * p['life']), 0)
            particle_surface = pygame.Surface((4, 4), pygame.SRCALPHA)
            particle_surface.fill((*p['color'], alpha))
            surface.blit(particle_surface, (int(p['x']), int(p['y'])))

# Класс Меню (Главное меню и Game Over меню)
class Menu:
    def __init__(self, title, options, on_select):
        """
        title: Текст заголовка меню.
        options: Список текстов опций.
        on_select: Функция обратного вызова, принимающая индекс выбранной опции.
        """
        self.title = title
        self.options = options
        self.on_select = on_select
        self.selected_index = 0
        self.alpha = 0
        self.show_menu_surface = True
        self.fade_in = True
        self.fade_out = False
        self.menu_surface = pygame.Surface((WIDTH, HEIGHT), pygame.SRCALPHA)
        self.setup_buttons()
        print(f"[DEBUG] '{self.title}' Menu initialized")

    def setup_buttons(self):
        # Расположение кнопок по вертикали
        self.buttons = []
        button_width = 200
        button_height = 50
        spacing = 20
        total_height = len(self.options) * (button_height + spacing) - spacing
        start_y = HEIGHT // 2 - total_height // 2

        for idx, option in enumerate(self.options):
            rect = pygame.Rect(WIDTH//2 - button_width//2, start_y + idx * (button_height + spacing), button_width, button_height)
            self.buttons.append(rect)
        print(f"[DEBUG] Buttons for '{self.title}' Menu set up")

    def handle_event(self, event):
        if event.type == pygame.MOUSEBUTTONDOWN:
            for idx, button in enumerate(self.buttons):
                if button.collidepoint(event.pos):
                    print(f"[DEBUG] '{self.title}' Menu option '{self.options[idx]}' clicked")
                    self.fade_out = True

        elif event.type == pygame.KEYDOWN:
            if event.key == pygame.K_UP:
                self.selected_index = (self.selected_index - 1) % len(self.options)
                print(f"[DEBUG] '{self.title}' Menu selection moved to {self.selected_index}")
            elif event.key == pygame.K_DOWN:
                self.selected_index = (self.selected_index + 1) % len(self.options)
                print(f"[DEBUG] '{self.title}' Menu selection moved to {self.selected_index}")
            elif event.key == pygame.K_RETURN or event.key == pygame.K_SPACE:
                print(f"[DEBUG] '{self.title}' Menu option '{self.options[self.selected_index]}' selected via keyboard")
                self.on_select(self.selected_index)
                self.fade_out = True

    def update(self):
        if self.fade_in and self.alpha < 255:
            self.alpha += 5
            if self.alpha >= 255:
                self.alpha = 255
                self.fade_in = False
                print(f"[DEBUG] '{self.title}' Menu fade-in completed")
        elif self.fade_out and self.alpha > 0:
            self.alpha -= 5
            if self.alpha <= 0:
                self.alpha = 0
                self.show_menu_surface = False
                print(f"[DEBUG] '{self.title}' Menu fade-out completed")
                # Вызов callback-функции после завершения анимации
                if self.on_select:
                    if self.title == "Главное Меню":
                        # Для главного меню используем текущий выбор
                        self.on_select(self.selected_index)
                    elif self.title == "Game Over":
                        # Для Game Over меню выбранный пункт индексу
                        self.on_select(self.selected_index)

        # Заполнение поверхности меню
        self.menu_surface.fill((0, 0, 0, self.alpha))

    def draw(self, surface):
        if self.show_menu_surface:
            surface.blit(self.menu_surface, (0, 0))
            # Рисуем заголовок
            title_text = BIG_FONT.render(self.title, True, (255, 255, 255))
            surface.blit(title_text, (WIDTH//2 - title_text.get_width()//2, HEIGHT//2 - 200))
            # Рисуем кнопки
            for idx, button in enumerate(self.buttons):
                color = (150, 150, 150) if idx != self.selected_index else (100, 100, 255)
                pygame.draw.rect(surface, color, button)
                # Рисуем текст кнопки
                option_text = FONT.render(self.options[idx], True, (255, 255, 255))
                surface.blit(option_text, (button.x + button.width//2 - option_text.get_width()//2,
                                           button.y + button.height//2 - option_text.get_height()//2))

# Класс Паузы
class PauseMenu:
    def __init__(self, on_resume, on_main_menu):
        self.on_resume = on_resume
        self.on_main_menu = on_main_menu
        self.show_menu = False
        self.menu_surface = pygame.Surface((WIDTH, HEIGHT), pygame.SRCALPHA)
        self.resume_button = pygame.Rect(WIDTH//2 - 100, HEIGHT//2 - 50, 200, 50)
        self.main_menu_button = pygame.Rect(WIDTH//2 - 100, HEIGHT//2 + 20, 200, 50)
        print("[DEBUG] PauseMenu initialized")

    def handle_event(self, event):
        if event.type == pygame.MOUSEBUTTONDOWN:
            if self.resume_button.collidepoint(event.pos):
                print("[DEBUG] Resume button clicked in PauseMenu")
                self.on_resume()
            elif self.main_menu_button.collidepoint(event.pos):
                print("[DEBUG] Main Menu button clicked in PauseMenu")
                self.on_main_menu()

    def show_menu_func(self, menu=True):
        self.show_menu = menu
        print(f"[DEBUG] PauseMenu show_menu set to {self.show_menu}")

    def hide(self):
        self.show_menu = False
        print("[DEBUG] PauseMenu hidden")

    def update(self):
        pass  # Место для анимации или других обновлений паузы

    def draw(self, surface):
        if self.show_menu:
            # Полупрозрачный фон
            self.menu_surface.fill((0, 0, 0, 150))
            surface.blit(self.menu_surface, (0, 0))
            # Текст "Paused"
            paused_text = BIG_FONT.render("Paused", True, (255, 255, 255))
            surface.blit(paused_text, (WIDTH//2 - paused_text.get_width()//2, HEIGHT//2 - 200))
            # Кнопки
            pygame.draw.rect(surface, (100, 100, 100), self.resume_button)
            pygame.draw.rect(surface, (100, 100, 100), self.main_menu_button)
            # Рисуем текст кнопок
            resume_text = FONT.render("Resume", True, (255, 255, 255))
            main_menu_text = FONT.render("Main Menu", True, (255, 255, 255))
            surface.blit(resume_text, (self.resume_button.x + self.resume_button.width//2 - resume_text.get_width()//2,
                                       self.resume_button.y + self.resume_button.height//2 - resume_text.get_height()//2))
            surface.blit(main_menu_text, (self.main_menu_button.x + self.main_menu_button.width//2 - main_menu_text.get_width()//2,
                                          self.main_menu_button.y + self.main_menu_button.height//2 - main_menu_text.get_height()//2))

# Класс Game Over Меню
class GameOverMenu(Menu):
    def __init__(self, score, high_score, on_restart, on_main_menu):
        self.score = score
        self.high_score = high_score
        self.on_restart = on_restart
        self.on_main_menu = on_main_menu
        options = ["Перезапустить", "Главное Меню"]
        super().__init__("Game Over", options, self.handle_selection)
        print("[DEBUG] GameOverMenu initialized")

    def handle_selection(self, selected_index):
        if selected_index == 0:
            print("[DEBUG] GameOverMenu: Restart selected")
            self.on_restart()
        elif selected_index == 1:
            print("[DEBUG] GameOverMenu: Main Menu selected")
            self.on_main_menu()

    def draw(self, surface):
        if self.show_menu_surface:
            surface.blit(self.menu_surface, (0, 0))
            # Рисуем заголовок
            title_text = BIG_FONT.render(self.title, True, (255, 255, 255))
            surface.blit(title_text, (WIDTH//2 - title_text.get_width()//2, HEIGHT//2 - 250))
            # Рисуем счёт
            score_text = FONT.render(f"Ваш счёт: {self.score}", True, (255, 255, 255))
            high_score_text = FONT.render(f"Рекорд: {self.high_score}", True, (255, 255, 255))
            surface.blit(score_text, (WIDTH//2 - score_text.get_width()//2, HEIGHT//2 - 150))
            surface.blit(high_score_text, (WIDTH//2 - high_score_text.get_width()//2, HEIGHT//2 - 120))
            # Рисуем кнопки
            for idx, button in enumerate(self.buttons):
                color = (150, 150, 150) if idx != self.selected_index else (100, 100, 255)
                pygame.draw.rect(surface, color, button)
                # Рисуем текст кнопки
                option_text = FONT.render(self.options[idx], True, (255, 255, 255))
                surface.blit(option_text, (button.x + button.width//2 - option_text.get_width()//2,
                                           button.y + button.height//2 - option_text.get_height()//2))

# Класс Игры
class Game:
    def __init__(self, on_return_to_menu):
        self.snake = Snake(WIDTH / 2, HEIGHT / 2)
        self.food = Food()
        self.particles = ParticleSystem()
        self.score = 0
        self.high_score = self.load_high_score()
        self.is_paused = False
        self.active = True
        self.game_over = False
        self.start_delay_frames = FPS * 2  # 2 секунды задержки
        print("[DEBUG] Game initialized")

        # Определяем функции обратного вызова
        def on_main_menu():
            self.save_high_score()
            on_return_to_menu()

        def on_restart():
            self.reset_game()

        self.pause_menu = PauseMenu(self.resume, on_main_menu)
        self.game_over_menu = None  # Изначально отсутствует
        print("[DEBUG] PauseMenu initialized")

    def reset_game(self):
        self.snake = Snake(WIDTH / 2, HEIGHT / 2)
        self.food = Food()
        self.particles = ParticleSystem()
        self.score = 0
        self.is_paused = False
        self.active = True
        self.game_over = False
        self.start_delay_frames = FPS * 2  # 2 секунды задержки
        self.game_over_menu = None
        print("[DEBUG] Game has been reset for a new session")

    def load_high_score(self):
        if not os.path.exists(HIGH_SCORE_FILE):
            # Создаём файл с начальным значением 0
            try:
                with open(HIGH_SCORE_FILE, 'w') as f:
                    f.write('0')
                print(f"[DEBUG] High score file created at {HIGH_SCORE_FILE}")
            except Exception as e:
                print(f"[ERROR] Failed to create high score file: {e}")
            return 0
        try:
            with open(HIGH_SCORE_FILE, 'r') as f:
                high_score = int(f.read())
                print(f"[DEBUG] High score loaded: {high_score}")
                return high_score
        except ValueError:
            print(f"[DEBUG] High score file corrupted. Resetting high score to 0.")
            return 0
        except Exception as e:
            print(f"[ERROR] Error loading high score: {e}")
            return 0

    def save_high_score(self):
        try:
            with open(HIGH_SCORE_FILE, 'w') as f:
                f.write(str(self.high_score))
            print(f"[DEBUG] High score {self.high_score} saved to {HIGH_SCORE_FILE}")
        except Exception as e:
            print(f"[ERROR] Error saving high score: {e}")

    def handle_events(self, events):
        if self.game_over and self.game_over_menu:
            self.game_over_menu.handle_event(pygame.event.Event(pygame.USEREVENT))
            for event in events:
                self.game_over_menu.handle_event(event)
            return

        for event in events:
            if event.type == pygame.QUIT:
                self.save_high_score()
                pygame.quit()
                sys.exit()
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_ESCAPE:
                    self.toggle_pause()
                elif not self.is_paused and not self.game_over:
                    if event.key == pygame.K_LEFT:
                        self.snake.turn(-1)
                        print("[DEBUG] Left arrow key pressed")
                    elif event.key == pygame.K_RIGHT:
                        self.snake.turn(1)
                        print("[DEBUG] Right arrow key pressed")
            elif event.type == pygame.KEYUP:
                if not self.is_paused and not self.game_over:
                    if event.key in [pygame.K_LEFT, pygame.K_RIGHT]:
                        self.snake.turn(0)
                        print("[DEBUG] Arrow key released")
            if self.is_paused:
                self.pause_menu.handle_event(event)

    def toggle_pause(self):
        self.is_paused = not self.is_paused
        if self.is_paused:
            self.pause_menu.show_menu_func(True)
            print("[DEBUG] Game paused")
        else:
            self.pause_menu.show_menu_func(False)
            print("[DEBUG] Game resumed")

    def resume(self):
        self.is_paused = False
        self.pause_menu.show_menu_func(False)
        print("[DEBUG] Resumed from pause")

    def return_to_menu(self):
        self.save_high_score()
        self.active = False
        print("[DEBUG] Returning to MainMenu")

    def update(self):
        if not self.is_paused and not self.game_over:
            collision_with_wall = self.snake.update()
            self.particles.update()
            self.check_collisions()
            if collision_with_wall:
                self.handle_game_over()

            if self.start_delay_frames <= 0:
                if self.snake.check_self_collision(self.start_delay_frames):
                    self.handle_game_over()
            else:
                self.start_delay_frames -=1
                if self.start_delay_frames == 0:
                    print("[DEBUG] Collision detection enabled")

    def handle_game_over(self):
        self.save_high_score()
        self.game_over = True
        self.active = False
        self.game_over_menu = GameOverMenu(
            score=self.score,
            high_score=self.high_score,
            on_restart=self.reset_game,
            on_main_menu=self.return_to_menu
        )
        print("[DEBUG] Game Over triggered")

    def check_collisions(self):
        head = self.snake.head_pos
        distance = math.hypot(head['x'] - self.food.x, head['y'] - self.food.y)
        if distance < self.snake.radius + self.food.radius:
            print("[DEBUG] Snake ate the food")
            self.snake.grow()
            self.particles.emit(self.food.x, self.food.y)
            self.food.respawn()
            self.score += 10
            if self.score > self.high_score:
                self.high_score = self.score
                print(f"[DEBUG] New high score: {self.high_score}")

    def draw(self, surface):
        surface.fill(BACKGROUND_COLOR)
        self.food.draw(surface)
        self.snake.draw(surface)
        self.particles.draw(surface)
        # Рисуем счёт
        score_text = FONT.render(f"Score: {self.score}", True, (255, 255, 255))
        high_score_text = FONT.render(f"High Score: {self.high_score}", True, (255, 255, 255))
        surface.blit(score_text, (20, 20))
        surface.blit(high_score_text, (20, 50))

        if self.is_paused:
            self.pause_menu.draw(surface)
        if self.game_over and self.game_over_menu:
            self.game_over_menu.draw(surface)

# Класс Главного Меню
class MainMenu:
    def __init__(self):
        self.menu = Menu(
            title="Главное Меню",
            options=["Start Game", "Exit"],
            on_select=self.handle_selection
        )
        self.active = True
        self.game_instance = None  # Изначально игра отсутствует
        print("[DEBUG] MainMenu initialized")

    def handle_selection(self, selected_index):
        if selected_index == 0:
            self.start_game()
        elif selected_index == 1:
            self.exit_game()

    def start_game(self):
        print("[DEBUG] Start Game button clicked in MainMenu")
        print("[DEBUG] Starting the game from MainMenu")
        self.game_instance = Game(self.show_menu)
        self.active = False

    def exit_game(self):
        print("[DEBUG] Exit button clicked in MainMenu")
        print("[DEBUG] Exiting the game from MainMenu")
        pygame.quit()
        sys.exit()

    def handle_events(self, events):
        for event in events:
            self.menu.handle_event(event)

    def update(self):
        self.menu.update()

    def draw(self, surface):
        self.menu.draw(surface)

    def show_menu(self):
        print("[DEBUG] Returning to MainMenu")
        if self.game_instance:
            self.game_instance.save_high_score()
            self.game_instance = None
        self.active = True

# Класс Главного Цикла
class GameOverScreen:
    def __init__(self, game_over_menu):
        self.game_over_menu = game_over_menu

    def handle_events(self, events):
        for event in events:
            self.game_over_menu.handle_event(event)

    def update(self):
        self.game_over_menu.update()

    def draw(self, surface):
        self.game_over_menu.draw(surface)

def main():
    main_menu = MainMenu()
    game_over_screen = None

    while True:
        events = pygame.event.get()
        if main_menu.active:
            main_menu.handle_events(events)
            main_menu.update()
            main_menu.draw(SCREEN)
        elif main_menu.game_instance and main_menu.game_instance.active:
            main_menu.game_instance.handle_events(events)
            main_menu.game_instance.update()
            main_menu.game_instance.draw(SCREEN)
        elif main_menu.game_instance and main_menu.game_instance.game_over_menu:
            # Игра окончена, показываем GameOverMenu
            game_over_screen = GameOverScreen(main_menu.game_instance.game_over_menu)
            game_over_screen.handle_events(events)
            game_over_screen.update()
            game_over_screen.draw(SCREEN)
        else:
            main_menu.show_menu()

        pygame.display.flip()
        CLOCK.tick(FPS)

# Запуск Игры
if __name__ == "__main__":
    main()

После проигрыша и нажатия на клавишу главного меню игра перезапускается, а не загружает главное меню. И ещё при нажатии на Esc и последующем нажатии на кнопку главного меню игра крашится с ошибкой а не загружает меню.