import pygame
import random

# Initialize Pygame
pygame.init()

# Set up the screenw
WIDTH, HEIGHT = 800, 600
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Pong")

# Colors
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)

# Define the paddles
PADDLE_WIDTH = 10
PADDLE_HEIGHT = 100
PADDLE_SPEED = 10

class Paddle:
    def __init__(self, x, y):
        self.rect = pygame.Rect(x, y, PADDLE_WIDTH, PADDLE_HEIGHT)

    def move_up(self):
        self.rect.y -= PADDLE_SPEED

    def move_down(self):
        self.rect.y += PADDLE_SPEED

    def draw(self):
        pygame.draw.rect(screen, WHITE, self.rect)

# Define the ball
BALL_SIZE = 20
BALL_SPEED = 5

class Ball:
    def __init__(self):
        self.rect = pygame.Rect(WIDTH // 2, HEIGHT // 2, BALL_SIZE, BALL_SIZE)
        self.dx = BALL_SPEED * random.choice([-1, 1])
        self.dy = BALL_SPEED * random.choice([-1, 1])

    def move(self):
        self.rect.x += self.dx
        self.rect.y += self.dy

        # Bounce off the top and bottom edges
        if self.rect.y <= 0 or self.rect.y >= HEIGHT - BALL_SIZE:
            self.dy *= -1

    def draw(self):
        pygame.draw.rect(screen, WHITE, self.rect)

# Create paddles and ball
paddle1 = Paddle(50, HEIGHT // 2 - PADDLE_HEIGHT // 2)
paddle2 = Paddle(WIDTH - 50 - PADDLE_WIDTH, HEIGHT // 2 - PADDLE_HEIGHT // 2)
ball = Ball()

# Game variables
player1_score = 0
player2_score = 0
rounds_to_win = 5

clock = pygame.time.Clock()
running = True

def move_ai_paddle():
    # Simple AI that occasionally makes a mistake
    if random.randint(1, 100) < 34:
        return 0  # Make a mistake
    elif ball.rect.centery < paddle2.rect.centery:
        return -PADDLE_SPEED
    elif ball.rect.centery > paddle2.rect.centery:
        return PADDLE_SPEED
    else:
        return 0

while running:
    screen.fill(BLACK)

    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False

    keys = pygame.key.get_pressed()
    if keys[pygame.K_w]:
        paddle1.move_up()
    if keys[pygame.K_s]:
        paddle1.move_down()

    paddle2_speed = move_ai_paddle()
    paddle2.rect.y += paddle2_speed

    ball.move()

    # Check for collisions with paddles
    if ball.rect.colliderect(paddle1.rect) or ball.rect.colliderect(paddle2.rect):
        ball.dx *= -1

    # Check if ball goes out of bounds
    if ball.rect.x <= 0:
        player2_score += 1
        if player2_score == rounds_to_win:
            print("Player 2 wins!")
            running = False
        else:
            ball = Ball()
    elif ball.rect.x >= WIDTH - BALL_SIZE:
        player1_score += 1
        if player1_score == rounds_to_win:
            print("Player 1 wins!")
            running = False
        else:
            ball = Ball()

    paddle1.draw()
    paddle2.draw()
    ball.draw()

    pygame.display.flip()
    clock.tick(60)

pygame.quit()