import pygame
import sys

# Initialize Pygame
pygame.init()

# Game Constants
WIDTH, HEIGHT = 800, 600
GRID_SIZE = 20  # Size of each pixel in the grid
GRID_WIDTH = WIDTH // GRID_SIZE
GRID_HEIGHT = HEIGHT // GRID_SIZE
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Pixel Art Drawing Game")
clock = pygame.time.Clock()

# Colors (Expanded Palette)
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
RED = (255, 0, 0)
GREEN = (0, 255, 0)
BLUE = (0, 0, 255)
YELLOW = (255, 255, 0)
CYAN = (0, 255, 255)
MAGENTA = (255, 0, 255)
ORANGE = (255, 165, 0)
PURPLE = (128, 0, 128)
PINK = (255, 105, 180)
BROWN = (139, 69, 19)
GRAY = (169, 169, 169)
LIME = (0, 255, 0)
TURQUOISE = (64, 224, 208)
GOLD = (255, 215, 0)

COLOR_PALETTE = [
    BLACK, RED, GREEN, BLUE, YELLOW, CYAN, MAGENTA, ORANGE, 
    PURPLE, PINK, BROWN, GRAY, LIME, TURQUOISE, GOLD, WHITE
]

# Canvas (grid)
canvas = [[WHITE for _ in range(GRID_WIDTH)] for _ in range(GRID_HEIGHT)]

# Current drawing color
current_color = BLACK

# Eraser mode flag
eraser_mode = False

# Functions
def draw_grid():
    for y in range(GRID_HEIGHT):
        for x in range(GRID_WIDTH):
            pygame.draw.rect(screen, canvas[y][x], (x * GRID_SIZE, y * GRID_SIZE, GRID_SIZE, GRID_SIZE))
            pygame.draw.rect(screen, BLACK, (x * GRID_SIZE, y * GRID_SIZE, GRID_SIZE, GRID_SIZE), 1)  # Grid lines

def draw_palette():
    for i, color in enumerate(COLOR_PALETTE):
        if color == BLACK:
            pygame.draw.rect(screen, color, (WIDTH - 100, 50 + i * 40, 40, 40))
            pygame.draw.rect(screen, (255, 255, 255), (WIDTH - 100, 50 + i * 40, 40, 40), 5)  # White border for black color
        else:
            pygame.draw.rect(screen, color, (WIDTH - 100, 50 + i * 40, 40, 40))

    # Draw eraser button
    pygame.draw.rect(screen, WHITE, (WIDTH - 100, HEIGHT - 50, 40, 40))
    pygame.draw.rect(screen, BLACK, (WIDTH - 100, HEIGHT - 50, 40, 40), 3)  # Eraser button border
    pygame.draw.line(screen, BLACK, (WIDTH - 95, HEIGHT - 45), (WIDTH - 85, HEIGHT - 35), 3)  # Draw an "X" for eraser

def handle_drawing():
    global canvas
    mouse_x, mouse_y = pygame.mouse.get_pos()
    grid_x = mouse_x // GRID_SIZE
    grid_y = mouse_y // GRID_SIZE
    if 0 <= grid_x < GRID_WIDTH and 0 <= grid_y < GRID_HEIGHT:
        if pygame.mouse.get_pressed()[0]:  # Left mouse button to draw
            if eraser_mode:
                canvas[grid_y][grid_x] = WHITE  # Erase by setting to white
            else:
                canvas[grid_y][grid_x] = current_color  # Draw with selected color

def handle_palette_click():
    global current_color, eraser_mode
    mouse_x, mouse_y = pygame.mouse.get_pos()
    if WIDTH - 100 <= mouse_x <= WIDTH - 60:
        for i, color in enumerate(COLOR_PALETTE):
            if 50 + i * 40 <= mouse_y <= 90 + i * 40:
                current_color = color
                eraser_mode = False  # Turn off eraser mode when selecting a color
    # Check if the eraser button was clicked
    if WIDTH - 100 <= mouse_x <= WIDTH - 60 and HEIGHT - 50 <= mouse_y <= HEIGHT - 10:
        eraser_mode = not eraser_mode  # Toggle eraser mode

def save_canvas():
    pygame.image.save(screen, "pixel_art.png")
    print("Saved the artwork as 'pixel_art.png'!")

def draw():
    screen.fill(WHITE)
    draw_grid()
    draw_palette()
    handle_drawing()
    pygame.display.flip()

# Game Loop
def game_loop():
    running = True
    while running:
        clock.tick(60)  # 60 FPS

        # Handle events
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
            if event.type == pygame.MOUSEBUTTONDOWN:
                handle_palette_click()
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_s:  # Press 'S' to save the canvas
                    save_canvas()

        draw()

    pygame.quit()
    sys.exit()

# Start the game
game_loop()