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 冰球 - Arcade Style")
clock = pygame.time.Clock()

# 颜色
ICE_COLOR = (220, 235, 245)
LINE_COLOR = (180, 200, 220)
GOAL_COLOR = (200, 50, 50)
PADDLE1_COLOR = (50, 100, 200)
PADDLE2_COLOR = (200, 50, 50)
PUCK_COLOR = (30, 30, 30)
TEXT_COLOR = (50, 50, 50)

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

# 配置
PADDLE_RADIUS = 30
PUCK_RADIUS = 15
GOAL_WIDTH = 120
GOAL_DEPTH = 20
WINNING_SCORE = 10
FRICTION = 0.985  # 冰面摩擦力
BOUNCE_DAMPING = 0.8  # 撞墙能量损耗
MAX_PUCK_SPEED = 20

# --- 游戏对象 ---

class Paddle:
    def __init__(self, x, color):
        self.x = x
        self.y = HEIGHT // 2
        self.color = color
        self.vx = 0
        self.vy = 0
        self.prev_x = x
        self.prev_y = HEIGHT // 2

    def move(self, dx, dy):
        self.prev_x = self.x
        self.prev_y = self.y
        self.x += dx
        self.y += dy
        
        # 边界限制（不能出界，不能进自己球门）
        self.x = max(PADDLE_RADIUS, min(WIDTH - PADDLE_RADIUS, self.x))
        self.y = max(PADDLE_RADIUS, min(HEIGHT - PADDLE_RADIUS, self.y))
        
        # 计算速度（用于击球力度）
        self.vx = self.x - self.prev_x
        self.vy = self.y - self.prev_y

    def draw(self, surface):
        pygame.draw.circle(surface, self.color, (int(self.x), int(self.y)), PADDLE_RADIUS)
        pygame.draw.circle(surface, (255, 255, 255), (int(self.x), int(self.y)), PADDLE_RADIUS, 3)

class Puck:
    def __init__(self):
        self.x = WIDTH // 2
        self.y = HEIGHT // 2
        self.vx = 0
        self.vy = 0
        self.reset()

    def reset(self, direction=None):
        self.x = WIDTH // 2
        self.y = HEIGHT // 2
        self.vx = random.uniform(-3, 3)
        self.vy = random.uniform(-2, 2)
        if direction:
            self.vx = direction * 4

    def update(self):
        self.x += self.vx
        self.y += self.vy
        
        # 摩擦力
        self.vx *= FRICTION
        self.vy *= FRICTION
        
        # 限速
        speed = math.hypot(self.vx, self.vy)
        if speed > MAX_PUCK_SPEED:
            self.vx = (self.vx / speed) * MAX_PUCK_SPEED
            self.vy = (self.vy / speed) * MAX_PUCK_SPEED

        # 上下墙壁反弹
        if self.y - PUCK_RADIUS <= 0:
            self.y = PUCK_RADIUS
            self.vy *= -BOUNCE_DAMPING
        if self.y + PUCK_RADIUS >= HEIGHT:
            self.y = HEIGHT - PUCK_RADIUS
            self.vy *= -BOUNCE_DAMPING

        # 左右墙壁（球门区域除外）
        goal_top = HEIGHT // 2 - GOAL_WIDTH // 2
        goal_bottom = HEIGHT // 2 + GOAL_WIDTH // 2
        
        # 左墙
        if self.x - PUCK_RADIUS <= 0:
            if not (goal_top <= self.y <= goal_bottom):
                self.x = PUCK_RADIUS
                self.vx *= -BOUNCE_DAMPING
        # 右墙
        if self.x + PUCK_RADIUS >= WIDTH:
            if not (goal_top <= self.y <= goal_bottom):
                self.x = WIDTH - PUCK_RADIUS
                self.vx *= -BOUNCE_DAMPING

    def draw(self, surface):
        pygame.draw.circle(surface, PUCK_COLOR, (int(self.x), int(self.y)), PUCK_RADIUS)
        pygame.draw.circle(surface, (100, 100, 100), (int(self.x), int(self.y)), PUCK_RADIUS, 2)

# --- 碰撞检测 ---
def circle_collision(c1_x, c1_y, c1_r, c2_x, c2_y, c2_r):
    dist = math.hypot(c1_x - c2_x, c1_y - c2_y)
    return dist < c1_r + c2_r

def resolve_paddle_puck(paddle, puck):
    if not circle_collision(paddle.x, paddle.y, PADDLE_RADIUS, puck.x, puck.y, PUCK_RADIUS):
        return False
    
    # 计算碰撞法线
    dx = puck.x - paddle.x
    dy = puck.y - paddle.y
    dist = math.hypot(dx, dy)
    if dist == 0: return False
    
    nx = dx / dist
    ny = dy / dist
    
    # 分离球拍和球
    overlap = (PADDLE_RADIUS + PUCK_RADIUS) - dist
    puck.x += nx * overlap
    puck.y += ny * overlap
    
    # 相对速度
    rel_vx = puck.vx - paddle.vx
    rel_vy = puck.vy - paddle.vy
    
    # 法线方向相对速度
    rel_vn = rel_vx * nx + rel_vy * ny
    
    # 如果正在分离，不处理
    if rel_vn > 0: return False
    
    # 弹性碰撞（球拍质量大，简化处理）
    impulse = -2 * rel_vn
    puck.vx += impulse * nx
    puck.vy += impulse * ny
    
    # 加上球拍速度（击球力度）
    puck.vx += paddle.vx * 0.5
    puck.vy += paddle.vy * 0.5
    
    return True

# --- 主程序 ---

def main():
    p1 = Paddle(100, PADDLE1_COLOR)
    p2 = Paddle(WIDTH - 100, PADDLE2_COLOR)
    puck = Puck()
    
    score1 = 0
    score2 = 0
    game_over = False
    goal_flash = 0  # 进球闪光计时

    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_r and game_over:
                    score1 = 0
                    score2 = 0
                    game_over = False
                    puck.reset()

        if not game_over:
            # 玩家1控制
            keys = pygame.key.get_pressed()
            dx1, dy1 = 0, 0
            if keys[pygame.K_w]: dy1 -= 6
            if keys[pygame.K_s]: dy1 += 6
            if keys[pygame.K_a]: dx1 -= 6
            if keys[pygame.K_d]: dx1 += 6
            p1.move(dx1, dy1)
            
            # 玩家2控制
            dx2, dy2 = 0, 0
            if keys[pygame.K_UP]: dy2 -= 6
            if keys[pygame.K_DOWN]: dy2 += 6
            if keys[pygame.K_LEFT]: dx2 -= 6
            if keys[pygame.K_RIGHT]: dx2 += 6
            p2.move(dx2, dy2)

            # 更新球
            puck.update()
            
            # 碰撞检测
            resolve_paddle_puck(p1, puck)
            resolve_paddle_puck(p2, puck)
            
            # 进球判定
            goal_top = HEIGHT // 2 - GOAL_WIDTH // 2
            goal_bottom = HEIGHT // 2 + GOAL_WIDTH // 2
            
            if puck.x - PUCK_RADIUS <= 0 and goal_top <= puck.y <= goal_bottom:
                score2 += 1
                goal_flash = 30
                if score2 >= WINNING_SCORE: game_over = True
                else: puck.reset(1)
                
            if puck.x + PUCK_RADIUS >= WIDTH and goal_top <= puck.y <= goal_bottom:
                score1 += 1
                goal_flash = 30
                if score1 >= WINNING_SCORE: game_over = True
                else: puck.reset(-1)

        if goal_flash > 0: goal_flash -= 1

        # --- 绘图 ---
        screen.fill(ICE_COLOR)
        
        # 中线
        pygame.draw.line(screen, LINE_COLOR, (WIDTH//2, 0), (WIDTH//2, HEIGHT), 3)
        # 中圈
        pygame.draw.circle(screen, LINE_COLOR, (WIDTH//2, HEIGHT//2), 60, 3)
        
        # 球门
        goal_top = HEIGHT // 2 - GOAL_WIDTH // 2
        goal_bottom = HEIGHT // 2 + GOAL_WIDTH // 2
        pygame.draw.rect(screen, GOAL_COLOR, (0, goal_top, GOAL_DEPTH, GOAL_WIDTH))
        pygame.draw.rect(screen, GOAL_COLOR, (WIDTH - GOAL_DEPTH, goal_top, GOAL_DEPTH, GOAL_WIDTH))
        
        # 进球闪光
        if goal_flash > 0:
            overlay = pygame.Surface((WIDTH, HEIGHT), pygame.SRCALPHA)
            overlay.fill((255, 255, 255, goal_flash * 8))
            screen.blit(overlay, (0, 0))

        # 球拍 & 球
        p1.draw(screen)
        p2.draw(screen)
        puck.draw(screen)

        # 分数
        s1_txt = FONT.render(str(score1), True, PADDLE1_COLOR)
        s2_txt = FONT.render(str(score2), True, PADDLE2_COLOR)
        screen.blit(s1_txt, (WIDTH//4 - 30, 30))
        screen.blit(s2_txt, (3*WIDTH//4 - 30, 30))

        # 提示
        hint = SMALL_FONT.render("WASD / Arrows | R: Restart", True, (150, 150, 150))
        screen.blit(hint, (WIDTH//2 - 120, HEIGHT - 40))

        if game_over:
            winner = "BLUE" if score1 >= WINNING_SCORE else "RED"
            txt = FONT.render(f"{winner} WINS!", True, (50, 50, 50))
            sub = SMALL_FONT.render("Press R to Restart", True, (100, 100, 100))
            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()