import pygame
import random
import time

# Initialize Pygame
pygame.init()

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

# Screen dimensions
WIDTH, HEIGHT = 800, 600

# Define constants
BALL_SIZE = 20
PADDLE_WIDTH, PADDLE_HEIGHT = 10, 200
BALL_SPEED_X = 7
PADDLE_SPEED = 7
COUNTDOWN_TIME = 3
WINNING_SCORE = 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 = BALL_SPEED_X * random.choice((1, -1))
        self.speed_y = BALL_SPEED_X * random.choice((1, -1))

    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_y *= random.uniform(1.1, 1.2)  # Increase speed on bounce

        # Ball out of bounds
        if self.rect.left <= 0:
            return "player2_scored"
        elif self.rect.right >= WIDTH:
            return "player1_scored"
        
        return None

    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
running = True
while running:
    # Handle events
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False

    # Countdown
    for i in range(COUNTDOWN_TIME, 0, -1):
        screen.fill(BLACK)
        draw_text(str(i), font, WHITE, WIDTH // 2, HEIGHT // 2)
        pygame.display.flip()
        time.sleep(1)

    # Main game loop
    game_over = False
    while not game_over:
        # Handle events
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False

        # Move paddles
        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
        score_event = ball.update()
        if score_event == "player1_scored":
            player1_score += 1
            ball.reset()
            game_over = True
        elif score_event == "player2_scored":
            player2_score += 1
            ball.reset()
            game_over = True

        # 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)

        # Check for winning condition
        if player1_score >= WINNING_SCORE or player2_score >= WINNING_SCORE:
            game_over = True

# Quit Pygame
pygame.quit()