import pygame
import sys
import math
import os
from datetime import datetime

# 初始化Pygame
pygame.init()

# 窗口设置
WINDOW_WIDTH = 1200
WINDOW_HEIGHT = 800
BACKGROUND_COLOR = (240, 240, 245)
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
RED = (255, 50, 50)
GREEN = (50, 255, 50)
BLUE = (50, 150, 255)
YELLOW = (255, 255, 50)
PURPLE = (200, 50, 200)
ORANGE = (255, 165, 0)
PINK = (255, 105, 180)
CYAN = (50, 255, 200)
GRAY = (128, 128, 128)
DARK_GRAY = (64, 64, 64)
LIGHT_GRAY = (200, 200, 200)

# 工具栏设置
TOOLBAR_WIDTH = 200
DRAWING_AREA = pygame.Rect(
    TOOLBAR_WIDTH, 0, WINDOW_WIDTH - TOOLBAR_WIDTH, WINDOW_HEIGHT)


class Button:
    def __init__(self, x, y, width, height, text, color, hover_color, font):
        self.rect = pygame.Rect(x, y, width, height)
        self.text = text
        self.color = color
        self.hover_color = hover_color
        self.current_color = color
        self.font = font

    def draw(self, screen):
        pygame.draw.rect(screen, self.current_color, self.rect)
        pygame.draw.rect(screen, DARK_GRAY, self.rect, 2)
        text_surface = self.font.render(self.text, True, BLACK)
        text_rect = text_surface.get_rect(center=self.rect.center)
        screen.blit(text_surface, text_rect)

    def update(self, mouse_pos):
        if self.rect.collidepoint(mouse_pos):
            self.current_color = self.hover_color
            return True
        else:
            self.current_color = self.color
            return False

    def is_clicked(self, mouse_pos, mouse_click):
        return self.rect.collidepoint(mouse_pos) and mouse_click[0]


class ColorButton:
    def __init__(self, x, y, size, color):
        self.rect = pygame.Rect(x, y, size, size)
        self.color = color
        self.selected = False

    def draw(self, screen):
        pygame.draw.rect(screen, self.color, self.rect)
        if self.selected:
            pygame.draw.rect(screen, BLACK, self.rect, 3)
            pygame.draw.rect(screen, WHITE, self.rect.inflate(-6, -6), 2)
        else:
            pygame.draw.rect(screen, DARK_GRAY, self.rect, 2)

    def is_clicked(self, pos):
        return self.rect.collidepoint(pos)


class DrawingBoard:
    def __init__(self):
        # 创建画布
        self.canvas = pygame.Surface((DRAWING_AREA.width, DRAWING_AREA.height))
        self.canvas.fill(WHITE)
        self.canvas_rect = DRAWING_AREA
        self.last_pos = None
        self.current_color = BLACK
        self.brush_size = 5
        self.tool = "brush"  # brush, eraser, rectangle, circle, line
        self.start_pos = None
        self.temp_surface = None

    def draw(self, screen):
        # 绘制画布
        screen.blit(self.canvas, (TOOLBAR_WIDTH, 0))

        # 绘制临时图形（预览）
        if self.temp_surface:
            screen.blit(self.temp_surface, (TOOLBAR_WIDTH, 0))

        # 绘制画布边框
        pygame.draw.rect(screen, DARK_GRAY, self.canvas_rect, 3)

    def handle_mouse_down(self, pos):
        x, y = pos
        if self.canvas_rect.collidepoint(x, y):
            canvas_x = x - TOOLBAR_WIDTH
            canvas_y = y
            self.last_pos = (canvas_x, canvas_y)
            self.start_pos = (canvas_x, canvas_y)

            if self.tool == "brush" or self.tool == "eraser":
                color = WHITE if self.tool == "eraser" else self.current_color
                pygame.draw.circle(self.canvas, color,
                                   (canvas_x, canvas_y), self.brush_size)

    def handle_mouse_move(self, pos):
        if self.last_pos is None:
            return

        x, y = pos
        if self.canvas_rect.collidepoint(x, y):
            canvas_x = x - TOOLBAR_WIDTH
            canvas_y = y
            current_pos = (canvas_x, canvas_y)

            if self.tool == "brush" or self.tool == "eraser":
                color = WHITE if self.tool == "eraser" else self.current_color
                # 绘制两点之间的线条
                pygame.draw.line(self.canvas, color, self.last_pos,
                                 current_pos, self.brush_size * 2)
                self.last_pos = current_pos
            elif self.tool in ["rectangle", "circle", "line"]:
                # 预览图形
                self.temp_surface = self.canvas.copy()
                if self.tool == "rectangle":
                    rect = pygame.Rect(self.start_pos[0], self.start_pos[1],
                                       current_pos[0] - self.start_pos[0],
                                       current_pos[1] - self.start_pos[1])
                    pygame.draw.rect(self.temp_surface,
                                     self.current_color, rect, self.brush_size)
                elif self.tool == "circle":
                    radius = int(math.hypot(current_pos[0] - self.start_pos[0],
                                            current_pos[1] - self.start_pos[1]))
                    pygame.draw.circle(
                        self.temp_surface, self.current_color, self.start_pos, radius, self.brush_size)
                elif self.tool == "line":
                    pygame.draw.line(self.temp_surface, self.current_color,
                                     self.start_pos, current_pos, self.brush_size)
        else:
            self.last_pos = None

    def handle_mouse_up(self, pos):
        if self.start_pos is not None and self.tool in ["rectangle", "circle", "line"]:
            x, y = pos
            if self.canvas_rect.collidepoint(x, y):
                canvas_x = x - TOOLBAR_WIDTH
                canvas_y = y
                end_pos = (canvas_x, canvas_y)

                if self.tool == "rectangle":
                    rect = pygame.Rect(self.start_pos[0], self.start_pos[1],
                                       end_pos[0] - self.start_pos[0],
                                       end_pos[1] - self.start_pos[1])
                    pygame.draw.rect(
                        self.canvas, self.current_color, rect, self.brush_size)
                elif self.tool == "circle":
                    radius = int(math.hypot(end_pos[0] - self.start_pos[0],
                                            end_pos[1] - self.start_pos[1]))
                    pygame.draw.circle(
                        self.canvas, self.current_color, self.start_pos, radius, self.brush_size)
                elif self.tool == "line":
                    pygame.draw.line(self.canvas, self.current_color,
                                     self.start_pos, end_pos, self.brush_size)

        self.last_pos = None
        self.start_pos = None
        self.temp_surface = None

    def clear(self):
        self.canvas.fill(WHITE)

    def save(self):
        """保存画作"""
        if not os.path.exists("drawings"):
            os.makedirs("drawings")
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"drawings/drawing_{timestamp}.png"
        pygame.image.save(self.canvas, filename)
        return filename


def get_chinese_font(size):
    """获取支持中文的字体"""
    font_paths = [
        "C:/Windows/Fonts/simhei.ttf",
        "C:/Windows/Fonts/msyh.ttc",
        "/System/Library/Fonts/PingFang.ttc",
    ]
    for path in font_paths:
        if os.path.exists(path):
            try:
                return pygame.font.Font(path, size)
            except:
                continue
    return pygame.font.Font(None, size)


def main():
    screen = pygame.display.set_mode((WINDOW_WIDTH, WINDOW_HEIGHT))
    pygame.display.set_caption("画图板 - 绘画小游戏")
    clock = pygame.time.Clock()

    # 字体
    font = get_chinese_font(20)
    title_font = get_chinese_font(28)

    # 创建画板
    board = DrawingBoard()

    # 工具按钮
    tools = [
        ("画笔", "brush", 20, 100),
        ("橡皮", "eraser", 20, 150),
        ("矩形", "rectangle", 20, 200),
        ("圆形", "circle", 20, 250),
        ("直线", "line", 20, 300),
    ]

    tool_buttons = []
    for text, tool, x, y in tools:
        btn = Button(x, y, 160, 40, text, LIGHT_GRAY, (180, 180, 220), font)
        tool_buttons.append((btn, tool))

    # 颜色按钮
    colors = [
        (BLACK, 20, 380), (RED, 70, 380), (GREEN, 120, 380),
        (BLUE, 170, 380), (YELLOW, 20, 430), (PURPLE, 70, 430),
        (ORANGE, 120, 430), (PINK, 170, 430), (CYAN, 20, 480)
    ]

    color_buttons = []
    for color, x, y in colors:
        btn = ColorButton(x, y, 40, color)
        color_buttons.append(btn)

    # 功能按钮
    clear_button = Button(20, 540, 160, 45, "清空画板",
                          (200, 100, 100), (230, 130, 130), font)
    save_button = Button(20, 600, 160, 45, "保存图片",
                         (100, 200, 100), (130, 230, 130), font)

    # 笔刷大小滑块
    slider_x, slider_y = 20, 680
    slider_width, slider_height = 160, 20
    slider_knob_radius = 10
    dragging_slider = False

    # 状态显示
    status_text = ""
    status_timer = 0

    # 当前工具和颜色
    current_tool = "brush"
    board.current_color = BLACK

    # 标记颜色按钮选中状态
    for btn in color_buttons:
        btn.selected = (btn.color == board.current_color)

    running = True
    drawing = False

    while running:
        current_time = pygame.time.get_ticks()
        mouse_pos = pygame.mouse.get_pos()
        mouse_click = pygame.mouse.get_pressed()

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

            # 鼠标事件
            if event.type == pygame.MOUSEBUTTONDOWN:
                if event.button == 1:  # 左键
                    # 检查工具按钮
                    for btn, tool in tool_buttons:
                        if btn.is_clicked(mouse_pos, mouse_click):
                            current_tool = tool
                            board.tool = tool
                            status_text = f"切换到：{btn.text}"
                            status_timer = current_time

                    # 检查颜色按钮
                    for btn in color_buttons:
                        if btn.is_clicked(mouse_pos):
                            board.current_color = btn.color
                            for b in color_buttons:
                                b.selected = (b == btn)
                            status_text = f"选择颜色"
                            status_timer = current_time

                    # 检查功能按钮
                    if clear_button.is_clicked(mouse_pos, mouse_click):
                        board.clear()
                        status_text = "清空画板"
                        status_timer = current_time

                    if save_button.is_clicked(mouse_pos, mouse_click):
                        filename = board.save()
                        status_text = f"已保存：{filename}"
                        status_timer = current_time

                    # 检查滑块
                    knob_rect = pygame.Rect(slider_x + (board.brush_size / 50) * slider_width - slider_knob_radius,
                                            slider_y - slider_knob_radius,
                                            slider_knob_radius * 2, slider_knob_radius * 2)
                    if knob_rect.collidepoint(mouse_pos):
                        dragging_slider = True

                    # 开始绘画
                    if DRAWING_AREA.collidepoint(mouse_pos):
                        drawing = True
                        board.handle_mouse_down(mouse_pos)

                elif event.button == 4:  # 滚轮向上 - 增大笔刷
                    board.brush_size = min(50, board.brush_size + 2)
                    status_text = f"笔刷大小：{board.brush_size}"
                    status_timer = current_time
                elif event.button == 5:  # 滚轮向下 - 减小笔刷
                    board.brush_size = max(1, board.brush_size - 2)
                    status_text = f"笔刷大小：{board.brush_size}"
                    status_timer = current_time

            elif event.type == pygame.MOUSEBUTTONUP:
                drawing = False
                dragging_slider = False
                board.handle_mouse_up(mouse_pos)

            elif event.type == pygame.MOUSEMOTION:
                if dragging_slider:
                    # 更新滑块位置
                    x = max(slider_x, min(
                        slider_x + slider_width, mouse_pos[0]))
                    board.brush_size = int(
                        ((x - slider_x) / slider_width) * 50)
                    board.brush_size = max(1, min(50, board.brush_size))
                    status_text = f"笔刷大小：{board.brush_size}"
                    status_timer = current_time
                elif drawing and DRAWING_AREA.collidepoint(mouse_pos):
                    board.handle_mouse_move(mouse_pos)

        # 绘制界面
        screen.fill(BACKGROUND_COLOR)

        # 绘制工具栏背景
        toolbar_rect = pygame.Rect(0, 0, TOOLBAR_WIDTH, WINDOW_HEIGHT)
        pygame.draw.rect(screen, (220, 220, 230), toolbar_rect)
        pygame.draw.line(screen, DARK_GRAY, (TOOLBAR_WIDTH, 0),
                         (TOOLBAR_WIDTH, WINDOW_HEIGHT), 3)

        # 标题
        title = title_font.render("画图板", True, DARK_GRAY)
        screen.blit(title, (TOOLBAR_WIDTH // 2 - title.get_width() // 2, 20))

        # 工具标签
        tool_label = font.render("绘画工具", True, DARK_GRAY)
        screen.blit(tool_label, (20, 70))

        # 绘制工具按钮
        for btn, _ in tool_buttons:
            btn.draw(screen)

        # 颜色标签
        color_label = font.render("颜色选择", True, DARK_GRAY)
        screen.blit(color_label, (20, 355))

        # 绘制颜色按钮
        for btn in color_buttons:
            btn.draw(screen)

        # 绘制功能按钮
        clear_button.draw(screen)
        save_button.draw(screen)

        # 笔刷大小控制
        size_label = font.render("笔刷大小", True, DARK_GRAY)
        screen.blit(size_label, (20, 655))

        # 绘制滑块轨道
        pygame.draw.rect(screen, GRAY, (slider_x, slider_y,
                         slider_width, slider_height))
        # 绘制滑块
        knob_x = slider_x + (board.brush_size / 50) * slider_width
        pygame.draw.circle(
            screen, BLUE, (int(knob_x), slider_y + slider_height // 2), slider_knob_radius)
        pygame.draw.circle(screen, BLACK, (int(
            knob_x), slider_y + slider_height // 2), slider_knob_radius, 2)

        # 显示笔刷大小数值
        size_value = font.render(str(board.brush_size), True, DARK_GRAY)
        screen.blit(size_value, (slider_x + slider_width + 10, slider_y))

        # 显示当前工具
        current_tool_text = font.render(
            f"当前工具: {current_tool}", True, DARK_GRAY)
        screen.blit(current_tool_text, (20, 510))

        # 显示状态信息
        if status_text and current_time - status_timer < 2000:
            status_surface = font.render(status_text, True, (50, 150, 50))
            status_rect = status_surface.get_rect(
                center=(TOOLBAR_WIDTH // 2, 760))
            screen.blit(status_surface, status_rect)

        # 绘制画板
        board.draw(screen)

        # 显示提示
        tip_text = font.render("提示：鼠标滚轮调节笔刷大小", True, GRAY)
        screen.blit(tip_text, (TOOLBAR_WIDTH + 10, WINDOW_HEIGHT - 30))

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

    pygame.quit()
    sys.exit()


if __name__ == "__main__":
    main()