import pygame
import random
import math

pygame.init()
WIDTH, HEIGHT = 960, 640
screen = pygame.display.set_mode((WIDTH, HEIGHT))
clock = pygame.time.Clock()
pygame.display.set_caption("小鳄鱼爱洗澡")

# ===== 颜色 =====
BG_COLOR = (25, 45, 70)
WATER_COLOR = (70, 160, 255)
WATER_LIGHT = (120, 190, 255)
DIRT_COLOR = (150, 110, 70)
STONE_COLOR = (80, 80, 90)
SPONGE_COLOR = (220, 170, 90)
BLOCK_OPEN = (60, 160, 80)
BLOCK_CLOSE = (200, 70, 70)
WHITE = (255, 255, 255)
GREEN = (30, 180, 60)

# 字体前置加载（防止报错）
try:
    font = pygame.font.SysFont("simhei", 20)
    font_big = pygame.font.SysFont("simhei", 32)
except:
    font = pygame.font.Font(None, 20)
    font_big = pygame.font.Font(None, 32)

# 网格【泥土缩小 CELL_SIZE = 8】
CELL_SIZE = 8
COLS = WIDTH // CELL_SIZE
ROWS = HEIGHT // CELL_SIZE

EMPTY = 0
DIRT = 1
STONE = 2
SPONGE = 3

# 关卡数据
current_level = 1
MAX_LEVEL = 3

# 水流粒子【提高流动速度版本】
class WaterParticle:
    def __init__(self, x, y):
        self.x = x
        self.y = y
        self.vx = random.uniform(-0.5, 0.5)
        self.vy = 0
        self.radius = 3.2
        self.life = 900
        self.on_floor = False

    def update(self, all_water):
        # 重力略微加大
        self.vy += 0.12
        self.x += self.vx
        self.y += self.vy
        self.life -= 1

        self.on_floor = False
        cx = int(self.x // CELL_SIZE)
        cy = int(self.y // CELL_SIZE)

        # 检测下方是否有固体
        if 0 <= cx < COLS and 0 <= cy < ROWS:
            if grid[cy][cx] != EMPTY:
                self.y = cy * CELL_SIZE - self.radius
                self.vy = 0
                self.on_floor = True

        # 地面上的水：向周围低洼扩散
        if self.on_floor:
            # 邻近水流互相挤压摊开
            for other in all_water:
                if other == self:
                    continue
                dx = self.x - other.x
                dy = self.y - other.y
                dist = math.hypot(dx, dy)
                if dist < self.radius * 2.2 and dist > 0:
                    force = 0.035  # 增大粒子相互推力
                    self.vx += dx / dist * force
                    self.vy += dy / dist * force

            # 摩擦力降低，滑动更快
            self.vx *= 0.96
            # 增大随机侧向流动幅度
            if abs(self.vx) < 0.32:
                self.vx += random.uniform(-0.32, 0.32)

        # 提高速度上限
        self.vx = max(-3.0, min(3.0, self.vx))

    def draw(self):
        pygame.draw.circle(screen, WATER_LIGHT, (int(self.x), int(self.y)), self.radius)
        pygame.draw.circle(screen, WATER_COLOR, (int(self.x), int(self.y)), self.radius - 1)

# 挡板机关类
class Gate:
    def __init__(self, grid_x, grid_y):
        self.gx = grid_x
        self.gy = grid_y
        self.open = False
        self.rect = pygame.Rect(grid_x*CELL_SIZE, grid_y*CELL_SIZE, CELL_SIZE, CELL_SIZE*6)

    def toggle(self):
        self.open = not self.open

    def draw(self):
        if self.open:
            pygame.draw.rect(screen, BLOCK_OPEN, self.rect)
        else:
            pygame.draw.rect(screen, BLOCK_CLOSE, self.rect)

    def is_block(self):
        return not self.open

# 全局对象
grid = []
gates = []
water_particles = []
source_x = source_y = 0
croc_rect = None
water_in_bath = 0
WIN_THRESHOLD = 8
game_state = "playing"
spawn_timer = 0
MAX_WATER = 520

# 加载关卡
def load_level(lv):
    global grid, gates, source_x, source_y, croc_rect, water_particles, water_in_bath, game_state, spawn_timer
    grid = [[EMPTY for _ in range(COLS)] for _ in range(ROWS)]
    gates.clear()
    water_particles.clear()
    water_in_bath = 0
    game_state = "playing"
    spawn_timer = 0

    # 围墙
    for x in range(COLS):
        grid[0][x] = STONE
        grid[ROWS - 1][x] = STONE
    for y in range(ROWS):
        grid[y][0] = STONE
        grid[y][COLS - 1] = STONE

    if lv == 1:
        # 第1关 新手关
        source_x, source_y = 9 * CELL_SIZE, 3 * CELL_SIZE
        croc_rect = pygame.Rect(WIDTH - 160, HEIGHT - 140, 100, 80)
        for y in range(4, ROWS - 4):
            for x in range(6, COLS - 6):
                if not (8 <= x <= 12 and 3 <= y <= 6):
                    grid[y][x] = DIRT
        grid[14][24] = STONE
        grid[15][24] = STONE
        grid[16][24] = STONE
        grid[14][25] = STONE
        grid[22][30] = SPONGE
        grid[23][30] = SPONGE

    elif lv == 2:
        # 第2关 加入挡板
        source_x, source_y = 12 * CELL_SIZE, 3 * CELL_SIZE
        croc_rect = pygame.Rect(WIDTH - 180, HEIGHT - 130, 90, 70)
        for y in range(4, ROWS - 5):
            for x in range(5, COLS - 5):
                grid[y][x] = DIRT
        # 留出水源口
        for x in range(11,15):
            grid[3][x] = EMPTY
        gates.append(Gate(28, 12))
        grid[20][16] = SPONGE
        grid[21][16] = SPONGE

    elif lv == 3:
        # 第3关 海绵+多个挡板
        source_x, source_y = 8 * CELL_SIZE, 3 * CELL_SIZE
        croc_rect = pygame.Rect(WIDTH - 140, HEIGHT - 150, 90, 70)
        for y in range(4, ROWS - 4):
            for x in range(6, COLS - 6):
                grid[y][x] = DIRT
        for x in range(7,11):
            grid[3][x] = EMPTY
        gates.append(Gate(22,10))
        gates.append(Gate(32,18))
        grid[16][20] = SPONGE
        grid[17][20] = SPONGE
        grid[26][28] = SPONGE

load_level(current_level)

running = True
while running:
    clock.tick(60)
    screen.fill(BG_COLOR)
    mx, my = pygame.mouse.get_pos()
    mouse_down = pygame.mouse.get_pressed()[0]

    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        if event.type == pygame.MOUSEBUTTONDOWN:
            # 点击挡板切换开关
            for g in gates:
                if g.rect.collidepoint(mx, my):
                    g.toggle()
        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_SPACE:
                load_level(current_level)

    # 挖土（小网格细腻挖掘）
    if mouse_down and game_state == "playing":
        m_cx = mx // CELL_SIZE
        m_cy = my // CELL_SIZE
        brush_r = 2
        for dy in range(-brush_r, brush_r + 1):
            for dx in range(-brush_r, brush_r + 1):
                nx = m_cx + dx
                ny = m_cy + dy
                if 0 <= nx < COLS and 0 <= ny < ROWS:
                    if grid[ny][nx] == DIRT:
                        grid[ny][nx] = EMPTY

    # 绘制网格
    for y in range(ROWS):
        for x in range(COLS):
            val = grid[y][x]
            rect = pygame.Rect(x * CELL_SIZE, y * CELL_SIZE, CELL_SIZE - 1, CELL_SIZE - 1)
            if val == DIRT:
                pygame.draw.rect(screen, DIRT_COLOR, rect)
            elif val == STONE:
                pygame.draw.rect(screen, STONE_COLOR, rect)
            elif val == SPONGE:
                pygame.draw.rect(screen, SPONGE_COLOR, rect)

    # 绘制挡板
    for g in gates:
        g.draw()

    # 鳄鱼浴缸
    pygame.draw.rect(screen, (10, 100, 140), croc_rect)
    screen.blit(font.render("🐊浴室", True, WHITE), (croc_rect.x + 5, croc_rect.y + 20))

    if game_state == "playing":
        # 生成水流
        spawn_timer += 1
        if spawn_timer >= 8 and len(water_particles) < MAX_WATER:
            water_particles.append(WaterParticle(source_x, source_y))
            spawn_timer = 0

        remove_list = []
        for p in water_particles:
            p.update(water_particles)
            p.draw()

            # 碰到海绵直接吸收
            cx = int(p.x // CELL_SIZE)
            cy = int(p.y // CELL_SIZE)
            if 0 <= cx < COLS and 0 <= cy < ROWS:
                if grid[cy][cx] == SPONGE:
                    remove_list.append(p)
                    continue

            # 碰到关闭的挡板
            for g in gates:
                if g.is_block() and g.rect.collidepoint(p.x, p.y):
                    p.vy = -p.vy * 0.2
                    p.vx *= -0.5

            # 流入浴缸
            if croc_rect.collidepoint(p.x, p.y):
                water_in_bath += 1
                remove_list.append(p)
                if water_in_bath >= WIN_THRESHOLD:
                    game_state = "win"

            if p.life <= 0 or p.y > HEIGHT + 30:
                remove_list.append(p)

        for p in remove_list:
            if p in water_particles:
                water_particles.remove(p)

        # 水源标记
        pygame.draw.circle(screen, WHITE, (source_x, source_y), 7)
        screen.blit(font.render("水源", True, WHITE), (source_x + 12, source_y - 10))

        # UI提示
        screen.blit(font.render(f"第{current_level}关 | 水量 {water_in_bath}/{WIN_THRESHOLD}", True, WHITE), (10, 8))
        screen.blit(font.render("鼠标拖动挖土 | 点击红绿挡板切换开关", True, WHITE), (10, 32))
        screen.blit(font.render("⚠黄色海绵会吸水！不要让水流碰到！", True, (255,220,100)), (10, 56))

    if game_state == "win":
        if current_level < MAX_LEVEL:
            text_win = font_big.render("🎉过关！即将进入下一关", True, GREEN)
            tip = font.render("按空格键进入下一关", True, WHITE)
            current_level += 1
            load_level(current_level)
        else:
            text_win = font_big.render("🏆全部关卡通关！", True, GREEN)
            tip = font.render("按空格键重新从第1关开始", True, WHITE)
            current_level = 1
        screen.blit(text_win, (WIDTH // 2 - 240, HEIGHT // 2 - 40))
        screen.blit(tip, (WIDTH // 2 - 160, HEIGHT // 2 + 10))

    pygame.display.flip()

pygame.quit()