import pygame
import sys

# 初始化pygame
pygame.init()
WIDTH, HEIGHT = 800, 600
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Python简易画图板")

# 颜色定义
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
RED = (255, 0, 0)
GREEN = (0, 255, 0)
BLUE = (0, 0, 255)
YELLOW = (255, 255, 0)

# 初始参数
screen.fill(WHITE)
current_color = BLACK
brush_size = 5
is_drawing = False
last_pos = None

clock = pygame.time.Clock()

while True:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            pygame.quit()
            sys.exit()

        # 鼠标按下 开始绘画
        if event.type == pygame.MOUSEBUTTONDOWN:
            is_drawing = True
            last_pos = event.pos
        # 鼠标松开 停止绘画
        if event.type == pygame.MOUSEBUTTONUP:
            is_drawing = False
            last_pos = None
        # 鼠标移动画线
        if event.type == pygame.MOUSEMOTION and is_drawing:
            if last_pos:
                pygame.draw.line(screen, current_color, last_pos, event.pos, brush_size)
            last_pos = event.pos

        # 键盘快捷键
        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_1:
                current_color = BLACK
            elif event.key == pygame.K_2:
                current_color = RED
            elif event.key == pygame.K_3:
                current_color = GREEN
            elif event.key == pygame.K_4:
                current_color = BLUE
            elif event.key == pygame.K_5:
                current_color = YELLOW
            elif event.key == pygame.K_e:
                # 橡皮擦
                current_color = WHITE
            elif event.key == pygame.K_c:
                # 清空画布
                screen.fill(WHITE)
            elif event.key == pygame.K_UP:
                brush_size += 2
            elif event.key == pygame.K_DOWN:
                if brush_size > 1:
                    brush_size -= 2

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