#!/usr/bin/env python3

import pygame
from pygame import Color, Rect
from pygame.locals import *
import datetime
import array
import requests

def upload_file(url, filename):
    with open(filename, "rb") as file:
        response = requests.post(url, files={"file": file})
    
    # Check the response status code and handle accordingly
    if response.status_code == 200:
        print("File upload successful")
    else:
        print("File upload failed")

# Initialize pygame
pygame.init()

# Set up the display
pixel_size = 40
width = 16 * pixel_size
height = 16 * pixel_size + 100  # Increased height for the color palette
screen = pygame.display.set_mode((width, height))

# Create a 16x16 pixel array to store the colors
pixels = [[(0, 0, 0) for _ in range(16)] for _ in range(16)]

# Create a stack to store previous pixel states for undo
undo_stack = []

# Set up the drawing flag
drawing = False

# Color palette (24-bit RGB colors)
palette = [
    Color(0, 0, 0),          # Black
    Color(255, 255, 255),    # White
    Color(255, 0, 0),        # Red
    Color(0, 255, 0),        # Green
    Color(0, 0, 255),        # Blue
    Color(255, 255, 0),      # Yellow
    Color(255, 0, 255),      # Magenta
    Color(0, 255, 255)       # Cyan
]

# Set up the color palette dimensions
palette_x = 20
palette_y = height - 80
palette_width = width - 40
palette_height = 60
color_width = palette_width // len(palette)

# Set up the initial selected color
selected_color = palette[0]

# Main loop
running = True
while running:
    for event in pygame.event.get():
        if event.type == QUIT:
            running = False
        elif event.type == MOUSEBUTTONDOWN:
            if event.button == 1:  # Left mouse button
                x, y = event.pos
                # Check if the click is inside the color palette
                if palette_y <= y < palette_y + palette_height:
                    color_index = (x - palette_x) // color_width
                    # Update the selected color for drawing
                    if 0 <= color_index < len(palette):
                        selected_color = palette[color_index]
                else:
                    row = (y - 20) // pixel_size
                    col = x // pixel_size
                    if 0 <= row < 16 and 0 <= col < 16:
                        # Store the current pixel state for undo
                        undo_stack.append([row, col, pixels[row][col]])
                        drawing = True
                        pixels[row][col] = selected_color
        elif event.type == MOUSEBUTTONUP:
            if event.button == 1:  # Left mouse button
                drawing = False
        elif event.type == MOUSEMOTION:
            if drawing:
                x, y = event.pos
                row = (y - 20) // pixel_size
                col = x // pixel_size
                if 0 <= row < 16 and 0 <= col < 16:
                    # Store the current pixel state for undo
                    undo_stack.append([row, col, pixels[row][col]])
                    pixels[row][col] = selected_color
        elif event.type == KEYDOWN:
            if event.key == K_s and pygame.key.get_mods() & KMOD_CTRL:
                # Create a new surface with the pixel art
                pixel_art = pygame.Surface((width, height))
                for row in range(16):
                    for col in range(16):
                        pygame.draw.rect(pixel_art, pixels[row][col], Rect(col * pixel_size, row * pixel_size + 20, pixel_size, pixel_size))

                # Save the pixel art as an image
                timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
                filename = f"art.png"
                pygame.image.save(pixel_art, filename, "PNG")
                print(f"Pixel art saved as {filename}!")

                # Convert RGB values to binary format
                bitstream = array.array("B")
                for y in range(16):
                    if y % 2 == 0:  # Every second line (odd-indexed line)
                        pixels_row = pixels[y][::-1]  # Flip the line
                    else:
                        pixels_row = pixels[y]

                    # Iterate over the RGB values of the line
                    for pixel in pixels_row:
                        r, g, b, *_ = pixel  # Unpack the first three values and ignore the rest

                        # Apply gamma correction if needed (optional)
                        # r = int(pow(r / 255, 2.8) * 255)
                        # g = int(pow(g / 255, 2.8) * 255)
                        # b = int(pow(b / 255, 2.8) * 255)

                        # Write the color components to the binary array in GRB order for WS2812
                        bitstream.extend([g, r, b])

                # Save the binary data as a file
                filename = f"art.grb"
                with open(filename, "wb") as file:
                    file.write(bitstream)

                upload_file("http://192.168.1.106", "art.grb")
                print("Uploaded")
                


                print(f"Pixel art saved as {filename}!")

            if event.key == K_u and pygame.key.get_mods() & KMOD_CTRL:
                # Undo the previous modification if available
                if undo_stack:
                    row, col, color = undo_stack.pop()
                    pixels[row][col] = color
                    print("Undo")

    # Clear the screen
    screen.fill((255, 255, 255))

    # Draw the outline of the canvas
    pygame.draw.rect(screen, (0, 0, 0), Rect(0, 20, width, height - 100), 1)

    # Draw the color palette
    for i, color in enumerate(palette):
        pygame.draw.rect(screen, color, Rect(palette_x + i * color_width, palette_y, color_width, palette_height))

    # Draw the pixels on the screen
    for row in range(16):
        for col in range(16):
            pygame.draw.rect(screen, pixels[row][col], Rect(col * pixel_size, row * pixel_size + 20, pixel_size, pixel_size))

    pygame.display.update()

# Quit the program
pygame.quit()