import pygame
import sys
import math
import random
from collections import deque

# ========== 初始化 ==========
pygame.init()
WIDTH, HEIGHT = 800, 600
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("🔧 Pipe Puzzle – Connect the Flow")
clock = pygame.time.Clock()

# ========== 配色方案 ==========
BG_COLOR = (28, 30, 40)
GRID_BG = (20, 22, 30)
PIPE_COLOR = (100, 180, 255)
PIPE_CONNECTED = (255, 220, 80)
START_COLOR = (80, 230, 100)
END_COLOR = (255, 100, 100)
TEXT_COLOR = (220, 220, 240)
TITLE_COLOR = (255, 215, 0)
WHITE = (255, 255, 255)

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

# ========== 管道形状（上，右，下，左） ==========
PIPE_SHAPES = [
    (False, True, False, True),   # 水平
    (True, False, True, False),   # 垂直
    (True, True, False, False),   # 右上直角
    (False, True, True, False),   # 右下直角
    (False, False, True, True),   # 左下直角
    (True, False, False, True),   # 左上直角
    (True, True, True, True),     # 十字
    (True, True, False, True),    # T上右下
    (False, True, True, True),    # T右下左
    (True, False, True, True),    # T下左上
    (True, True, True, False),    # T上右左
]

def rotate_shape(shape, rot):
    return shape[(4 - rot) % 4:] + shape[:(4 - rot) % 4]

# ========== 预设关卡 ==========
PRESET_LEVELS = [
    {
        "rows": 3, "cols": 3,
        "start": (0, 0), "end": (2, 2),
        "pipes": [
            [(2, 0), (0, 1), (5, 0)],
            [(1, 0), (2, 1), (0, 0)],
            [(4, 0), (0, 1), (3, 0)],
        ]
    },
    {
        "rows": 4, "cols": 4,
        "start": (0, 0), "end": (3, 3),
        "pipes": [
            [(2, 0), (0, 0), (5, 0), (2, 1)],
            [(1, 0), (2, 3), (0, 1), (4, 1)],
            [(0, 0), (1, 1), (3, 0), (0, 1)],
            [(4, 0), (0, 0), (0, 1), (1, 0)],
        ]
    },
    {
        "rows": 5, "cols": 5,
        "start": (0, 2), "end": (4, 2),
        "pipes": [
            [(0,1), (2,0), (1,0), (5,0), (0,1)],
            [(1,0), (2,1), (0,0), (2,2), (1,0)],
            [(3,1), (1,0), (6,0), (1,0), (4,2)],
            [(1,0), (0,0), (2,3), (0,0), (1,0)],
            [(0,1), (3,2), (1,0), (4,0), (0,1)],
        ]
    }
]

# ========== 游戏类 ==========
class PipeGame:
    def __init__(self):
        self.particles = []
        self.connected = False
        self.level_index = 0
        self.load_level(0)

    def load_level(self, idx):
        level = PRESET_LEVELS[idx]
        self.rows = level["rows"]
        self.cols = level["cols"]
        self.start = level["start"]
        self.end = level["end"]
        self.grid = [[list(cell) for cell in row] for row in level["pipes"]]
        self.connected = False
        self.particles.clear()
        self.connected = self.check_connectivity()
        if self.connected:
            self.generate_flow_particles()

    # ---------- 工具函数 ----------
    def cell_size(self):
        # 根据窗口大小和行列数自动计算格子大小
        return min(80, (WIDTH - 100) // self.cols, (HEIGHT - 120) // self.rows)

    def grid_origin(self):
        size = self.cell_size()
        grid_w = self.cols * size
        grid_h = self.rows * size
        ox = (WIDTH - grid_w) // 2
        oy = (HEIGHT - grid_h) // 2 + 20
        return ox, oy

    def get_cell_connections(self, row, col):
        shape_idx, rotation = self.grid[row][col]
        base = PIPE_SHAPES[shape_idx]
        return rotate_shape(base, rotation)

    def is_connected_to(self, r1, c1, r2, c2):
        if not (0 <= r1 < self.rows and 0 <= c1 < self.cols): return False
        if not (0 <= r2 < self.rows and 0 <= c2 < self.cols): return False
        conn1 = self.get_cell_connections(r1, c1)
        conn2 = self.get_cell_connections(r2, c2)
        if r2 == r1 - 1: return conn1[0] and conn2[2]
        if r2 == r1 + 1: return conn1[2] and conn2[0]
        if c2 == c1 - 1: return conn1[3] and conn2[1]
        if c2 == c1 + 1: return conn1[1] and conn2[3]
        return False

    def check_connectivity(self):
        visited = set()
        q = deque()
        sr, sc = self.start
        q.append((sr, sc))
        visited.add((sr, sc))
        while q:
            r, c = q.popleft()
            if (r, c) == self.end:
                return True
            for dr, dc in [(-1,0), (1,0), (0,-1), (0,1)]:
                nr, nc = r+dr, c+dc
                if (nr, nc) not in visited and self.is_connected_to(r, c, nr, nc):
                    visited.add((nr, nc))
                    q.append((nr, nc))
        return False

    def handle_click(self, pos):
        ox, oy = self.grid_origin()
        size = self.cell_size()
        x, y = pos
        col = (x - ox) // size
        row = (y - oy) // size
        if 0 <= row < self.rows and 0 <= col < self.cols:
            self.grid[row][col][1] = (self.grid[row][col][1] + 1) % 4
            self.connected = self.check_connectivity()
            if self.connected:
                self.generate_flow_particles()

    def generate_flow_particles(self):
        self.particles.clear()
        # 查找路径
        parent = {}
        q = deque()
        sr, sc = self.start
        q.append((sr, sc))
        parent[(sr, sc)] = None
        while q:
            r, c = q.popleft()
            if (r, c) == self.end:
                break
            for dr, dc in [(-1,0),(1,0),(0,-1),(0,1)]:
                nr, nc = r+dr, c+dc
                if (0 <= nr < self.rows and 0 <= nc < self.cols and (nr, nc) not in parent):
                    if self.is_connected_to(r, c, nr, nc):
                        parent[(nr, nc)] = (r, c)
                        q.append((nr, nc))
        path = []
        cur = self.end
        while cur is not None:
            path.append(cur)
            cur = parent.get(cur)
        path.reverse()

        ox, oy = self.grid_origin()
        size = self.cell_size()
        for i, (r, c) in enumerate(path):
            px = ox + c * size + size // 2
            py = oy + r * size + size // 2
            for _ in range(4):
                p = {
                    "x": px + random.uniform(-6, 6),
                    "y": py + random.uniform(-6, 6),
                    "speed": random.uniform(0.5, 1.5),
                    "progress": 0.0,
                    "color": PIPE_CONNECTED,
                    "size": random.randint(3, 6),
                    "path": path,
                    "path_idx": i
                }
                self.particles.append(p)

    def update(self):
        if not self.connected:
            return
        ox, oy = self.grid_origin()
        size = self.cell_size()
        for p in self.particles[:]:
            p["progress"] += p["speed"] * 0.03
            if p["progress"] >= 1.0:
                p["path_idx"] += 1
                if p["path_idx"] >= len(p["path"]):
                    self.particles.remove(p)
                    continue
                p["progress"] = 0.0
            idx = p["path_idx"]
            if idx >= len(p["path"]):
                continue
            r, c = p["path"][idx]
            cx = ox + c * size + size // 2
            cy = oy + r * size + size // 2
            if idx < len(p["path"]) - 1:
                nr, nc = p["path"][idx + 1]
                nx = ox + nc * size + size // 2
                ny = oy + nr * size + size // 2
                p["x"] = cx + (nx - cx) * p["progress"]
                p["y"] = cy + (ny - cy) * p["progress"]
            else:
                p["x"] = cx
                p["y"] = cy

    def draw_grid(self):
        ox, oy = self.grid_origin()
        size = self.cell_size()

        # 背景格
        for r in range(self.rows):
            for c in range(self.cols):
                rect = pygame.Rect(ox + c*size, oy + r*size, size, size)
                pygame.draw.rect(screen, GRID_BG, rect, border_radius=8)
                pygame.draw.rect(screen, (60, 60, 70), rect, 2, border_radius=8)

        # 管道
        for r in range(self.rows):
            for c in range(self.cols):
                cx = ox + c * size + size // 2
                cy = oy + r * size + size // 2
                conn = self.get_cell_connections(r, c)
                pipe_color = PIPE_CONNECTED if self.connected else PIPE_COLOR
                pygame.draw.circle(screen, pipe_color, (cx, cy), size//6)
                arm_len = size//3
                arm_w = size//6
                dirs = [
                    (0, -arm_len, conn[0]),
                    (arm_len, 0, conn[1]),
                    (0, arm_len, conn[2]),
                    (-arm_len, 0, conn[3])
                ]
                for dx, dy, active in dirs:
                    if active:
                        end_x = cx + dx
                        end_y = cy + dy
                        pygame.draw.line(screen, pipe_color, (cx, cy), (end_x, end_y), arm_w)
                        pygame.draw.circle(screen, pipe_color, (end_x, end_y), arm_w//2)

        # 起点/终点
        sr, sc = self.start
        er, ec = self.end
        sx = ox + sc * size + size // 2
        sy = oy + sr * size + size // 2
        ex = ox + ec * size + size // 2
        ey = oy + er * size + size // 2
        pygame.draw.circle(screen, START_COLOR, (sx, sy), size//5)
        pygame.draw.circle(screen, WHITE, (sx, sy), size//5, 2)
        s_text = font_small.render("S", True, WHITE)
        screen.blit(s_text, (sx-6, sy-12))
        pygame.draw.circle(screen, END_COLOR, (ex, ey), size//5)
        pygame.draw.circle(screen, WHITE, (ex, ey), size//5, 2)
        e_text = font_small.render("E", True, WHITE)
        screen.blit(e_text, (ex-6, ey-12))

        # 粒子
        for p in self.particles:
            pygame.draw.circle(screen, p["color"], (int(p["x"]), int(p["y"])), p["size"])

    def draw_ui(self):
        title = font_large.render("Pipe Puzzle", True, TITLE_COLOR)
        screen.blit(title, (WIDTH // 2 - title.get_width() // 2, 15))
        if self.connected:
            msg = font_med.render("Connected! Press SPACE for next level", True, (100, 255, 100))
        else:
            msg = font_small.render("Click a pipe to rotate. Connect S to E.", True, TEXT_COLOR)
        screen.blit(msg, (WIDTH // 2 - msg.get_width() // 2, HEIGHT - 40))

    def run(self):
        running = True
        while running:
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    running = False
                elif event.type == pygame.MOUSEBUTTONDOWN and event.button == 1:
                    self.handle_click(event.pos)
                elif event.type == pygame.KEYDOWN:
                    if event.key == pygame.K_SPACE and self.connected:
                        self.level_index = (self.level_index + 1) % len(PRESET_LEVELS)
                        self.load_level(self.level_index)
                    elif event.key == pygame.K_r:
                        self.load_level(self.level_index)

            self.update()
            screen.fill(BG_COLOR)
            self.draw_grid()
            self.draw_ui()
            pygame.display.flip()
            clock.tick(60)
        pygame.quit()
        sys.exit()

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