import pygame
import sys
import math
import random

# --- 初始化 ---
pygame.init()

# --- 常量设置 ---
WIDTH, HEIGHT = 800, 600
SCREEN = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("🏒 简易冰球游戏")

CLOCK = pygame.time.Clock()
FPS = 60

# 颜色定义 (R, G, B)
COLOR_BG = (20, 20, 30)          # 深蓝色背景
COLOR_ICE = (220, 235, 245)      # 冰面颜色
COLOR_LINE = (180, 200, 220)     # 冰场线
COLOR_PUCK = (20, 20, 20)        # 冰球黑色
COLOR_PLAYER = (0, 100, 255)     # 玩家蓝色
COLOR_AI = (255, 50, 50)         # AI 红色
COLOR_TEXT = (255, 255, 255)

# 游戏对象参数
PADDLE_RADIUS = 30
PUCK_RADIUS = 15
PADDLE_SPEED = 6
FRICTION = 0.98  # 冰面摩擦力，让球慢慢停下

# --- 游戏对象类 ---
class Paddle:
    def __init__(self, x, y, color, is_ai=False):
        self.x = x
        self.y = y
        self.color = color
        self.is_ai = is_ai
        self.radius = PADDLE_RADIUS
        self.speed = PADDLE_SPEED

    def move(self, dx, dy):
        """限制移动范围，防止出界"""
        self.x += dx
        self.y += dy
        self.x = max(self.radius, min(WIDTH - self.radius, self.x))
        self.y = max(self.radius, min(HEIGHT - self.radius, self.y))

    def update_ai(self, puck):
        """简单的 AI 逻辑：追踪冰球，但只在自己半场活动"""
        if self.is_ai:
            # 只有球在 AI 半场或向 AI 移动时才反应
            if puck.x > WIDTH // 2 or puck.vx > 0:
                dx = puck.x - self.x
                dy = puck.y - self.y
                dist = math.hypot(dx, dy)
                if dist > 1:
                    dx /= dist
                    dy /= dist
                    # AI 速度稍慢，给玩家机会
                    self.move(dx * self.speed * 0.8, dy * self.speed * 0.8)

    def draw(self, surface):
        pygame.draw.circle(surface, self.color, (int(self.x), int(self.y)), self.radius)
        # 画个内圈增加立体感
        pygame.draw.circle(surface, (255, 255, 255), (int(self.x), int(self.y)), self.radius - 5, 2)

class Puck:
    def __init__(self):
        self.reset()

    def reset(self):
        self.x = WIDTH // 2
        self.y = HEIGHT // 2
        self.vx = random.choice([-3, 3])  # 随机向左或向右发球
        self.vy = random.uniform(-2, 2)
        self.radius = PUCK_RADIUS

    def update(self):
        self.x += self.vx
        self.y += self.vy
        
        # 冰面摩擦
        self.vx *= FRICTION
        self.vy *= FRICTION

        # 上下墙壁反弹
        if self.y - self.radius <= 0 or self.y + self.radius >= HEIGHT:
            self.vy *= -1
            # 防止卡在墙里
            self.y = max(self.radius, min(HEIGHT - self.radius, self.y))

        # 左右边界（进球判定）
        if self.x - self.radius <= 0:
            return "AI_SCORE"
        if self.x + self.radius >= WIDTH:
            return "PLAYER_SCORE"
            
        return None

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

# --- 核心逻辑函数 ---
def check_collision(paddle, puck):
    """圆与圆的碰撞检测"""
    dx = puck.x - paddle.x
    dy = puck.y - paddle.y
    distance = math.hypot(dx, dy)
    
    if distance < paddle.radius + puck.radius:
        # 1. 分离物体，防止粘连
        overlap = (paddle.radius + puck.radius) - distance
        angle = math.atan2(dy, dx)
        puck.x += math.cos(angle) * overlap
        puck.y += math.sin(angle) * overlap
        
        # 2. 物理反弹：基于碰撞角度和球拍速度
        # 简单的弹性碰撞模型
        puck.vx = math.cos(angle) * 8 + paddle.speed * 0.5 * (1 if paddle.x < WIDTH//2 else -1)
        puck.vy = math.sin(angle) * 8
        
        return True
    return False

def draw_ice_rink(surface):
    """绘制冰场背景"""
    surface.fill(COLOR_BG)
    # 冰面
    pygame.draw.rect(surface, COLOR_ICE, (20, 20, WIDTH-40, HEIGHT-40), border_radius=20)
    # 中线
    pygame.draw.line(surface, COLOR_LINE, (WIDTH//2, 20), (WIDTH//2, HEIGHT-20), 3)
    # 中圈
    pygame.draw.circle(surface, COLOR_LINE, (WIDTH//2, HEIGHT//2), 60, 3)
    # 球门线
    pygame.draw.line(surface, COLOR_AI, (20, HEIGHT//2 - 50), (20, HEIGHT//2 + 50), 5)
    pygame.draw.line(surface, COLOR_PLAYER, (WIDTH-20, HEIGHT//2 - 50), (WIDTH-20, HEIGHT//2 + 50), 5)

# --- 主游戏循环 ---
def main():
    player = Paddle(WIDTH - 100, HEIGHT // 2, COLOR_PLAYER)
    ai = Paddle(100, HEIGHT // 2, COLOR_AI, is_ai=True)
    puck = Puck()
    
    score_player = 0
    score_ai = 0
    font = pygame.font.SysFont("Arial", 40)

    while True:
        CLOCK.tick(FPS)
        
        # 1. 事件处理
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()

        # 2. 玩家控制 (WASD 或 方向键)
        keys = pygame.key.get_pressed()
        dx, dy = 0, 0
        if keys[pygame.K_LEFT] or keys[pygame.K_a]: dx -= 1
        if keys[pygame.K_RIGHT] or keys[pygame.K_d]: dx += 1
        if keys[pygame.K_UP] or keys[pygame.K_w]: dy -= 1
        if keys[pygame.K_DOWN] or keys[pygame.K_s]: dy += 1
        
        # 归一化斜向移动速度
        if dx != 0 or dy != 0:
            length = math.hypot(dx, dy)
            dx /= length
            dy /= length
            
        player.move(dx * player.speed, dy * player.speed)
        ai.update_ai(puck)

        # 3. 更新物理
        puck.update()
        
        # 碰撞检测
        check_collision(player, puck)
        check_collision(ai, puck)

        # 4. 进球判定
        goal_status = puck.x  # 简单判断，实际用 puck.update 返回值更好
        if puck.x < 0:
            score_ai += 1
            puck.reset()
        elif puck.x > WIDTH:
            score_player += 1
            puck.reset()

        # 5. 绘制
        draw_ice_rink(SCREEN)
        puck.draw(SCREEN)
        player.draw(SCREEN)
        ai.draw(SCREEN)
        
        # 计分板
        text = font.render(f"{score_ai}  -  {score_player}", True, COLOR_TEXT)
        SCREEN.blit(text, (WIDTH//2 - text.get_width()//2, 30))

        pygame.display.flip()

if __name__ == "__main__":
    main()