import pygame
import sys
import math
import random

# --- 初始化 ---
pygame.init()
WIDTH, HEIGHT = 900, 600
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Pygame 足球 - Arcade Style")
clock = pygame.time.Clock()

# 颜色
FIELD_COLOR = (40, 120, 40)
LINE_COLOR = (255, 255, 255)
BALL_COLOR = (255, 255, 255)
TEAM1_COLOR = (50, 100, 200)
TEAM2_COLOR = (200, 50, 50)
GOAL_COLOR = (200, 200, 200, 100)
TEXT_COLOR = (255, 255, 255)

FONT = pygame.font.SysFont("consolas", 40)
BIG_FONT = pygame.font.SysFont("consolas", 80)

# 配置
FIELD_LEFT = 50
FIELD_RIGHT = WIDTH - 50
FIELD_TOP = 50
FIELD_BOTTOM = HEIGHT - 50
GOAL_WIDTH = 120
GOAL_DEPTH = 30
PLAYER_RADIUS = 15
BALL_RADIUS = 10
FRICTION = 0.96
MAX_SPEED = 6
SHOT_POWER = 12

# --- 游戏对象 ---

class Ball:
    def __init__(self):
        self.x = WIDTH // 2
        self.y = HEIGHT // 2
        self.vx = 0
        self.vy = 0
        self.trail = []

    def reset(self):
        self.x = WIDTH // 2
        self.y = HEIGHT // 2
        self.vx = 0
        self.vy = 0
        self.trail = []

    def update(self):
        self.x += self.vx
        self.y += self.vy
        self.vx *= FRICTION
        self.vy *= FRICTION

        # 边界
        goal_top = HEIGHT // 2 - GOAL_WIDTH // 2
        goal_bottom = HEIGHT // 2 + GOAL_WIDTH // 2

        # 上下
        if self.y - BALL_RADIUS < FIELD_TOP:
            self.y = FIELD_TOP + BALL_RADIUS
            self.vy *= -0.7
        if self.y + BALL_RADIUS > FIELD_BOTTOM:
            self.y = FIELD_BOTTOM - BALL_RADIUS
            self.vy *= -0.7

        # 左右（球门除外）
        if self.x - BALL_RADIUS < FIELD_LEFT:
            if not (goal_top <= self.y <= goal_bottom):
                self.x = FIELD_LEFT + BALL_RADIUS
                self.vx *= -0.7
        if self.x + BALL_RADIUS > FIELD_RIGHT:
            if not (goal_top <= self.y <= goal_bottom):
                self.x = FIELD_RIGHT - BALL_RADIUS
                self.vx *= -0.7

        # 拖尾
        self.trail.append((self.x, self.y))
        if len(self.trail) > 10: self.trail.pop(0)

    def draw(self, surface):
        # 拖尾
        for i, (tx, ty) in enumerate(self.trail):
            alpha = int(100 * (i / len(self.trail)))
            r = int(BALL_RADIUS * (i / len(self.trail)))
            pygame.draw.circle(surface, (200, 200, 200), (int(tx), int(ty)), max(2, r))
        
        pygame.draw.circle(surface, BALL_COLOR, (int(self.x), int(self.y)), BALL_RADIUS)
        pygame.draw.circle(surface, (100, 100, 100), (int(self.x), int(self.y)), BALL_RADIUS, 2)

class Player:
    def __init__(self, x, y, color, team, is_ai=False, is_gk=False):
        self.x = x
        self.y = y
        self.vx = 0
        self.vy = 0
        self.color = color
        self.team = team  # 1 or 2
        self.is_ai = is_ai
        self.is_gk = is_gk
        self.kick_cd = 0

    def move(self, dx, dy):
        self.vx = dx * 0.3
        self.vy = dy * 0.3
        self.x += self.vx
        self.y += self.vy
        # 限制在球场
        self.x = max(FIELD_LEFT + PLAYER_RADIUS, min(FIELD_RIGHT - PLAYER_RADIUS, self.x))
        self.y = max(FIELD_TOP + PLAYER_RADIUS, min(FIELD_BOTTOM - PLAYER_RADIUS, self.y))
        if self.kick_cd > 0: self.kick_cd -= 1

    def ai_update(self, ball, teammates, opponents):
        if not self.is_ai: return
        if self.kick_cd > 0: return

        # 门将逻辑
        if self.is_gk:
            target_y = ball.y
            target_y = max(HEIGHT//2 - GOAL_WIDTH//2 + 20, min(HEIGHT//2 + GOAL_WIDTH//2 - 20, target_y))
            target_x = FIELD_LEFT + 40 if self.team == 1 else FIELD_RIGHT - 40
            dx = target_x - self.x
            dy = target_y - self.y
            dist = math.hypot(dx, dy)
            if dist > 5:
                self.move(dx/dist * 4, dy/dist * 4)
            return

        # 普通AI：追球或跑位
        dx = ball.x - self.x
        dy = ball.y - self.y
        dist = math.hypot(dx, dy)

        # 离球最近的人追球
        min_dist = min(math.hypot(ball.x - t.x, ball.y - t.y) for t in teammates if t != self)
        if dist < min_dist or dist < 100:
            if dist > 30:
                self.move(dx/dist * 4, dy/dist * 4)
            else:
                # 踢球
                if self.kick_cd <= 0:
                    # 射门或传球
                    if (self.team == 1 and self.x > WIDTH * 0.6) or (self.team == 2 and self.x < WIDTH * 0.4):
                        # 射门
                        goal_x = FIELD_RIGHT if self.team == 1 else FIELD_LEFT
                        goal_y = HEIGHT // 2
                        angle = math.atan2(goal_y - ball.y, goal_x - ball.x)
                        ball.vx = math.cos(angle) * SHOT_POWER
                        ball.vy = math.sin(angle) * SHOT_POWER
                    else:
                        # 传球：找前方队友
                        target = None
                        min_d = float('inf')
                        for t in teammates:
                            if t == self: continue
                            d = math.hypot(t.x - self.x, t.y - self.y)
                            if (self.team == 1 and t.x > self.x) or (self.team == 2 and t.x < self.x):
                                if d < min_d:
                                    min_d = d
                                    target = t
                        if target:
                            angle = math.atan2(target.y - ball.y, target.x - ball.x)
                            ball.vx = math.cos(angle) * 8
                            ball.vy = math.sin(angle) * 8
                        else:
                            ball.vx = (1 if self.team == 1 else -1) * 6
                            ball.vy = random.uniform(-2, 2)
                    self.kick_cd = 20
        else:
            # 跑位：往对方球门方向
            target_x = WIDTH // 2 + (100 if self.team == 1 else -100)
            target_y = self.y + random.uniform(-50, 50)
            dx = target_x - self.x
            dy = target_y - self.y
            dist = math.hypot(dx, dy)
            if dist > 20:
                self.move(dx/dist * 3, dy/dist * 3)

    def draw(self, surface):
        pygame.draw.circle(surface, self.color, (int(self.x), int(self.y)), PLAYER_RADIUS)
        pygame.draw.circle(surface, (255, 255, 255), (int(self.x), int(self.y)), PLAYER_RADIUS, 2)
        if self.is_gk:
            pygame.draw.circle(surface, (255, 215, 0), (int(self.x), int(self.y)), PLAYER_RADIUS - 4)

# --- 主程序 ---

def main():
    ball = Ball()
    
    # 队伍1（左）
    team1 = [
        Player(FIELD_LEFT + 40, HEIGHT//2, TEAM1_COLOR, 1, is_ai=True, is_gk=True),
        Player(WIDTH//4, HEIGHT//3, TEAM1_COLOR, 1, is_ai=True),
        Player(WIDTH//4, HEIGHT//2, TEAM1_COLOR, 1, is_ai=False),  # 玩家
        Player(WIDTH//4, HEIGHT*2//3, TEAM1_COLOR, 1, is_ai=True),
    ]
    
    # 队伍2（右）
    team2 = [
        Player(FIELD_RIGHT - 40, HEIGHT//2, TEAM2_COLOR, 2, is_ai=True, is_gk=True),
        Player(WIDTH*3//4, HEIGHT//3, TEAM2_COLOR, 2, is_ai=True),
        Player(WIDTH*3//4, HEIGHT//2, TEAM2_COLOR, 2, is_ai=True),
        Player(WIDTH*3//4, HEIGHT*2//3, TEAM2_COLOR, 2, is_ai=True),
    ]

    score1 = 0
    score2 = 0
    goal_flash = 0
    goal_msg = ""

    running = True
    while running:
        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:
                    score1 = 0
                    score2 = 0
                    ball.reset()
                    for p in team1 + team2:
                        p.x = WIDTH//4 if p.team == 1 else WIDTH*3//4
                        p.y = HEIGHT//2
                if event.key == pygame.K_SPACE:
                    # 玩家踢球
                    player = team1[2]
                    if player.kick_cd <= 0:
                        dist = math.hypot(ball.x - player.x, ball.y - player.y)
                        if dist < PLAYER_RADIUS + BALL_RADIUS + 10:
                            # 射门或传球
                            if player.x > WIDTH * 0.6:
                                # 射门
                                goal_x = FIELD_RIGHT
                                goal_y = HEIGHT // 2
                                angle = math.atan2(goal_y - ball.y, goal_x - ball.x)
                                ball.vx = math.cos(angle) * SHOT_POWER
                                ball.vy = math.sin(angle) * SHOT_POWER
                                goal_flash = 30
                            else:
                                # 传球：找最近前方队友
                                target = None
                                min_d = float('inf')
                                for t in team1:
                                    if t == player: continue
                                    d = math.hypot(t.x - player.x, t.y - player.y)
                                    if t.x > player.x and d < min_d:
                                        min_d = d
                                        target = t
                                if target:
                                    angle = math.atan2(target.y - ball.y, target.x - ball.x)
                                    ball.vx = math.cos(angle) * 8
                                    ball.vy = math.sin(angle) * 8
                                else:
                                    ball.vx = 6
                                    ball.vy = random.uniform(-1, 1)
                            player.kick_cd = 20

        # 更新
        keys = pygame.key.get_pressed()
        dx, dy = 0, 0
        if keys[pygame.K_w]: dy -= 5
        if keys[pygame.K_s]: dy += 5
        if keys[pygame.K_a]: dx -= 5
        if keys[pygame.K_d]: dx += 5
        team1[2].move(dx, dy)

        for p in team1: p.ai_update(ball, team1, team2)
        for p in team2: p.ai_update(ball, team2, team1)
        ball.update()

        # 进球判定
        goal_top = HEIGHT // 2 - GOAL_WIDTH // 2
        goal_bottom = HEIGHT // 2 + GOAL_WIDTH // 2
        if ball.x - BALL_RADIUS < FIELD_LEFT and goal_top <= ball.y <= goal_bottom:
            score2 += 1
            goal_msg = "GOAL! P2"
            goal_flash = 60
            ball.reset()
        if ball.x + BALL_RADIUS > FIELD_RIGHT and goal_top <= ball.y <= goal_bottom:
            score1 += 1
            goal_msg = "GOAL! P1"
            goal_flash = 60
            ball.reset()

        if goal_flash > 0: goal_flash -= 1

        # --- 绘图 ---
        screen.fill(FIELD_COLOR)
        
        # 球场线
        pygame.draw.rect(screen, LINE_COLOR, (FIELD_LEFT, FIELD_TOP, FIELD_RIGHT-FIELD_LEFT, FIELD_BOTTOM-FIELD_TOP), 3)
        pygame.draw.line(screen, LINE_COLOR, (WIDTH//2, FIELD_TOP), (WIDTH//2, FIELD_BOTTOM), 3)
        pygame.draw.circle(screen, LINE_COLOR, (WIDTH//2, HEIGHT//2), 60, 3)
        
        # 球门
        pygame.draw.rect(screen, GOAL_COLOR, (FIELD_LEFT - GOAL_DEPTH, goal_top, GOAL_DEPTH, GOAL_WIDTH))
        pygame.draw.rect(screen, GOAL_COLOR, (FIELD_RIGHT, goal_top, GOAL_DEPTH, GOAL_WIDTH))

        # 球员 & 球
        for p in team1: p.draw(screen)
        for p in team2: p.draw(screen)
        ball.draw(screen)

        # UI
        score_txt = FONT.render(f"P1: {score1} | P2: {score2}", True, TEXT_COLOR)
        screen.blit(score_txt, (WIDTH//2 - 100, 10))
        hint = FONT.render("WASD: Move | SPACE: Kick | R: Reset", True, (200, 200, 200))
        screen.blit(hint, (WIDTH//2 - 200, HEIGHT - 40))

        # 进球特效
        if goal_flash > 0:
            overlay = pygame.Surface((WIDTH, HEIGHT), pygame.SRCALPHA)
            overlay.fill((255, 255, 255, goal_flash * 4))
            screen.blit(overlay, (0, 0))
            txt = BIG_FONT.render(goal_msg, True, (255, 215, 0))
            screen.blit(txt, txt.get_rect(center=(WIDTH//2, HEIGHT//2)))

        pygame.display.flip()

    pygame.quit()

if __name__ == "__main__":
    main()