import pygame
import sys

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

# Константы
WIDTH, HEIGHT = 800, 600
PADDLE_WIDTH, PADDLE_HEIGHT = 10, 100
BALL_SIZE = 15
WHITE = (255, 255, 255)
FPS = 60
WINNING_SCORE = 10  # Основной счёт для победы
TIE_BREAK_REQUIRED_DIFFERENCE = 2  # Разница для победы в тай-брейке
AI_SPEED = 0.7  # Уровень "умения" ИИ (0.5 - очень медленный, 1.0 - идеальный)

# Создание экрана
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Pong с ИИ")
clock = pygame.time.Clock()

# Функция для сброса игры
def reset_game():
    return {
        "player_paddle": pygame.Rect(20, HEIGHT//2 - PADDLE_HEIGHT//2, PADDLE_WIDTH, PADDLE_HEIGHT),
        "ai_paddle": pygame.Rect(WIDTH - 30, HEIGHT//2 - PADDLE_HEIGHT//2, PADDLE_WIDTH, PADDLE_HEIGHT),
        "ball": pygame.Rect(WIDTH//2 - BALL_SIZE//2, HEIGHT//2 - BALL_SIZE//2, BALL_SIZE, BALL_SIZE),
        "ball_speed_x": 5,
        "ball_speed_y": 5,
        "player_score": 0,
        "ai_score": 0,
        "game_over": False,
        "tie_break_active": False,
        "win_message": None,
        "show_win_screen": False,
        "restart_timer": 0
    }

# Состояние игры
game_state = reset_game()

# Шрифты
score_font = pygame.font.Font(None, 74)
win_font = pygame.font.Font(None, 100)

def move_ai_paddle(paddle, ball):
    """Логика ИИ: ракетка следует за мячом с задержкой"""
    if ball.centerx > WIDTH/2:  # Реагируем только на мяч, приближающийся к ИИ
        if paddle.centery < ball.centery:
            paddle.y += min(7 * AI_SPEED, ball.centery - paddle.centery)
        elif paddle.centery > ball.centery:
            paddle.y -= min(7 * AI_SPEED, paddle.centery - ball.centery)
        
        # Ограничение границ экрана
        paddle.y = max(0, min(HEIGHT - PADDLE_HEIGHT, paddle.y))

def calculate_ball_angle(contact_point, paddle_center):
    """Рассчитывает угол отскока мяча в зависимости от точки контакта"""
    contact_diff = contact_point - paddle_center
    angle_factor = contact_diff / (PADDLE_HEIGHT/2)  # От -1 до +1
    
    # Более резкий угол для ударов по краям ракетки
    return angle_factor * 0.5  # 0.5 - коэффициент максимального изменения угла

# Основной цикл игры
while True:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            pygame.quit()
            sys.exit()
    
    if not game_state["show_win_screen"]:
        # Обработка управления
        keys = pygame.key.get_pressed()
        if keys[pygame.K_w] and game_state["player_paddle"].top > 0:
            game_state["player_paddle"].y -= 7
        if keys[pygame.K_s] and game_state["player_paddle"].bottom < HEIGHT:
            game_state["player_paddle"].y += 7

        # Движение мяча
        game_state["ball"].x += game_state["ball_speed_x"]
        game_state["ball"].y += game_state["ball_speed_y"]

        # Отскок от стен
        if game_state["ball"].top <= 0 or game_state["ball"].bottom >= HEIGHT:
            game_state["ball_speed_y"] *= -1

        # Проверка на гол
        if game_state["ball"].left <= 0:
            game_state["ai_score"] += 1
            
            # Проверка условий победы после гола
            if game_state["ai_score"] == WINNING_SCORE:
                if game_state["player_score"] == WINNING_SCORE - 1:
                    game_state["tie_break_active"] = True
                elif game_state["player_score"] <= WINNING_SCORE - 2:
                    game_state["win_message"] = "КОМПЬЮТЕР ПОБЕДИЛ!"
                    game_state["show_win_screen"] = True
                    game_state["game_over"] = True
                    game_state["restart_timer"] = pygame.time.get_ticks() + 1000
            
            # Проверка условий победы в тай-брейке
            if game_state["tie_break_active"] and abs(game_state["ai_score"] - game_state["player_score"]) >= TIE_BREAK_REQUIRED_DIFFERENCE:
                game_state["win_message"] = "КОМПЬЮТЕР ПОБЕДИЛ!"
                game_state["show_win_screen"] = True
                game_state["game_over"] = True
                game_state["restart_timer"] = pygame.time.get_ticks() + 1000
            
            game_state["ball"].center = (WIDTH//2, HEIGHT//2)
            game_state["ball_speed_x"] = 5
            game_state["ball_speed_y"] = 5 * random.choice([-1, 1])
            pygame.time.delay(500)
            
        elif game_state["ball"].right >= WIDTH:
            game_state["player_score"] += 1
            
            # Проверка условий победы после гола
            if game_state["player_score"] == WINNING_SCORE:
                if game_state["ai_score"] == WINNING_SCORE - 1:
                    game_state["tie_break_active"] = True
                elif game_state["ai_score"] <= WINNING_SCORE - 2:
                    game_state["win_message"] = "ТЫ ПОБЕДИЛ!"
                    game_state["show_win_screen"] = True
                    game_state["game_over"] = True
                    game_state["restart_timer"] = pygame.time.get_ticks() + 1000
            
            # Проверка условий победы в тай-брейке
            if game_state["tie_break_active"] and abs(game_state["player_score"] - game_state["ai_score"]) >= TIE_BREAK_REQUIRED_DIFFERENCE:
                game_state["win_message"] = "ТЫ ПОБЕДИЛ!"
                game_state["show_win_screen"] = True
                game_state["game_over"] = True
                game_state["restart_timer"] = pygame.time.get_ticks() + 1000
            
            game_state["ball"].center = (WIDTH//2, HEIGHT//2)
            game_state["ball_speed_x"] = -5
            game_state["ball_speed_y"] = 5 * random.choice([-1, 1])
            pygame.time.delay(500)

        # Отскок от ракеток с учётом угла
        if game_state["ball"].colliderect(game_state["player_paddle"]):
            contact_point = game_state["ball"].centery
            paddle_center = game_state["player_paddle"].centery
            
            angle_change = calculate_ball_angle(contact_point, paddle_center)
            game_state["ball_speed_y"] += angle_change
            game_state["ball_speed_x"] *= -1.1
            
        elif game_state["ball"].colliderect(game_state["ai_paddle"]):
            contact_point = game_state["ball"].centery
            paddle_center = game_state["ai_paddle"].centery
            
            angle_change = calculate_ball_angle(contact_point, paddle_center)
            game_state["ball_speed_y"] += angle_change
            game_state["ball_speed_x"] *= -1.1

        # Обновление позиции ИИ
        if not game_state["game_over"]:
            move_ai_paddle(game_state["ai_paddle"], game_state["ball"])

    # Обработка события нажатия клавиши для перезапуска
    if game_state["show_win_screen"]:
        keys = pygame.key.get_pressed()
        if any(keys):  # Любая клавиша
            game_state = reset_game()

    # Отрисовка
    screen.fill((0, 0, 0))
    
    if not game_state["show_win_screen"]:
        # Отображение счета
        score_text = score_font.render(f"{game_state['player_score']}  {game_state['ai_score']}", True, WHITE)
        screen.blit(score_text, (WIDTH//2 - 50, 10))
        
        # Текст тай-брейка
        if game_state["tie_break_active"]:
            tie_text = score_font.render("TIE BREAK!", True, WHITE)
            screen.blit(tie_text, (WIDTH//2 - 100, HEIGHT//2))
        
        # Основные элементы игры
        pygame.draw.rect(screen, WHITE, game_state["player_paddle"])
        pygame.draw.rect(screen, WHITE, game_state["ai_paddle"])
        pygame.draw.ellipse(screen, WHITE, game_state["ball"])
        pygame.draw.aaline(screen, WHITE, (WIDTH//2, 0), (WIDTH//2, HEIGHT))
    else:
        # Экран победы
        win_text = win_font.render(game_state["win_message"], True, WHITE)
        text_rect = win_text.get_rect(center=(WIDTH//2, HEIGHT//2))
        screen.blit(win_text, text_rect)
    
    pygame.display.flip()
    clock.tick(FPS)
