import pygame
import sys
import random
import math

# --- 初始化 ---
pygame.init()
WIDTH, HEIGHT = 900, 600
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Pygame 乒乓球 - 经典复刻")
clock = pygame.time.Clock()

# 颜色
BG_COLOR = (20, 20, 30)
PADDLE_COLOR = (0, 255, 200)
BALL_COLOR = (255, 255, 100)
TEXT_COLOR = (255, 255, 255)
LINE_COLOR = (100, 100, 120)

FONT = pygame.font.SysFont("consolas", 60)
SMALL_FONT = pygame.font.SysFont("consolas", 30)

# 配置
PADDLE_W, PADDLE_H = 15, 100
BALL_SIZE = 16
PADDLE_SPEED = 7
BALL_START_SPEED = 6
WINNING_SCORE = 10

# --- 游戏对象 ---

class Paddle:
    def __init__(self, x, is_ai=False):
        self.rect = pygame.Rect(x, HEIGHT//2 - PADDLE_H//2, PADDLE_W, PADDLE_H)
        self.speed = 0
        self.is_ai = is_ai
        self.score = 0

    def move(self, dy):
        self.rect.y += dy
        self.rect.top = max(0, self.rect.top)
        self.rect.bottom = min(HEIGHT, self.rect.bottom)

    def ai_update(self, ball):
        if not self.is_ai: return
        # 简单AI：球过来才动，且有反应延迟
        if ball.vx > 0:  # 球向右飞
            center = self.rect.centery
            if center < ball.rect.centery - 20:
                self.move(PADDLE_SPEED * 0.8)
            elif center > ball.rect.centery + 20:
                self.move(-PADDLE_SPEED * 0.8)
        else:
            # 球远离时慢慢回中
            center = self.rect.centery
            if center < HEIGHT//2 - 10: self.move(3)
            elif center > HEIGHT//2 + 10: self.move(-3)

class Ball:
    def __init__(self):
        self.rect = pygame.Rect(WIDTH//2 - BALL_SIZE//2, HEIGHT//2 - BALL_SIZE//2, BALL_SIZE, BALL_SIZE)
        self.reset()

    def reset(self, direction=None):
        self.rect.center = (WIDTH//2, HEIGHT//2)
        angle = random.uniform(-0.5, 0.5)  # 随机角度
        dir_x = direction if direction else random.choice([-1, 1])
        self.speed = BALL_START_SPEED
        self.vx = dir_x * self.speed * math.cos(angle)
        self.vy = self.speed * math.sin(angle)

    def update(self, left_paddle, right_paddle):
        self.rect.x += int(self.vx)
        self.rect.y += int(self.vy)

        # 上下墙壁反弹
        if self.rect.top <= 0 or self.rect.bottom >= HEIGHT:
            self.vy *= -1
            self.rect.top = max(0, self.rect.top)
            self.rect.bottom = min(HEIGHT, self.rect.bottom)

        # 球拍碰撞
        if self.rect.colliderect(left_paddle.rect) and self.vx < 0:
            self._bounce(left_paddle, 1)
        if self.rect.colliderect(right_paddle.rect) and self.vx > 0:
            self._bounce(right_paddle, -1)

        # 得分判定
        if self.rect.left < 0:
            right_paddle.score += 1
            return 1  # 右得分
        if self.rect.right > WIDTH:
            left_paddle.score += 1
            return -1  # 左得分
        return 0

    def _bounce(self, paddle, dir_x):
        # 根据击中位置计算角度
        hit_pos = (self.rect.centery - paddle.rect.centery) / (PADDLE_H / 2)
        hit_pos = max(-1, min(1, hit_pos))  # 限制在 -1 到 1
        
        self.speed *= 1.05  # 每次击球加速
        angle = hit_pos * (math.pi / 3)  # 最大60度角
        
        self.vx = dir_x * self.speed * math.cos(angle)
        self.vy = self.speed * math.sin(angle)
        
        # 防止球卡在球拍里
        if dir_x > 0: self.rect.left = paddle.rect.right
        else: self.rect.right = paddle.rect.left

# --- 主程序 ---

def main():
    left_paddle = Paddle(30)
    right_paddle = Paddle(WIDTH - 45, is_ai=True)
    ball = Ball()
    
    ai_enabled = True
    paused = False
    game_over = False

    while True:
        clock.tick(60)
        
        for event in pygame.event.get():
            if event.type == pygame.QUIT: pygame.quit(); sys.exit()
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_ESCAPE: pygame.quit(); sys.exit()
                if event.key == pygame.K_p: paused = not paused
                if event.key == pygame.K_r:
                    left_paddle.score = 0
                    right_paddle.score = 0
                    ball.reset()
                    game_over = False
                if event.key == pygame.K_a:
                    ai_enabled = not ai_enabled
                    right_paddle.is_ai = ai_enabled

        if not paused and not game_over:
            # 玩家1控制
            keys = pygame.key.get_pressed()
            if keys[pygame.K_w]: left_paddle.move(-PADDLE_SPEED)
            if keys[pygame.K_s]: left_paddle.move(PADDLE_SPEED)
            
            # 玩家2/AI控制
            if not ai_enabled:
                if keys[pygame.K_UP]: right_paddle.move(-PADDLE_SPEED)
                if keys[pygame.K_DOWN]: right_paddle.move(PADDLE_SPEED)
            else:
                right_paddle.ai_update(ball)

            # 更新球
            result = ball.update(left_paddle, right_paddle)
            if result != 0:
                ball.reset(-result)  # 向失分方发球
                if left_paddle.score >= WINNING_SCORE or right_paddle.score >= WINNING_SCORE:
                    game_over = True

        # --- 绘图 ---
        screen.fill(BG_COLOR)
        
        # 中线
        for y in range(0, HEIGHT, 40):
            pygame.draw.rect(screen, LINE_COLOR, (WIDTH//2 - 2, y, 4, 20))

        # 球拍 & 球
        pygame.draw.rect(screen, PADDLE_COLOR, left_paddle.rect)
        pygame.draw.rect(screen, PADDLE_COLOR, right_paddle.rect)
        pygame.draw.ellipse(screen, BALL_COLOR, ball.rect)

        # 分数
        l_txt = FONT.render(str(left_paddle.score), True, TEXT_COLOR)
        r_txt = FONT.render(str(right_paddle.score), True, TEXT_COLOR)
        screen.blit(l_txt, (WIDTH//4 - 30, 50))
        screen.blit(r_txt, (3*WIDTH//4 - 30, 50))

        # UI提示
        mode_txt = SMALL_FONT.render(f"Mode: {'AI' if ai_enabled else 'PVP'} | P:Pause | R:Reset | A:Toggle AI", True, (150,150,150))
        screen.blit(mode_txt, (20, HEIGHT - 40))

        if paused:
            txt = FONT.render("PAUSED", True, (255, 200, 0))
            screen.blit(txt, txt.get_rect(center=(WIDTH//2, HEIGHT//2)))
        
        if game_over:
            winner = "LEFT" if left_paddle.score >= WINNING_SCORE else "RIGHT"
            txt = FONT.render(f"{winner} WINS!", True, (0, 255, 100))
            sub = SMALL_FONT.render("Press R to Restart", True, TEXT_COLOR)
            screen.blit(txt, txt.get_rect(center=(WIDTH//2, HEIGHT//2 - 30)))
            screen.blit(sub, sub.get_rect(center=(WIDTH//2, HEIGHT//2 + 30)))

        pygame.display.flip()

if __name__ == "__main__":
    main()