import gradio as gr
from PIL import Image, ImageDraw, ImageFont, ImageGrab, ImageEnhance
from requests import get
from sty import bg, rs
import numpy as np
from typing import Optional, Sequence, Tuple

# ASCII characters from dark to light
CONVERSION_CHARACTERS = "@%#*+=-:. "

# Use the default font provided by PIL
base_font = ImageFont.load_default()

def sizeof(text: str, font: ImageFont.FreeTypeFont = base_font) -> Tuple[int, int]:
    draw = ImageDraw.Draw(Image.new("RGB", (1, 1)))
    _, _, width, height = draw.textbbox((0, 0), text, font)
    return width, height

def ascii_to_image(text: str, font: ImageFont.FreeTypeFont = base_font) -> Image.Image:
    text_image = Image.new("RGB", sizeof(text, font))
    draw = ImageDraw.Draw(text_image)
    draw.text((0, 0), text, (255, 255, 255), font)
    return text_image

def get_brightness_of_char(char: str, font: ImageFont.FreeTypeFont = base_font) -> int:
    image = ascii_to_image(char, font)
    return (np.array(image) != 0).sum().item()

sorted_letters = sorted(
    CONVERSION_CHARACTERS,
    key=lambda char: (
        get_brightness_of_char(char),
        char
    )
)

def image_to_ascii(
    image: Image.Image | str,
    size: Optional[Tuple[int, int]] = None,
    charset: Optional[Sequence[str]] = None,
    fix_scaling: bool = True,
    scale: float | Tuple[float, float] = 1,
    sharpness: float = 1,
    brightness: float = 1,
    sort_chars: bool = False,
    colorful: bool = False,
) -> str:
    if isinstance(image, str):
        if image.lower() in ("clip", "clipboard"):
            image = ImageGrab.grabclipboard()
            if not isinstance(image, Image.Image):
                raise ValueError("Unable to load image from clipboard")
        else:
            try:
                image = Image.open(image)
            except FileNotFoundError:
                raise ValueError("Unable to load image from path")
            except Exception:
                try:
                    image = Image.open(get(image, stream=True).raw)
                except Exception:
                    raise ValueError("Unable to load image from URL")
    if not isinstance(image, Image.Image):
        raise ValueError("Invalid image path or URL")

    image = image.convert("RGB")

    if sort_chars and charset:
        charset = sorted(
            charset,
            key=lambda char: (
                get_brightness_of_char(char),
                char
            )
        )
    charset = charset or sorted_letters

    image_width, image_height = size or image.size

    if isinstance(scale, (int, float)):
        scale = (scale,) * 2
    image_width = int(image_width * scale[0] * (bool(fix_scaling) + 1))
    image_height = int(image_height * scale[1])

    scaled_image = image.resize((image_width, image_height))
    brightened_image = ImageEnhance.Brightness(scaled_image).enhance(brightness)
    sharpened_image = ImageEnhance.Sharpness(brightened_image).enhance(sharpness)
    image_array = np.array(sharpened_image.convert("L"), dtype=int) * len(charset) // 256
    ascii_converted = np.vectorize(charset.__getitem__)(image_array)

    output = "\n".join(map("".join, ascii_converted))
    if colorful:
        colors = []
        for row in np.array(sharpened_image.convert("RGB")):
            for pixel in row:
                colors.append(bg(*pixel))
            colors.append(rs.bg)
        output = "".join(color + char for color, char in zip(colors, output))
    return output

def gradio_image_to_ascii(image, size, scale, sharpness, brightness, sort_chars, colorful):
    try:
        ascii_art = image_to_ascii(
            image,
            size=(size, size),  # Assuming size is a single value for simplicity
            scale=scale,
            sharpness=sharpness,
            brightness=brightness,
            sort_chars=sort_chars,
            colorful=colorful
        )
        return ascii_art
    except Exception as e:
        return str(e)

# Gradio interface
demo = gr.Interface(
    fn=gradio_image_to_ascii,
    inputs=[
        gr.Image(type="pil", label="Upload Image"),
        gr.Slider(10, 200, value=50, step=10, label="Size"),
        gr.Slider(0.1, 2.0, value=1.0, step=0.1, label="Scale"),
        gr.Slider(0.1, 2.0, value=1.0, step=0.1, label="Sharpness"),
        gr.Slider(0.1, 2.0, value=1.0, step=0.1, label="Brightness"),
        gr.Checkbox(label="Sort Characters by Brightness"),
        gr.Checkbox(label="Colorful Output"),
    ],
    outputs=gr.Textbox(label="ASCII Art"),
    title="Image to ASCII Generator",
    description="Upload an image and adjust parameters to generate ASCII art.",
)

demo.launch()