import pygame
import sys

pygame.init()
WIDTH, HEIGHT = 800, 600
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Pong")
clock = pygame.time.Clock()

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

ball = pygame.Rect(WIDTH//2 - 10, HEIGHT//2 - 10, 20, 20)
cpu_paddle = pygame.Rect(20, HEIGHT//2 - 60, 10, 120)
player_paddle = pygame.Rect(WIDTH - 30, HEIGHT//2 - 60, 10, 120)

ball_speed = [5, 5]
paddle_speed = 6
cpu_speed = 5

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

    # Player input (WASD, Arrows, HJKL)
    keys = pygame.key.get_pressed()
    if (keys[pygame.K_w] or keys[pygame.K_UP] or keys[pygame.K_k]) and player_paddle.top > 0:
        player_paddle.y -= paddle_speed
    if (keys[pygame.K_s] or keys[pygame.K_DOWN] or keys[pygame.K_j]) and player_paddle.bottom < HEIGHT:
        player_paddle.y += paddle_speed

    # CPU movement (simple AI)
    if ball.centery > cpu_paddle.centery and cpu_paddle.bottom < HEIGHT:
        cpu_paddle.y += cpu_speed
    if ball.centery < cpu_paddle.centery and cpu_paddle.top > 0:
        cpu_paddle.y -= cpu_speed

    # Ball movement
    ball.x += ball_speed[0]
    ball.y += ball_speed[1]

    # Wall collision
    if ball.top <= 0 or ball.bottom >= HEIGHT:
        ball_speed[1] *= -1

    # Paddle collision
    if ball.colliderect(cpu_paddle) or ball.colliderect(player_paddle):
        ball_speed[0] *= -1

    # Reset ball
    if ball.left <= 0 or ball.right >= WIDTH:
        ball.center = (WIDTH//2, HEIGHT//2)
        ball_speed[0] *= -1

    # Draw
    screen.fill(BLACK)
    pygame.draw.rect(screen, WHITE, cpu_paddle)
    pygame.draw.rect(screen, WHITE, player_paddle)
    pygame.draw.ellipse(screen, WHITE, ball)
    pygame.draw.aaline(screen, WHITE, (WIDTH//2, 0), (WIDTH//2, HEIGHT))
    pygame.display.flip()

    clock.tick(60)