import pygame
import sys
import math
import random

# ========== 初始化 ==========
pygame.init()
WIDTH, HEIGHT = 900, 550
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("🎱 台球游戏 – 轻松版")
clock = pygame.time.Clock()

# ========== 高级色彩方案 ==========
BG_COLOR = (15, 15, 25)
TABLE_EDGE = (60, 35, 20)
TABLE_WOOD = (120, 70, 35)
TABLE_FELT = (34, 139, 34)
RAIL_WOOD = (160, 100, 50)
POCKET_BLACK = (10, 10, 10)
POCKET_LEATHER = (80, 50, 30)
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
GOLD = (255, 215, 0)
RED = (220, 20, 60)
YELLOW = (255, 215, 0)
BLUE = (100, 149, 237)
PURPLE = (160, 32, 240)
ORANGE = (255, 165, 0)
MAROON = (128, 0, 0)
SHADOW_COLOR = (0, 0, 0, 80)

# ========== 字体 ==========
font_small = pygame.font.Font(None, 28)
font_med = pygame.font.Font(None, 36)
font_large = pygame.font.Font(None, 48)

# ========== 台球桌参数 ==========
TABLE_LEFT = 150
TABLE_TOP = 50
TABLE_WIDTH = 600
TABLE_HEIGHT = 300
TABLE_RIGHT = TABLE_LEFT + TABLE_WIDTH
TABLE_BOTTOM = TABLE_TOP + TABLE_HEIGHT

# 增大袋口半径，更容易进球
POCKET_RADIUS = 22
pockets = [
    (TABLE_LEFT + 4, TABLE_TOP + 4),
    (TABLE_RIGHT - 4, TABLE_TOP + 4),
    (TABLE_LEFT + 4, TABLE_BOTTOM - 4),
    (TABLE_RIGHT - 4, TABLE_BOTTOM - 4),
    (TABLE_LEFT + TABLE_WIDTH//2, TABLE_TOP + 2),
    (TABLE_LEFT + TABLE_WIDTH//2, TABLE_BOTTOM - 2)
]

BALL_RADIUS = 11
FRICTION = 0.99            # 降低摩擦，球走得更远
MIN_SPEED = 0.05
CUE_BALL_START = (TABLE_LEFT + TABLE_WIDTH // 4, TABLE_TOP + TABLE_HEIGHT // 2)

# 球体颜色映射
ball_colors = {
    0: (240, 240, 240),
    1: (255, 200, 50), 2: (50, 100, 220), 3: (220, 50, 50),
    4: (150, 50, 220), 5: (240, 130, 50), 6: (130, 50, 50),
    7: (160, 60, 20), 8: (20, 20, 20), 9: (255, 210, 50),
    10: (60, 120, 240), 11: (230, 60, 60), 12: (160, 60, 230),
    13: (240, 140, 60), 14: (140, 60, 60), 15: (60, 130, 60)
}

# ========== Ball类 ==========
class Ball:
    def __init__(self, number, x, y):
        self.number = number
        self.x = x
        self.y = y
        self.vx = 0.0
        self.vy = 0.0
        self.radius = BALL_RADIUS
        self.pocketed = False
        self.color = ball_colors.get(number, WHITE)

    def update(self):
        if self.pocketed: return
        self.x += self.vx
        self.y += self.vy
        self.vx *= FRICTION
        self.vy *= FRICTION
        if abs(self.vx) < MIN_SPEED and abs(self.vy) < MIN_SPEED:
            self.vx = 0
            self.vy = 0

    def is_moving(self):
        return abs(self.vx) > 0.01 or abs(self.vy) > 0.01

    def draw(self, surface):
        if self.pocketed: return
        pygame.draw.circle(surface, SHADOW_COLOR, (int(self.x + 2), int(self.y + 2)), self.radius)
        pygame.draw.circle(surface, self.color, (int(self.x), int(self.y)), self.radius)
        if self.number != 0:
            num_bg_r = self.radius - 3
            pygame.draw.circle(surface, WHITE, (int(self.x), int(self.y)), num_bg_r)
            num_text = font_small.render(str(self.number), True, BLACK)
            num_rect = num_text.get_rect(center=(int(self.x), int(self.y)))
            surface.blit(num_text, num_rect)
        else:
            pygame.draw.circle(surface, BLACK, (int(self.x), int(self.y)), self.radius - 4, 2)
        highlight_x = int(self.x - self.radius * 0.3)
        highlight_y = int(self.y - self.radius * 0.3)
        highlight_r = int(self.radius * 0.3)
        pygame.draw.circle(surface, WHITE, (highlight_x, highlight_y), highlight_r)

# ========== 游戏状态 ==========
class PoolGame:
    def __init__(self):
        self.balls = []
        self.cue_ball = None
        self.reset_balls()
        self.aiming = False
        self.aim_start = (0, 0)
        self.power = 0.0
        self.max_power = 25.0
        self.state = "aiming"
        self.wait_for_stop = False

    def reset_balls(self):
        self.balls = []
        cx, cy = CUE_BALL_START
        self.cue_ball = Ball(0, cx, cy)
        self.balls.append(self.cue_ball)
        start_x = TABLE_LEFT + TABLE_WIDTH * 3 // 4
        start_y = TABLE_TOP + TABLE_HEIGHT // 2
        numbers = [1,2,3,4,5,6,7,8,9,10,11,12,13,14,15]
        random.shuffle(numbers)
        idx = 0
        for row in range(5):
            for col in range(row + 1):
                x = start_x + row * (BALL_RADIUS * 2 + 2)
                y = start_y - row * (BALL_RADIUS + 1) + col * (BALL_RADIUS * 2 + 2)
                self.balls.append(Ball(numbers[idx], x, y))
                idx += 1

    def all_balls_stopped(self):
        for b in self.balls:
            if b.is_moving(): return False
        return True

    def handle_event(self, event):
        if event.type == pygame.MOUSEBUTTONDOWN:
            if self.state == "aiming" and not self.wait_for_stop:
                self.aim_start = pygame.mouse.get_pos()
                self.aiming = True
            elif self.state == "idle" and not self.wait_for_stop:
                if self.cue_ball.pocketed:
                    self.cue_ball.x, self.cue_ball.y = CUE_BALL_START
                    self.cue_ball.vx = 0
                    self.cue_ball.vy = 0
                    self.cue_ball.pocketed = False
                    self.state = "aiming"
        elif event.type == pygame.MOUSEBUTTONUP:
            if self.aiming and not self.wait_for_stop:
                mouse_pos = pygame.mouse.get_pos()
                dx = self.aim_start[0] - mouse_pos[0]
                dy = self.aim_start[1] - mouse_pos[1]
                dist = math.hypot(dx, dy)
                if dist > 5:
                    self.power = min(dist * 0.2, self.max_power)
                    angle = math.atan2(-dy, -dx)
                    self.cue_ball.vx = math.cos(angle) * self.power
                    self.cue_ball.vy = math.sin(angle) * self.power
                    self.state = "shooting"
                    self.wait_for_stop = True
                self.aiming = False

    def update(self):
        if self.wait_for_stop:
            if self.all_balls_stopped():
                self.wait_for_stop = False
                self.state = "idle"
                if not self.cue_ball.pocketed:
                    self.state = "aiming"
            else:
                for ball in self.balls:
                    ball.update()
                self.check_collisions()
                self.check_pockets()

    def check_collisions(self):
        for i in range(len(self.balls)):
            if self.balls[i].pocketed: continue
            for j in range(i+1, len(self.balls)):
                if self.balls[j].pocketed: continue
                b1, b2 = self.balls[i], self.balls[j]
                dx = b1.x - b2.x
                dy = b1.y - b2.y
                dist = math.hypot(dx, dy)
                min_dist = b1.radius + b2.radius
                if dist < min_dist and dist > 0:
                    overlap = min_dist - dist
                    angle = math.atan2(dy, dx)
                    b1.x += math.cos(angle) * overlap / 2
                    b1.y += math.sin(angle) * overlap / 2
                    b2.x -= math.cos(angle) * overlap / 2
                    b2.y -= math.sin(angle) * overlap / 2
                    nx, ny = dx / dist, dy / dist
                    dvx = b1.vx - b2.vx
                    dvy = b1.vy - b2.vy
                    dot = dvx * nx + dvy * ny
                    if dot > 0:
                        b1.vx -= dot * nx
                        b1.vy -= dot * ny
                        b2.vx += dot * nx
                        b2.vy += dot * ny

        for ball in self.balls:
            if ball.pocketed: continue
            if ball.x - ball.radius < TABLE_LEFT + 5:
                ball.x = TABLE_LEFT + 5 + ball.radius
                ball.vx *= -0.8
            elif ball.x + ball.radius > TABLE_RIGHT - 5:
                ball.x = TABLE_RIGHT - 5 - ball.radius
                ball.vx *= -0.8
            if ball.y - ball.radius < TABLE_TOP + 5:
                ball.y = TABLE_TOP + 5 + ball.radius
                ball.vy *= -0.8
            elif ball.y + ball.radius > TABLE_BOTTOM - 5:
                ball.y = TABLE_BOTTOM - 5 - ball.radius
                ball.vy *= -0.8

    def check_pockets(self):
        for ball in self.balls:
            if ball.pocketed: continue
            for px, py in pockets:
                if math.hypot(ball.x - px, ball.y - py) < POCKET_RADIUS - 3:
                    ball.pocketed = True
                    ball.vx = 0
                    ball.vy = 0
                    break

    def draw_table(self):
        outer_rect = pygame.Rect(TABLE_LEFT - 20, TABLE_TOP - 20, TABLE_WIDTH + 40, TABLE_HEIGHT + 40)
        pygame.draw.rect(screen, TABLE_EDGE, outer_rect, border_radius=15)
        border_rect = pygame.Rect(TABLE_LEFT - 8, TABLE_TOP - 8, TABLE_WIDTH + 16, TABLE_HEIGHT + 16)
        pygame.draw.rect(screen, RAIL_WOOD, border_rect, border_radius=10)
        felt_rect = pygame.Rect(TABLE_LEFT, TABLE_TOP, TABLE_WIDTH, TABLE_HEIGHT)
        pygame.draw.rect(screen, TABLE_FELT, felt_rect)
        for x in range(TABLE_LEFT, TABLE_RIGHT, 20):
            pygame.draw.line(screen, (30, 120, 30), (x, TABLE_TOP), (x, TABLE_BOTTOM), 1)
        for y in range(TABLE_TOP, TABLE_BOTTOM, 20):
            pygame.draw.line(screen, (30, 120, 30), (TABLE_LEFT, y), (TABLE_RIGHT, y), 1)
        inner_rect = felt_rect.inflate(-8, -8)
        pygame.draw.rect(screen, RAIL_WOOD, inner_rect, 3)
        for px, py in pockets:
            pygame.draw.circle(screen, POCKET_BLACK, (px, py), POCKET_RADIUS)
            pygame.draw.circle(screen, POCKET_LEATHER, (px, py), POCKET_RADIUS, 2)
        head_line_x = TABLE_LEFT + TABLE_WIDTH // 4
        pygame.draw.line(screen, WHITE, (head_line_x, TABLE_TOP+3), (head_line_x, TABLE_BOTTOM-3), 1)
        center_x = TABLE_LEFT + TABLE_WIDTH // 2
        center_y = TABLE_TOP + TABLE_HEIGHT // 2
        pygame.draw.circle(screen, WHITE, (center_x, center_y), 3)
        foot_x = TABLE_LEFT + TABLE_WIDTH * 3 // 4
        pygame.draw.circle(screen, WHITE, (foot_x, center_y), 3)

    def draw_aim(self):
        if self.aiming and not self.wait_for_stop and self.state == "aiming":
            mouse_pos = pygame.mouse.get_pos()
            cx, cy = self.cue_ball.x, self.cue_ball.y
            # 绘制一条从白球出发穿过鼠标的瞄准线，并延伸出去显示预计撞击方向
            dx = mouse_pos[0] - cx
            dy = mouse_pos[1] - cy
            dist = math.hypot(dx, dy)
            if dist > 0:
                # 绘制向后拉的力量线（白色虚线）
                pull_x = cx - dx
                pull_y = cy - dy
                pygame.draw.line(screen, (255, 255, 200), (cx, cy), (pull_x, pull_y), 2)
                # 绘制向前的瞄准线（延伸出去）
                extend = 300
                ex = cx + (dx / dist) * extend
                ey = cy + (dy / dist) * extend
                pygame.draw.line(screen, (255, 255, 200, 100), (cx, cy), (ex, ey), 1)
            # 力量条
            power = min(dist * 0.2, self.max_power)
            bar_x = mouse_pos[0] + 15
            bar_y = mouse_pos[1] - 30
            pygame.draw.rect(screen, WHITE, (bar_x, bar_y, 8, 60), 2)
            fill_h = int((power / self.max_power) * 60)
            pygame.draw.rect(screen, RED, (bar_x, bar_y + 60 - fill_h, 8, fill_h))

    def draw_ui(self):
        remaining = sum(1 for b in self.balls if not b.pocketed and b.number != 0)
        text = font_med.render(f"剩余球: {remaining}", True, WHITE)
        screen.blit(text, (20, 20))
        if self.state == "aiming":
            hint = font_small.render("从白球拖动设置力度与方向", True, GOLD)
            screen.blit(hint, (20, HEIGHT - 30))
        elif self.cue_ball.pocketed:
            hint = font_small.render("点击放置白球", True, GOLD)
            screen.blit(hint, (20, HEIGHT - 30))

    def draw(self):
        screen.fill(BG_COLOR)
        self.draw_table()
        for ball in self.balls:
            ball.draw(screen)
        self.draw_aim()
        self.draw_ui()

    def run(self):
        running = True
        while running:
            dt = clock.tick(60)
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    running = False
                elif event.type == pygame.KEYDOWN:
                    if event.key == pygame.K_r:
                        self.reset_balls()
                        self.state = "aiming"
                        self.wait_for_stop = False
                else:
                    self.handle_event(event)
            self.update()
            self.draw()
            pygame.display.flip()
        pygame.quit()
        sys.exit()

if __name__ == "__main__":
    game = PoolGame()
    game.run()