import gradio as gr
from PIL import Image
import numpy as np
import os
import tempfile

def process_gif(gif_file, target_device):
    # Load config
    if target_device == "OLED 128x64 (Mono)":
        width, height, fmt = 128, 64, "mono"
    else:
        width, height, fmt = 170, 320, "rgb565"

    # Open GIF
    gif = Image.open(gif_file.name)
    frames = []
    durations = []

    try:
        while True:
            frame = gif.copy().convert("RGBA")
            # Resize with aspect ratio
            frame.thumbnail((width, height), Image.LANCZOS)
            new_frame = Image.new("RGBA", (width, height), (0,0,0,0))
            paste_x = (width - frame.width) // 2
            paste_y = (height - frame.height) // 2
            new_frame.paste(frame, (paste_x, paste_y), frame)

            # Convert
            if fmt == "mono":
                gray = new_frame.convert('L')
                bw = gray.point(lambda x: 0 if x < 128 else 255, '1')
                byte_data = image_to_monochrome_bytes(bw, width, height)
            else:
                rgb_frame = new_frame.convert('RGB')
                byte_data = image_to_rgb565_bytes(rgb_frame, width, height)

            frames.append(byte_data)
            durations.append(gif.info.get('duration', 100))
            gif.seek(gif.tell() + 1)
    except EOFError:
        pass

    # Generate .h content
    device_name = "oled_128x64" if fmt == "mono" else "tft_170x320"
    h_content = f"// Generated for {target_device}\n"
    h_content += f"#define {device_name.upper()}_FRAME_COUNT {len(frames)}\n"
    h_content += f"#define {device_name.upper()}_WIDTH {width}\n"
    h_content += f"#define {device_name.upper()}_HEIGHT {height}\n\n"

    for i, frame_bytes in enumerate(frames):
        h_content += f"const uint8_t {device_name}_frame_{i}[] = {{\n    "
        byte_strs = [f"0x{b:02X}" for b in frame_bytes]
        lines = [", ".join(byte_strs[j:j+16]) for j in range(0, len(byte_strs), 16)]
        h_content += ",\n    ".join(lines) + "\n};\n\n"

    h_content += f"const uint8_t* const {device_name}_frames[] = {{\n    "
    h_content += ",\n    ".join(f"{device_name}_frame_{i}" for i in range(len(frames)))
    h_content += "\n};\n"

    h_content += f"\nconst uint16_t {device_name}_durations[] = {{"
    h_content += ", ".join(str(d) for d in durations) + "};\n"

    # Save to temp file
    temp_dir = tempfile.mkdtemp()
    output_path = os.path.join(temp_dir, f"{device_name}.h")
    with open(output_path, "w") as f:
        f.write(h_content)

    return output_path

def image_to_monochrome_bytes(img, width, height):
    img_data = np.array(img).astype(np.uint8)
    byte_array = []
    for col in range(width):
        for page in range(0, height, 8):
            byte = 0
            for bit in range(8):
                row = page + bit
                if row < height:
                    pixel = img_data[row][col]
                    if pixel > 0:
                        byte |= (1 << bit)
            byte_array.append(byte)
    return bytes(byte_array)

def image_to_rgb565_bytes(img, width, height):
    pixels = np.array(img)
    byte_array = bytearray()
    for y in range(height):
        for x in range(width):
            r, g, b = pixels[y, x][:3]
            r5 = (r >> 3) & 0x1F
            g6 = (g >> 2) & 0x3F
            b5 = (b >> 3) & 0x1F
            rgb565 = (r5 << 11) | (g6 << 5) | b5
            byte_array.append(rgb565 & 0xFF)
            byte_array.append((rgb565 >> 8) & 0xFF)
    return bytes(byte_array)

# Gradio UI
demo = gr.Interface(
    fn=process_gif,
    inputs=[
        gr.File(label="Upload GIF"),
        gr.Dropdown(
            choices=["OLED 128x64 (Mono)", "TFT 170x320 (RGB565)"],
            label="Target Display"
        )
    ],
    outputs=gr.File(label="Download .h File"),
    title="🎨 GIF to ESP32 Display Converter",
    description="Convert any GIF to C byte arrays for ESP32 OLED or TFT displays. Outputs ready-to-use .h files."
)

if __name__ == "__main__":
    demo.launch()