import gradio as gr
from PIL import Image, ImageDraw, ImageFont
import numpy as np

# Define color ranges using RGB bounds
COLOR_RANGES = {
    "7": {"name": "Red", "min": (150, 0, 0), "max": (255, 100, 100)},     # Red
    "B": {"name": "Yellow", "min": (200, 200, 0), "max": (255, 255, 100)}, # Yellow
    "H": {"name": "Blue", "min": (0, 0, 150), "max": (100, 100, 255)},    # Blue
    "9": {"name": "Orange", "min": (200, 100, 0), "max": (255, 180, 50)},  # Orange
    "A": {"name": "Yellow-Orange", "min": (200, 150, 0), "max": (255, 255, 100)},  # Light orange/yellow
    "D": {"name": "Green", "min": (0, 150, 0), "max": (100, 255, 100)},   # Green
    "F": {"name": "Cyan", "min": (0, 150, 150), "max": (100, 255, 255)},  # Cyan
    "I": {"name": "Dark Blue", "min": (0, 0, 100), "max": (50, 50, 255)},  # Dark blue
    "M": {"name": "Magenta", "min": (150, 0, 150), "max": (255, 100, 255)}, # Magenta
    "L": {"name": "Purple", "min": (100, 0, 150), "max": (180, 50, 255)},  # Purple
    "0": {"name": "Black", "min": (0, 0, 0), "max": (50, 50, 50)},         # Black
    "W": {"name": "White", "min": (200, 200, 200), "max": (255, 255, 255)}, # White
    "6": {"name": "Brown", "min": (100, 50, 0), "max": (180, 100, 50)},    # Brown
}

BLOCK_SIZE = 20


def rgb_in_range(rgb, min_rgb, max_rgb):
    """Check if RGB value is within range (inclusive)."""
    r, g, b = rgb
    mr, mg, mb = min_rgb
    Mr, Mg, Mb = max_rgb
    return mr <= r <= Mr and mg <= g <= Mg and mb <= b <= Mb


def find_nearest_color_key(rgb):
    """Find the best matching color key based on ranges."""
    for key, data in COLOR_RANGES.items():
        if rgb_in_range(rgb, data["min"], data["max"]):
            return key
    # If no match, use black as fallback
    return "0"


def process_image(img):
    # Convert to RGB
    img = img.convert("RGB")
    w, h = img.size

    # Resize to max 50x50 for performance
    MAX_DIM = 50
    if max(w, h) > MAX_DIM:
        scale = MAX_DIM / max(w, h)
        img = img.resize((int(w * scale), int(h * scale)))
        w, h = img.size

    # Create output image
    output_img = Image.new("RGB", (w * BLOCK_SIZE, h * BLOCK_SIZE), (255, 255, 255))
    draw = ImageDraw.Draw(output_img)
    font = ImageFont.load_default()

    for y in range(h):
        for x in range(w):
            pixel = img.getpixel((x, y))
            key = find_nearest_color_key(pixel)
            color_data = COLOR_RANGES[key]
            # Use average of min/max as representative color
            color = (
                (color_data["max"][0] + color_data["min"][0]) // 2,
                (color_data["max"][1] + color_data["min"][1]) // 2,
                (color_data["max"][2] + color_data["min"][2]) // 2
            )

            # Draw block
            x0 = x * BLOCK_SIZE
            y0 = y * BLOCK_SIZE
            x1 = x0 + BLOCK_SIZE
            y1 = y0 + BLOCK_SIZE
            draw.rectangle([x0, y0, x1, y1], fill=color, outline=(0, 0, 0))

            # Center text
            bbox = draw.textbbox((0, 0), key, font=font)
            w_text = bbox[2] - bbox[0]
            h_text = bbox[3] - bbox[1]
            draw.text(
                (x0 + (BLOCK_SIZE - w_text) / 2, y0 + (BLOCK_SIZE - h_text) / 2),
                key,
                fill=(0, 0, 0),
                font=font
            )

    return output_img


# Interface — NO EXAMPLES, just clean upload
iface = gr.Interface(
    fn=process_image,
    inputs=gr.Image(type="pil"),
    outputs=gr.Image(type="pil"),
    title="🎨 Color By Number Pixel Art",
    description="Upload any image → Get a printable color-by-number pixel art with symbols. Colors are grouped by range (e.g., all reds → '7')."
)

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