import pygame
import sys
import random

# ---------------- 基础设置 ----------------
WIDTH, HEIGHT = 960, 600
FPS = 60
GROUND_Y = 500
GRAVITY = 0.85

pygame.init()

screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("森林冰火人大战 Forest IceFire Battle")

clock = pygame.time.Clock()
font = pygame.font.Font(None, 32)
big_font = pygame.font.Font(None, 72)

WHITE = (255, 255, 255)
BLACK = (15, 18, 25)
GREEN = (60, 140, 80)
DARK_GREEN = (30, 70, 40)
BROWN = (90, 60, 40)
GRAY = (120, 120, 120)

FIRE_COLOR = (240, 90, 50)
FIRE_LIGHT = (255, 170, 70)
ICE_COLOR = (90, 180, 240)
ICE_LIGHT = (200, 235, 255)


# ---------------- 工具函数 ----------------
def draw_text_center(text, font_obj, color, x, y):
    rendered = font_obj.render(text, True, color)
    screen.blit(rendered, (x - rendered.get_width() // 2, y))


# ---------------- 森林背景 ----------------
class ForestBackground:
    def __init__(self):
        self.trees = []

        for _ in range(14):
            self.trees.append({
                "x": random.randint(20, WIDTH - 20),
                "h": random.randint(110, 210),
                "w": random.randint(70, 130),
                "trunk": random.randint(14, 22),
                "shade": random.randint(-20, 20)
            })

        self.trees.sort(key=lambda t: t["h"])

    def draw(self):
        # 天空渐变
        for y in range(GROUND_Y):
            ratio = y / GROUND_Y
            r = int(20 + 30 * ratio)
            g = int(35 + 55 * ratio)
            b = int(55 + 40 * ratio)
            pygame.draw.line(screen, (r, g, b), (0, y), (WIDTH, y))

        # 远景树林
        for tree in self.trees:
            x = tree["x"]
            h = tree["h"]
            w = tree["w"]
            trunk = tree["trunk"]
            shade = tree["shade"]

            trunk_color = (
                max(40, min(120, BROWN[0] + shade)),
                max(30, min(90, BROWN[1] + shade)),
                max(20, min(70, BROWN[2] + shade))
            )

            leaf_color = (
                max(20, min(120, GREEN[0] + shade)),
                max(60, min(180, GREEN[1] + shade)),
                max(30, min(120, GREEN[2] + shade))
            )

            pygame.draw.rect(screen, trunk_color, (x - trunk // 2, GROUND_Y - h, trunk, h))

            points = [
                (x - w // 2, GROUND_Y - h + 40),
                (x + w // 2, GROUND_Y - h + 40),
                (x, GROUND_Y - h - 60)
            ]
            pygame.draw.polygon(screen, leaf_color, points)

        # 地面
        pygame.draw.rect(screen, DARK_GREEN, (0, GROUND_Y, WIDTH, HEIGHT - GROUND_Y))
        pygame.draw.line(screen, GREEN, (0, GROUND_Y), (WIDTH, GROUND_Y), 3)


# ---------------- 投射物 ----------------
class Projectile:
    def __init__(self, x, y, direction, element, owner):
        self.x = x
        self.y = y
        self.direction = direction
        self.element = element
        self.owner = owner

        self.speed = 9
        self.radius = 10
        self.damage = 18
        self.alive = True

        self.lifetime = 90
        self.trail = []

    @property
    def rect(self):
        return pygame.Rect(
            self.x - self.radius,
            self.y - self.radius,
            self.radius * 2,
            self.radius * 2
        )

    def update(self):
        if not self.alive:
            return

        self.trail.append((self.x, self.y))
        if len(self.trail) > 8:
            self.trail.pop(0)

        self.x += self.speed * self.direction
        self.lifetime -= 1

        if self.x < -20 or self.x > WIDTH + 20 or self.lifetime <= 0:
            self.alive = False

    def draw(self):
        if not self.alive:
            return

        if self.element == "fire":
            main_color = FIRE_COLOR
            light_color = FIRE_LIGHT
        else:
            main_color = ICE_COLOR
            light_color = ICE_LIGHT

        for i, (tx, ty) in enumerate(self.trail):
            alpha = int(120 * (i + 1) / len(self.trail))
            radius = max(2, self.radius // 2 + i // 2)
            pygame.draw.circle(screen, (*main_color, alpha), (int(tx), int(ty)), radius)

        pygame.draw.circle(screen, main_color, (int(self.x), int(self.y)), self.radius)
        pygame.draw.circle(screen, light_color, (int(self.x), int(self.y)), self.radius // 2)


# ---------------- 冰火人 ----------------
class IceFireMan:
    def __init__(self, x, y, element, name, color, light_color, controls):
        self.x = x
        self.y = y
        self.element = element
        self.name = name
        self.color = color
        self.light_color = light_color
        self.controls = controls

        self.width = 44
        self.height = 92

        self.vx = 0
        self.vy = 0
        self.speed = 5
        self.jump_power = 14

        self.on_ground = False
        self.facing = 1

        self.health = 100
        self.max_health = 100

        self.attack_cooldown = 0
        self.hit_cooldown = 0
        self.alive = True

        self.walk_phase = 0

    @property
    def rect(self):
        return pygame.Rect(
            self.x - self.width // 2,
            self.y - self.height,
            self.width,
            self.height
        )

    def handle_input(self, keys, projectiles):
        if not self.alive:
            return

        self.vx = 0

        if keys[self.controls["left"]]:
            self.vx = -self.speed
            self.facing = -1

        if keys[self.controls["right"]]:
            self.vx = self.speed
            self.facing = 1

        if keys[self.controls["jump"]] and self.on_ground:
            self.vy = -self.jump_power
            self.on_ground = False

        if keys[self.controls["attack"]] and self.attack_cooldown <= 0:
            projectile_x = self.x + self.facing * 28
            projectile_y = self.y - 55
            projectiles.append(Projectile(projectile_x, projectile_y, self.facing, self.element, self))
            self.attack_cooldown = 35

    def update(self):
        if not self.alive:
            return

        self.x += self.vx

        if self.x < 30:
            self.x = 30
        if self.x > WIDTH - 30:
            self.x = WIDTH - 30

        self.vy += GRAVITY
        self.y += self.vy

        if self.y >= GROUND_Y:
            self.y = GROUND_Y
            self.vy = 0
            self.on_ground = True

        if self.attack_cooldown > 0:
            self.attack_cooldown -= 1

        if self.hit_cooldown > 0:
            self.hit_cooldown -= 1

        if self.vx != 0 and self.on_ground:
            self.walk_phase += 0.25

    def take_damage(self, damage):
        if self.hit_cooldown > 0 or not self.alive:
            return False

        self.health -= damage
        self.hit_cooldown = 25

        if self.health <= 0:
            self.health = 0
            self.alive = False

        return True

    def draw(self):
        x, y = self.x, self.y

        if not self.alive:
            body_color = GRAY
            head_color = GRAY
        else:
            body_color = self.color
            head_color = self.light_color

            if self.hit_cooldown > 0 and self.hit_cooldown % 6 < 3:
                body_color = WHITE
                head_color = WHITE

        # 元素光环
        if self.alive:
            aura_color = (*self.color, 70)
            pygame.draw.circle(screen, aura_color, (x, y - 46), 52)

        # 头
        pygame.draw.circle(screen, head_color, (x, y - 76), 14)
        pygame.draw.circle(screen, body_color, (x, y - 76), 14, 3)

        # 身体
        pygame.draw.line(screen, body_color, (x, y - 62), (x, y - 26), 4)

        # 腿
        if self.vx != 0 and self.on_ground:
            leg_offset = int(10 * abs((self.walk_phase % 2) - 1))
            pygame.draw.line(screen, body_color, (x, y - 26), (x - 12 + leg_offset, y), 4)
            pygame.draw.line(screen, body_color, (x, y - 26), (x + 12 - leg_offset, y), 4)
        else:
            pygame.draw.line(screen, body_color, (x, y - 26), (x - 12, y), 4)
            pygame.draw.line(screen, body_color, (x, y - 26), (x + 12, y), 4)

        # 手臂
        if self.attack_cooldown > 25:
            arm_end_x = x + self.facing * 30
            arm_end_y = y - 55
            pygame.draw.line(screen, body_color, (x, y - 55), (arm_end_x, arm_end_y), 4)
        else:
            pygame.draw.line(screen, body_color, (x, y - 55), (x - 16, y - 38), 4)
            pygame.draw.line(screen, body_color, (x, y - 55), (x + 16, y - 38), 4)

        # 元素标志
        if self.element == "fire":
            flame_points = [
                (x, y - 98),
                (x - 6, y - 86),
                (x + 6, y - 86)
            ]
            pygame.draw.polygon(screen, FIRE_COLOR, flame_points)
            pygame.draw.polygon(screen, FIRE_LIGHT, [
                (x, y - 94),
                (x - 3, y - 88),
                (x + 3, y - 88)
            ])
        else:
            pygame.draw.line(screen, ICE_COLOR, (x, y - 98), (x, y - 86), 3)
            pygame.draw.line(screen, ICE_COLOR, (x - 6, y - 92), (x + 6, y - 92), 3)
            pygame.draw.line(screen, ICE_LIGHT, (x - 4, y - 96), (x + 4, y - 88), 2)
            pygame.draw.line(screen, ICE_LIGHT, (x + 4, y - 96), (x - 4, y - 88), 2)

        # 名字
        name_text = font.render(self.name, True, self.color)
        screen.blit(name_text, (x - name_text.get_width() // 2, y - 120))

        # 头顶血条
        bar_width = 80
        bar_height = 8
        bar_x = x - bar_width // 2
        bar_y = y - 108

        pygame.draw.rect(screen, GRAY, (bar_x, bar_y, bar_width, bar_height))
        health_width = int(bar_width * self.health / self.max_health)
        pygame.draw.rect(screen, self.color, (bar_x, bar_y, health_width, bar_height))
        pygame.draw.rect(screen, WHITE, (bar_x, bar_y, bar_width, bar_height), 1)


# ---------------- 游戏逻辑 ----------------
def create_players():
    fire_man = IceFireMan(
        x=260,
        y=GROUND_Y,
        element="fire",
        name="火人",
        color=FIRE_COLOR,
        light_color=FIRE_LIGHT,
        controls={
            "left": pygame.K_a,
            "right": pygame.K_d,
            "jump": pygame.K_w,
            "attack": pygame.K_f
        }
    )

    ice_man = IceFireMan(
        x=700,
        y=GROUND_Y,
        element="ice",
        name="冰人",
        color=ICE_COLOR,
        light_color=ICE_LIGHT,
        controls={
            "left": pygame.K_LEFT,
            "right": pygame.K_RIGHT,
            "jump": pygame.K_UP,
            "attack": pygame.K_SLASH
        }
    )

    return fire_man, ice_man


def get_damage(projectile, target):
    if projectile.element == target.element:
        return 8
    return 20


def draw_hud(fire_man, ice_man):
    bar_width = 320
    bar_height = 22

    pygame.draw.rect(screen, GRAY, (30, 20, bar_width, bar_height))
    pygame.draw.rect(
        screen,
        FIRE_COLOR,
        (30, 20, int(bar_width * fire_man.health / fire_man.max_health), bar_height)
    )
    pygame.draw.rect(screen, WHITE, (30, 20, bar_width, bar_height), 2)

    p1_text = font.render("火人：A/D移动 W跳跃 F攻击", True, FIRE_COLOR)
    screen.blit(p1_text, (30, 50))

    pygame.draw.rect(screen, GRAY, (WIDTH - 30 - bar_width, 20, bar_width, bar_height))
    pygame.draw.rect(
        screen,
        ICE_COLOR,
        (
            WIDTH - 30 - int(bar_width * ice_man.health / ice_man.max_health),
            20,
            int(bar_width * ice_man.health / ice_man.max_health),
            bar_height
        )
    )
    pygame.draw.rect(screen, WHITE, (WIDTH - 30 - bar_width, 20, bar_width, bar_height), 2)

    p2_text = font.render("冰人：←/→移动 ↑跳跃 /攻击", True, ICE_COLOR)
    screen.blit(p2_text, (WIDTH - 30 - p2_text.get_width(), 50))


def draw_winner(winner_name):
    overlay = pygame.Surface((WIDTH, HEIGHT), pygame.SRCALPHA)
    overlay.fill((0, 0, 0, 160))
    screen.blit(overlay, (0, 0))

    win_text = big_font.render(f"{winner_name} 胜利！", True, WHITE)
    screen.blit(win_text, (WIDTH // 2 - win_text.get_width() // 2, HEIGHT // 2 - 80))

    restart_text = font.render("按 R 重新开始", True, WHITE)
    screen.blit(restart_text, (WIDTH // 2 - restart_text.get_width() // 2, HEIGHT // 2 + 20))


# ---------------- 主循环 ----------------
def main():
    background = ForestBackground()
    fire_man, ice_man = create_players()
    projectiles = []
    winner = None

    running = True
    while running:
        clock.tick(FPS)

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False

            if event.type == pygame.KEYDOWN and event.key == pygame.K_r:
                fire_man, ice_man = create_players()
                projectiles.clear()
                winner = None

        keys = pygame.key.get_pressed()

        if winner is None:
            fire_man.handle_input(keys, projectiles)
            ice_man.handle_input(keys, projectiles)

            fire_man.update()
            ice_man.update()

            for projectile in projectiles:
                projectile.update()

            for projectile in projectiles[:]:
                if not projectile.alive:
                    projectiles.remove(projectile)
                    continue

                target = ice_man if projectile.owner == fire_man else fire_man

                if target.alive and projectile.rect.colliderect(target.rect):
                    damage = get_damage(projectile, target)
                    target.take_damage(damage)
                    projectile.alive = False
                    projectiles.remove(projectile)

            if not fire_man.alive:
                winner = "冰人"
            elif not ice_man.alive:
                winner = "火人"

        # 绘制
        background.draw()

        for projectile in projectiles:
            projectile.draw()

        fire_man.draw()
        ice_man.draw()

        draw_hud(fire_man, ice_man)

        if winner:
            draw_winner(winner)

        pygame.display.flip()

    pygame.quit()
    sys.exit()


if __name__ == "__main__":
    main()