import pygame
import sys
import random

# Initialize Pygame
pygame.init()

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

# Screen dimensions
WIDTH, HEIGHT = 800, 600

# Define constants
PADDLE_WIDTH, PADDLE_HEIGHT = 10, 100
BALL_SIZE = 20
PADDLE_SPEED = 7
INITIAL_BALL_SPEED_X = 5
INITIAL_BALL_SPEED_Y = 5

# Create the screen
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Pong")

# Clock for controlling the frame rate
clock = pygame.time.Clock()

# Font
font = pygame.font.Font(None, 36)

def draw_text(text, font, color, x, y):
    text_surface = font.render(text, True, color)
    text_rect = text_surface.get_rect()
    text_rect.center = (x, y)
    screen.blit(text_surface, text_rect)

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
        if self.rect.top < 0:
            self.rect.top = 0

    def move_down(self):
        self.rect.y += PADDLE_SPEED
        if self.rect.bottom > HEIGHT:
            self.rect.bottom = HEIGHT

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

class Ball:
    def __init__(self, x, y):
        self.rect = pygame.Rect(x, y, BALL_SIZE, BALL_SIZE)
        self.reset()

    def reset(self):
        self.rect.center = (WIDTH // 2, HEIGHT // 2)
        self.speed_x = random.choice([1, -1]) * INITIAL_BALL_SPEED_X
        self.speed_y = random.choice([1, -1]) * INITIAL_BALL_SPEED_Y

    def update(self):
        self.rect.x += self.speed_x
        self.rect.y += self.speed_y

        # Ball collision with walls
        if self.rect.top <= 0 or self.rect.bottom >= HEIGHT:
            self.speed_y *= -1

        # Ball collision with paddles
        if self.rect.colliderect(player1_paddle.rect) or self.rect.colliderect(player2_paddle.rect):
            self.speed_x *= -1
            self.speed_x *= 1.2  # Increase speed on each bounce
            self.speed_y *= 1.2

        # Ball out of bounds
        if self.rect.left <= 0 or self.rect.right >= WIDTH:
            return True  # Someone scored
        
        return False

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

# Create paddles
player1_paddle = Paddle(WIDTH - PADDLE_WIDTH - 20, HEIGHT // 2 - PADDLE_HEIGHT // 2)
player2_paddle = Paddle(20, HEIGHT // 2 - PADDLE_HEIGHT // 2)

# Create ball
ball = Ball(WIDTH // 2, HEIGHT // 2)

# Scores
player1_score = 0
player2_score = 0

# Main loop
while True:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            pygame.quit()
            sys.exit()

    keys = pygame.key.get_pressed()
    if keys[pygame.K_UP]:
        player1_paddle.move_up()
    if keys[pygame.K_DOWN]:
        player1_paddle.move_down()
    if keys[pygame.K_w]:
        player2_paddle.move_up()
    if keys[pygame.K_s]:
        player2_paddle.move_down()

    # Update ball and check for scoring
    if ball.update():
        if ball.rect.left <= 0:
            player1_score += 1
        elif ball.rect.right >= WIDTH:
            player2_score += 1

        if player1_score >= 5 or player2_score >= 5:
            # Someone won
            print("Game Over")
            pygame.quit()
            sys.exit()

        # Reset ball for the next round
        ball.reset()

    # Clear the screen
    screen.fill(BLACK)

    # Draw paddles and ball
    player1_paddle.draw()
    player2_paddle.draw()
    ball.draw()

    # Draw scores
    draw_text(str(player1_score), font, WHITE, WIDTH // 2 + 50, 50)
    draw_text(str(player2_score), font, WHITE, WIDTH // 2 - 50, 50)

    # Update the display
    pygame.display.flip()

    # Cap the frame rate
    clock.tick(60)