import pygame
import sys
import math

pygame.init()
WIDTH, HEIGHT = 1024, 768
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("森林冰火人 - 硬核协作版")
clock = pygame.time.Clock()

# --- 颜色与常量 ---
RED = (220, 50, 50)
BLUE = (50, 100, 220)
GREEN = (50, 200, 50)
GRAY = (80, 80, 80)
GOLD = (255, 215, 0)
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)

GRAVITY = 0.8
JUMP_FORCE = -18
MOVE_SPEED = 5
MAX_FALL = 12

# --- 核心机制类 ---

class Player(pygame.sprite.Sprite):
    """增强版玩家：支持精密跳跃与机关交互"""
    def __init__(self, x, y, color, controls, player_type):
        super().__init__()
        self.image = pygame.Surface((32, 48))
        self.image.fill(color)
        self.rect = self.image.get_rect(topleft=(x, y))
        self.color = color
        self.player_type = player_type
        self.controls = controls
        self.vel_x = 0
        self.vel_y = 0
        self.on_ground = False
        self.facing_right = True
        self.dead = False

    def update(self, platforms, hazards, switches, doors):
        if self.dead:
            return "dead"

        keys = pygame.key.get_pressed()

        # 精密水平移动 (带惯性)
        if keys[self.controls[0]]:
            self.vel_x = -MOVE_SPEED
            self.facing_right = False
        elif keys[self.controls[1]]:
            self.vel_x = MOVE_SPEED
            self.facing_right = True
        else:
            self.vel_x *= 0.8

        # 水平碰撞
        self.rect.x += self.vel_x
        for p in platforms:
            if self.rect.colliderect(p):
                if self.vel_x > 0:
                    self.rect.right = p.left
                elif self.vel_x < 0:
                    self.rect.left = p.right
                self.vel_x = 0

        # 精密垂直物理
        self.vel_y += GRAVITY
        self.vel_y = min(self.vel_y, MAX_FALL)
        self.rect.y += self.vel_y
        self.on_ground = False

        for p in platforms:
            if self.rect.colliderect(p):
                if self.vel_y > 0:
                    self.rect.bottom = p.top
                    self.vel_y = 0
                    self.on_ground = True
                elif self.vel_y < 0:
                    self.rect.top = p.bottom
                    self.vel_y = 0

        # 跳跃 (仅地面可跳)
        if keys[self.controls[2]] and self.on_ground:
            self.vel_y = JUMP_FORCE
            self.on_ground = False

        # 机关交互检测
        for sw in switches:
            if not sw.activated and self.rect.colliderect(sw.rect):
                sw.activated = True
                for door in doors:
                    if door.switch_id == sw.id:
                        door.open()

        # 危险检测
        for h in hazards:
            if self.rect.colliderect(h):
                if h.type == "lava" and self.player_type == "fire":
                    continue
                elif h.type == "water" and self.player_type == "ice":
                    continue
                elif h.type == "green":
                    continue
                else:
                    self.dead = True
                    return "dead"

        # 边界
        if self.rect.left < 0: self.rect.left = 0
        if self.rect.right > WIDTH: self.rect.right = WIDTH
        if self.rect.bottom > HEIGHT:
            self.dead = True
            return "dead"

        return "alive"


class Hazard:
    """危险区域：岩浆/水/酸液"""
    def __init__(self, x, y, w, h, type):
        self.rect = pygame.Rect(x, y, w, h)
        self.type = type

    def draw(self, surface):
        if self.type == "lava":
            color = RED
        elif self.type == "water":
            color = BLUE
        else:
            color = GREEN
        pygame.draw.rect(surface, color, self.rect)
        pygame.draw.rect(surface, WHITE, self.rect, 2)


class Switch:
    """机关开关：必须踩住才能开门"""
    def __init__(self, x, y, switch_id):
        self.rect = pygame.Rect(x, y, 40, 10)
        self.id = switch_id
        self.activated = False

    def draw(self, surface):
        color = GREEN if self.activated else GRAY
        pygame.draw.rect(surface, color, self.rect)
        pygame.draw.rect(surface, WHITE, self.rect, 2)


class Door:
    """机关门：对应开关激活时开启"""
    def __init__(self, x, y, w, h, switch_id):
        self.rect = pygame.Rect(x, y, w, h)
        self.switch_id = switch_id
        self.opened = False
        self.target_y = y - h

    def open(self):
        self.opened = True

    def update(self):
        # 门平滑升起
        if self.opened and self.rect.y > self.target_y:
            self.rect.y -= 3
        # 门平滑关闭 (已修复：使用 self.rect.height 替代未定义的 h)
        elif not self.opened and self.rect.y < self.target_y + self.rect.height:
            self.rect.y += 3

    def draw(self, surface):
        color = (50, 50, 50) if self.opened else (150, 50, 50)
        pygame.draw.rect(surface, color, self.rect)
        pygame.draw.rect(surface, WHITE, self.rect, 2)


class Diamond:
    """收集物：必须两人各自收集对应颜色"""
    def __init__(self, x, y, color_type):
        self.rect = pygame.Rect(x, y, 20, 20)
        self.color_type = color_type
        self.collected = False
        self.float_offset = 0

    def update(self):
        self.float_offset += 0.1
        self.rect.y += 2 * math.sin(self.float_offset)

    def draw(self, surface):
        if self.collected:
            return
        if self.color_type == "red":
            color = RED
        elif self.color_type == "blue":
            color = BLUE
        else:
            color = GREEN
        pygame.draw.polygon(surface, color, [
            (self.rect.centerx, self.rect.top),
            (self.rect.right, self.rect.centery),
            (self.rect.centerx, self.rect.bottom),
            (self.rect.left, self.rect.centery)
        ])


# --- 高难度关卡设计 (强制协作) ---
def create_hard_level():
    platforms = [
        pygame.Rect(0, 700, 1024, 64),
        pygame.Rect(100, 550, 200, 20),
        pygame.Rect(724, 550, 200, 20),
        pygame.Rect(350, 450, 324, 20),
        pygame.Rect(200, 300, 150, 20),
        pygame.Rect(674, 300, 150, 20),
        pygame.Rect(400, 150, 224, 20),
    ]

    hazards = [
        Hazard(300, 680, 424, 20, "lava"),
        Hazard(50, 680, 150, 20, "water"),
        Hazard(824, 680, 150, 20, "lava"),
        Hazard(400, 130, 224, 20, "green"),
    ]

    # 交叉开关：左开关控制右门，右开关控制左门
    switches = [
        Switch(150, 540, 1),
        Switch(774, 540, 2),
    ]
    doors = [
        Door(450, 350, 124, 100, 1),
        Door(450, 350, 124, 100, 2),
    ]

    diamonds = [
        Diamond(180, 270, "red"),
        Diamond(700, 270, "blue"),
        Diamond(512, 120, "green"),
    ]

    return platforms, hazards, switches, doors, diamonds


# --- 主循环 ---
def main():
    fireboy = Player(10, 600, RED, (pygame.K_a, pygame.K_d, pygame.K_w), "fire")
    icegirl = Player(10, 600, BLUE, (pygame.K_LEFT, pygame.K_RIGHT, pygame.K_UP), "ice")
    players = [fireboy, icegirl]

    platforms, hazards, switches, doors, diamonds = create_hard_level()

    font = pygame.font.SysFont("Arial", 20)
    small_font = pygame.font.SysFont("Arial", 16)

    running = True
    while running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
            if event.type == pygame.KEYDOWN and event.key == pygame.K_r:
                fireboy.rect.topleft = (50, 600)
                fireboy.vel_x = fireboy.vel_y = 0
                fireboy.dead = False
                icegirl.rect.topleft = (100, 600)
                icegirl.vel_x = icegirl.vel_y = 0
                icegirl.dead = False
                for sw in switches:
                    sw.activated = False
                for d in doors:
                    d.opened = False
                for dm in diamonds:
                    dm.collected = False

        screen.fill((20, 20, 30))

        for p in platforms:
            pygame.draw.rect(screen, GRAY, p)
        for h in hazards:
            h.draw(screen)
        for sw in switches:
            sw.draw(screen)
        for d in doors:
            d.update()
            d.draw(screen)
        for dm in diamonds:
            dm.update()
            dm.draw(screen)

        # 修复：无论生死都绘制角色，解决红色火娃不显示问题
        all_alive = True
        for p in players:
            status = p.update(platforms, hazards, switches, doors)
            if status == "dead":
                all_alive = False
            # 死亡状态下依然绘制角色，确保红色火娃可见
            screen.blit(p.image, p.rect)
            # 绘制朝向指示
            if p.facing_right:
                pygame.draw.circle(screen, WHITE, (p.rect.right - 4, p.rect.centery), 3)
            else:
                pygame.draw.circle(screen, WHITE, (p.rect.left + 4, p.rect.centery), 3)

        # 钻石收集检测
        for p in players:
            for dm in diamonds:
                if not dm.collected and p.rect.colliderect(dm.rect):
                    if dm.color_type == "green" or dm.color_type == p.player_type:
                        dm.collected = True

        # UI
        collected = sum(1 for d in diamonds if d.collected)
        total = len(diamonds)
        ui_text = font.render(f"Diamonds: {collected}/{total} | R: Reset", True, WHITE)
        screen.blit(ui_text, (10, 10))

        hint = small_font.render("FIRE: WASD | ICE: ARROWS | SWAP SWITCHES TO OPEN DOORS", True, (200, 200, 200))
        screen.blit(hint, (10, 40))

        if not all_alive:
            dead_text = font.render("DEAD! Press R to Reset", True, RED)
            screen.blit(dead_text, (WIDTH // 2 - 120, HEIGHT // 2))

        # 通关检测
        if all_alive and collected == total:
            win_text = font.render("LEVEL COMPLETE!", True, GOLD)
            screen.blit(win_text, (WIDTH // 2 - 100, HEIGHT // 2))

        pygame.display.flip()
        clock.tick(60)

    pygame.quit()
    sys.exit()

if __name__ == "__main__":
    main()