import pygame
import sys
import math
import os

# 初始化
pygame.init()
WINDOW_WIDTH, WINDOW_HEIGHT = 1200, 800
screen = pygame.display.set_mode((WINDOW_WIDTH, WINDOW_HEIGHT))
pygame.display.set_caption("图片网格 - 无限滚动")
clock = pygame.time.Clock()

# 颜色
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
GRAY = (200, 200, 200)
DARK_GRAY = (100, 100, 100)
BLUE = (70, 130, 230)
RED = (255, 100, 100)
LIGHT_BLUE = (173, 216, 230)

class ScrollBar:
    """垂直滚动条"""
    def __init__(self, x, y, height, content_height, view_height):
        self.rect = pygame.Rect(x, y, 16, height)
        self.content_height = content_height
        self.view_height = view_height
        self.dragging = False
        self.update_slider()
    
    def update_slider(self):
        """更新滑块大小和位置"""
        if self.content_height <= self.view_height:
            self.slider_height = self.rect.height
            self.slider_y = self.rect.y
            self.max_scroll = 0
            self.max_slider_move = 0
        else:
            ratio = self.view_height / self.content_height
            self.slider_height = max(30, int(self.rect.height * ratio))
            self.max_slider_move = self.rect.height - self.slider_height
            self.max_scroll = self.content_height - self.view_height
            if hasattr(self, 'scroll_ratio'):
                self.slider_y = self.rect.y + (self.scroll_ratio * self.max_slider_move)
            else:
                self.slider_y = self.rect.y
    
    def set_content_size(self, content_height, view_height):
        self.content_height = content_height
        self.view_height = view_height
        self.update_slider()
    
    def handle_event(self, event):
        if event.type == pygame.MOUSEBUTTONDOWN and event.button == 1:
            slider_rect = pygame.Rect(self.rect.x, self.slider_y, self.rect.width, self.slider_height)
            if slider_rect.collidepoint(event.pos):
                self.dragging = True
                self.drag_offset = event.pos[1] - self.slider_y
            elif self.rect.collidepoint(event.pos):
                self.slider_y = event.pos[1] - self.slider_height / 2
                self._clamp_slider()
                self.dragging = True
                self.drag_offset = self.slider_height / 2
                
        elif event.type == pygame.MOUSEBUTTONUP and event.button == 1:
            self.dragging = False
            
        elif event.type == pygame.MOUSEMOTION and self.dragging:
            self.slider_y = event.pos[1] - self.drag_offset
            self._clamp_slider()
    
    def _clamp_slider(self):
        self.slider_y = max(self.rect.y, min(self.slider_y, self.rect.y + self.max_slider_move))
    
    def get_scroll_offset(self):
        if self.content_height <= self.view_height:
            return 0
        
        slider_move = self.slider_y - self.rect.y
        self.scroll_ratio = slider_move / self.max_slider_move if self.max_slider_move > 0 else 0
        return int(self.scroll_ratio * self.max_scroll)
    
    def draw(self, screen):
        pygame.draw.rect(screen, (30, 30, 30), self.rect)
        pygame.draw.rect(screen, DARK_GRAY, self.rect, 1)
        slider_rect = pygame.Rect(self.rect.x, self.slider_y, self.rect.width, self.slider_height)
        color = BLUE if not self.dragging else RED
        pygame.draw.rect(screen, color, slider_rect)
        pygame.draw.rect(screen, DARK_GRAY, slider_rect, 1)

class InfiniteImageGrid:
    """无限滚动的图片网格（多列布局）"""
    def __init__(self, x, y, width, height, columns=2, thumbnail_size=120, spacing=15, name_height=20, bg_color=(30, 30, 30)):
        self.x = x
        self.y = y
        self.width = width
        self.height = height
        self.columns = columns  # 列数（这里默认2行，实际是2列）
        self.thumbnail_size = thumbnail_size
        self.spacing = spacing
        self.bg_color = bg_color
        
        self.images = []  # 存储图片Surface
        self.image_paths = []  # 存储路径
        self.image_names = []  # 存储名称
        self.original_images = []  # 存储原始图片
        self.original_images_size = []  # 存储原始图片大小
        
        self.scroll_offset = 0
        self.total_content_height = 0
        
        # 滚动条在右侧
        scrollbar_x = x + width + 8
        scrollbar_y = y
        self.scrollbar = ScrollBar(scrollbar_x, scrollbar_y, height, 0, height)


        self.name_height = name_height  # 名字显示区域高度
    
        # 计算每行高度（图片高度 + 间距 + 名字高度）
        self.row_height = thumbnail_size + spacing + name_height
        
        # 裁剪表面
        self.clip_surface = pygame.Surface((width, height))


        # 拖动相关
        self.dragging = False  # 是否有图片被拖动属性
        self.drag_index = -1
        self.drag_start_pos = (0, 0)
        self.drag_start_image_pos = (0, 0)
        self.drag_current_offset = (0, 0)
        self.drag_offset_x = 0
        self.drag_offset_y = 0

    def add_image_name(self, image_path):
        """设置指定图片的名字"""
        name = os.path.basename(image_path)
        name = os.path.splitext(name)[0]
        self.image_names.append(name)

    
    def add_image(self, image_path):
        """添加单张图片"""
        try:
            original = pygame.image.load(image_path)  # 原始图片
            original_size = original.get_size()
            if original_size[0] > 500: original_size = (500, original_size[1])  # 限制最大宽度
            if original_size[1] > 500: original_size = (original_size[0], 500)  # 限制最大高度
            
            # 缩放到显示大小
            scaled = pygame.transform.scale(original, (self.thumbnail_size, self.thumbnail_size))
            
            self.original_images.append(original)  # 保存原始图片
            self.images.append(scaled)             # 保存缩放后的图片
            self.image_paths.append(image_path)
            self.original_images_size.append(original_size)
            
            # ===== 关键：同步添加图片名称 =====
            name = os.path.basename(image_path)
            name = os.path.splitext(name)[0]
            self.image_names.append(name)
            
            self._update_content_height()
            return True
        except Exception as e:
            print(f"加载失败 {image_path}: {e}")
            return False
    
    def add_images(self, image_paths):
        """批量添加图片"""
        for path in image_paths:
            self.add_image(path)

    def set_image_to_original_size(self, index):
        """缩放指定索引的图片为原来大小"""
        if 0 <= index < len(self.original_images):
            original = self.original_images[index]
            scaled = pygame.transform.scale(original, self.original_images_size[index])
            self.images[index] = scaled
            return True
        return False
    
    def add_test_images(self, count):
        """添加测试图片（演示用）"""
        colors = [(255,100,100), (100,255,100), (100,100,255), (255,255,100),
                  (255,100,255), (100,255,255), (255,200,100), (200,100,255),
                  (100,200,255), (255,150,150)]
        
        for i in range(count):
            surf = pygame.Surface((self.thumbnail_size, self.thumbnail_size))
            surf.fill(colors[i % len(colors)])
            
            font = pygame.font.Font(None, 36)
            text = font.render(str(i+1), True, WHITE)
            text_rect = text.get_rect(center=(self.thumbnail_size//2, self.thumbnail_size//2))
            surf.blit(text, text_rect)
            pygame.draw.rect(surf, BLACK, (0, 0, self.thumbnail_size, self.thumbnail_size), 2)
            
            self.images.append(surf)
            self.image_paths.append("test")
        
        self._update_content_height()
    
    def _update_content_height(self):
        """根据图片数量计算总高度"""
        if not self.images:
            self.total_content_height = 0
        else:
            # 计算需要多少行
            rows = math.ceil(len(self.images) / self.columns)
            self.total_content_height = rows * self.row_height + self.spacing
        
        self.scrollbar.set_content_size(self.total_content_height, self.height)
    
    def handle_event(self, event):
        """处理事件（滚轮、滚动条、鼠标拖动）"""
        # 鼠标滚轮
        if event.type == pygame.MOUSEWHEEL:
            mouse_x, mouse_y = pygame.mouse.get_pos()
            if self.x <= mouse_x <= self.x + self.width and self.y <= mouse_y <= self.y + self.height:
                self.scroll_offset -= event.y * 25
                self._clamp_scroll()
                self._sync_scrollbar()
        
        # 鼠标按下 - 开始拖动
        if event.type == pygame.MOUSEBUTTONDOWN and event.button == 1:
            mouse_x, mouse_y = event.pos
            self.drag_index = self.get_clicked_image_index(mouse_x, mouse_y)
            if self.drag_index != -1:
                self.dragging = True
                
                # 获取被点击图片的实际屏幕位置
                rect = self.get_image_rect(self.drag_index)
                
                # 计算鼠标相对于图片左上角的偏移
                self.drag_offset_x = mouse_x - rect.x
                self.drag_offset_y = mouse_y - rect.y
                
                self.drag_start_pos = (mouse_x, mouse_y)
                self.drag_start_image_pos = (rect.x, rect.y)
        
        # 鼠标移动 - 拖动中
        if event.type == pygame.MOUSEMOTION and self.dragging:
            mouse_x, mouse_y = event.pos
            # 计算图片应该出现的位置（鼠标位置减去偏移量）
            new_x = mouse_x - self.drag_offset_x
            new_y = mouse_y - self.drag_offset_y
            
            # 计算相对于原始位置的偏移
            dx = new_x - self.drag_start_image_pos[0]
            dy = new_y - self.drag_start_image_pos[1]
            
            self.drag_current_offset = (dx, dy)
        
        # 鼠标释放 - 结束拖动
        if event.type == pygame.MOUSEBUTTONUP and event.button == 1 and self.dragging:
            self.dragging = False
            self.drag_index = -1
            self.drag_current_offset = (0, 0)
        
        # 滚动条事件
        self.scrollbar.handle_event(event)
        self.scroll_offset = self.scrollbar.get_scroll_offset()
    
    def _clamp_scroll(self):
        """限制滚动范围"""
        max_scroll = max(0, self.total_content_height - self.height)
        self.scroll_offset = max(0, min(self.scroll_offset, max_scroll))
    
    def _sync_scrollbar(self):
        """同步滚动条位置"""
        if self.total_content_height <= self.height:
            return
        
        max_scroll = self.total_content_height - self.height
        scroll_ratio = self.scroll_offset / max_scroll if max_scroll > 0 else 0
        self.scrollbar.slider_y = self.scrollbar.rect.y + (scroll_ratio * self.scrollbar.max_slider_move)
        self.scrollbar.scroll_ratio = scroll_ratio
    
    def draw(self, screen, font_path):
        """绘制网格"""
        self.clip_surface.fill(self.bg_color)
        
        if not self.images:
            font = pygame.font.Font(font_path, 36)
            text = font.render("暂无图片", True, GRAY)
            text_rect = text.get_rect(center=(self.width//2, self.height//2))
            self.clip_surface.blit(text, text_rect)
        else:
            # 计算可见区域
            start_y = self.scroll_offset
            end_y = self.scroll_offset + self.height
            
            start_row = max(0, start_y // self.row_height)
            end_row = min(math.ceil(len(self.images) / self.columns), (end_y // self.row_height) + 2)
            
            # 绘制可见的图片
            for row in range(start_row, end_row):
                for col in range(self.columns):
                    idx = row * self.columns + col
                    if idx >= len(self.images):
                        break
                    
                    # 如果是正在拖动的图片，跳过（单独绘制）
                    if self.dragging and idx == self.drag_index:
                        continue
                    
                    # 计算位置
                    img_x = self.spacing + col * (self.thumbnail_size + self.spacing)
                    img_y = self.spacing + row * self.row_height - self.scroll_offset
                    
                    # 只绘制可见的
                    if img_y + self.thumbnail_size > 0 and img_y < self.height:
                        # 背景框
                        bg_rect = pygame.Rect(img_x - 3, img_y - 3, 
                                            self.thumbnail_size + 6, self.thumbnail_size + 6)
                        pygame.draw.rect(self.clip_surface, LIGHT_BLUE, bg_rect)
                        pygame.draw.rect(self.clip_surface, DARK_GRAY, bg_rect, 1)
                        
                        # 图片
                        self.clip_surface.blit(self.images[idx], (img_x, img_y))
                        
                        # ===== 显示图片名字 =====
                        name = self.image_names[idx] if idx < len(self.image_names) else f"图片{idx+1}"
                        font = pygame.font.Font(font_path, 14)
                        text = font.render(name, True, BLACK)
                        
                        # 文字位置（图片下方居中）
                        text_x = img_x + (self.thumbnail_size - text.get_width()) // 2
                        text_y = img_y + self.thumbnail_size + 2
                        
                        # 文字背景
                        bg_rect = pygame.Rect(text_x - 2, text_y - 1, text.get_width() + 4, text.get_height() + 2)
                        pygame.draw.rect(self.clip_surface, WHITE, bg_rect)
                        pygame.draw.rect(self.clip_surface, DARK_GRAY, bg_rect, 1)
                        
                        # 绘制文字
                        self.clip_surface.blit(text, (text_x, text_y))
        
        # 边框
        pygame.draw.rect(self.clip_surface, BLACK, (0, 0, self.width, self.height), 2)
        
        # 绘制网格区域到屏幕
        screen.blit(self.clip_surface, (self.x, self.y))
        
        # 滚动条
        self.scrollbar.draw(screen)
        
        # 绘制被拖动的图片（如果正在拖动）
        if self.dragging and self.drag_index != -1:
            # ===== 最后绘制被拖动的图片（真正的最上层）=====
            idx = self.drag_index
            # 计算原始位置
            row = idx // self.columns
            col = idx % self.columns
            original_x = self.spacing + col * (self.thumbnail_size + self.spacing)
            original_y = self.spacing + row * self.row_height - self.scroll_offset
            
            # 加上偏移量，并转换到屏幕坐标
            img_x = self.x + original_x + self.drag_current_offset[0]
            img_y = self.y + original_y + self.drag_current_offset[1]
            
            # 绘制阴影
            shadow_rect = pygame.Rect(img_x + 5, img_y + 5, self.thumbnail_size, self.thumbnail_size)
            pygame.draw.rect(screen, (100, 100, 100), shadow_rect)
            
            # 绘制图片
            screen.blit(self.images[idx], (img_x, img_y))
            
            # 绘制高亮边框
            pygame.draw.rect(screen, RED, (img_x, img_y, self.thumbnail_size, self.thumbnail_size), 3)
            
            # ===== 拖动时也显示图片名字 =====
            name = self.image_names[idx] if idx < len(self.image_names) else f"图片{idx+1}"
            font = pygame.font.Font(font_path, 14)
            text = font.render(name, True, BLACK)
            
            # 文字位置（图片下方居中）
            text_x = img_x + (self.thumbnail_size - text.get_width()) // 2
            text_y = img_y + self.thumbnail_size + 2
            
            # 文字背景
            bg_rect = pygame.Rect(text_x - 2, text_y - 1, text.get_width() + 4, text.get_height() + 2)
            pygame.draw.rect(screen, WHITE, bg_rect)
            pygame.draw.rect(screen, DARK_GRAY, bg_rect, 1)
            
            # 绘制文字
            screen.blit(text, (text_x, text_y))


    def get_clicked_image_index(self, mouse_x, mouse_y):
        """根据鼠标坐标返回点击的图片索引，如果没有点击到返回-1"""
        # 先检查鼠标是否在网格可视区域内
        if not (self.x <= mouse_x <= self.x + self.width and self.y <= mouse_y <= self.y + self.height):
            return -1
        
        # 计算相对于网格区域的坐标（注意：这里不加 scroll_offset）
        local_x = mouse_x - self.x
        local_y = mouse_y - self.y  # 不加滚动偏移
        
        # 获取当前可视区域顶部对应的内容位置
        view_top = self.scroll_offset
        
        # 计算鼠标在内容中的实际位置
        content_y = view_top + local_y
        
        # 完整的行高
        full_row_height = self.thumbnail_size + self.spacing + self.name_height
        
        # 计算行号
        row = content_y // full_row_height
        col = local_x // (self.thumbnail_size + self.spacing)
        
        # 边界检查
        if col >= self.columns or row < 0 or row >= len(self.images) // self.columns + 1:
            return -1
        
        # 计算索引
        index = row * self.columns + col
        
        # 验证索引有效
        if 0 <= index < len(self.images):
            # 计算该图片在内容中的实际位置
            img_row = index // self.columns
            img_col = index % self.columns
            img_content_y = img_row * full_row_height + self.spacing
            img_content_x = img_col * (self.thumbnail_size + self.spacing) + self.spacing
            
            # 检查鼠标是否在图片区域内
            img_top = img_content_y - view_top  # 屏幕上的Y坐标
            img_left = img_content_x
            
            if (img_left <= local_x <= img_left + self.thumbnail_size and 
                img_top <= local_y <= img_top + self.thumbnail_size):
                return index
        
        return -1

    def get_image_rect(self, index):
        """获取图片在屏幕上的矩形区域"""
        row = index // self.columns
        col = index % self.columns
        
        img_x = self.x + self.spacing + col * (self.thumbnail_size + self.spacing)
        img_y = self.y + self.spacing + row * self.row_height - self.scroll_offset
        
        return pygame.Rect(img_x, img_y, self.thumbnail_size, self.thumbnail_size)
    
    def get_drag_image_rect(self, index):
        """获取图片在屏幕上的矩形区域（自动处理拖动时的偏移）"""
        row = index // self.columns
        col = index % self.columns
        
        img_x = self.x + self.spacing + col * (self.thumbnail_size + self.spacing)
        img_y = self.y + self.spacing + row * self.row_height - self.scroll_offset
        
        # 如果是正在被拖动的图片，加上偏移量
        if self.dragging and index == self.drag_index:
            img_x += self.drag_current_offset[0]
            img_y += self.drag_current_offset[1]
        
        return pygame.Rect(img_x, img_y, self.thumbnail_size, self.thumbnail_size)

    def get_image_path(self, index):
        """获取指定索引的图片路径"""
        if 0 <= index < len(self.image_paths):
            return self.image_paths[index]
        return None



    def set_image_name(self, index, new_name):
        """修改指定索引的图片名字"""
        if 0 <= index < len(self.image_names):
            self.image_names[index] = new_name
            return True
        return False

    def get_image_name(self, index):
        """获取指定索引的图片名字"""
        if 0 <= index < len(self.image_names):
            return self.image_names[index]
        return None


    def get_hover_image_index(self, mouse_x, mouse_y):
        """获取鼠标悬停的图片索引，如果没有返回-1"""
        # 先检查鼠标是否在网格区域内
        if not (self.x <= mouse_x <= self.x + self.width and self.y <= mouse_y <= self.y + self.height):
            return -1
        
        # 计算相对于网格区域的坐标（不加滚动偏移）
        local_x = mouse_x - self.x
        local_y = mouse_y - self.y
        
        # 获取当前可视区域顶部对应的内容位置
        view_top = self.scroll_offset
        
        # 计算鼠标在内容中的实际位置
        content_y = view_top + local_y
        
        # 完整的行高
        full_row_height = self.thumbnail_size + self.spacing + self.name_height
        
        # 计算行号和列号
        col = int(local_x // (self.thumbnail_size + self.spacing))
        row = int(content_y // full_row_height)
        
        # 边界检查
        if col >= self.columns or row < 0:
            return -1
        
        # 计算图片索引
        index = row * self.columns + col
        
        # 检查索引是否有效
        if 0 <= index < len(self.images):
            # 计算该图片在内容中的实际位置
            img_row = index // self.columns
            img_col = index % self.columns
            img_content_y = img_row * full_row_height + self.spacing
            img_content_x = img_col * (self.thumbnail_size + self.spacing) + self.spacing
            
            # 检查鼠标是否在图片区域内
            img_top = img_content_y - view_top
            img_left = img_content_x
            
            if (img_left <= local_x <= img_left + self.thumbnail_size and 
                img_top <= local_y <= img_top + self.thumbnail_size):
                return index
        
        return -1

    def quit():
        """退出"""
        pygame.quit()
        sys.exit()



if __name__ == "__main__":
    # 设置：显示区域在右下角
    # 宽度300，高度窗口的一半多一点（500）
    grid_width = 320  # 宽度（包含间距）
    grid_height = 500  # 高度
    
    # 位置：右下角，留出边距
    margin = 20
    x = WINDOW_WIDTH - grid_width - margin
    y = WINDOW_HEIGHT - grid_height - margin
    
    # 创建网格（2列显示）
    grid = InfiniteImageGrid(
        x=x, 
        y=y, 
        width=grid_width, 
        height=grid_height,
        columns=2,          # 2列
        thumbnail_size=130, # 缩略图大小
        spacing=12,         # 间距
        bg_color=(240, 240, 240)
    )

    font_path = os.path.join(os.path.dirname(__file__), "font/MSYHBD.TTC")
    
    # 添加图片的几种方式：
    
    # 方式1：自动生成测试图片（演示无限滚动）
    print("生成50张测试图片演示无限滚动...")
    grid.add_test_images(50)  # 可以改成100、200等任意数量
    
    # 方式2：添加真实图片（取消注释）
    # import glob
    # image_files = glob.glob("*.jpg") + glob.glob("*.png")
    # grid.add_images(image_files)
    
    # 方式3：手动添加单张
    # grid.add_image("photo1.jpg")
    
    running = True
    while running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
            
            # 传递给网格处理
            grid.handle_event(event)
        
        # 绘制
        screen.fill(WHITE)
        
        # 绘制标题和说明
        font_title = pygame.font.Font(font_path, 36)
        title = font_title.render("图片网格 - 2列显示 | 无限滚动", True, BLACK)
        screen.blit(title, (margin, margin))
        
        font_info = pygame.font.Font(font_path, 24)
        info1 = font_info.render(f"共 {len(grid.images)} 张图片", True, DARK_GRAY)
        info2 = font_info.render("操作：拖拽右侧滚动条 或 鼠标滚轮", True, DARK_GRAY)
        info3 = font_info.render(f"当前滚动: {grid.scroll_offset} / {max(0, grid.total_content_height - grid.height)}", True, DARK_GRAY)
        
        screen.blit(info1, (margin, margin + 50))
        screen.blit(info2, (margin, margin + 80))
        screen.blit(info3, (margin, margin + 110))
        
        # 绘制网格区域说明
        pygame.draw.rect(screen, (200, 200, 200), (x-5, y-5, grid_width+10, grid_height+10), 1)
        label = font_info.render("图片显示区域", True, GRAY)
        screen.blit(label, (x+10, y-25))
        
        # 绘制网格
        grid.draw(screen, font_path)
        
        pygame.display.flip()
        clock.tick(60)
    
    pygame.quit()
    sys.exit()
