import gradio as gr
from PIL import Image, ImageEnhance, ImageOps, ImageColor
import numpy as np
import cv2
import re

# Helper function to safely parse color
def parse_color(color):
    if isinstance(color, str):
        if color.startswith("rgba"):
            match = re.match(r"rgba\(([\d.]+), *([\d.]+), *([\d.]+), *[\d.]+\)", color)
            if match:
                r, g, b = match.groups()
                return (int(float(r)), int(float(g)), int(float(b)))
        elif color.startswith("#"):
            return color  # valid hex
    return "white"  # fallback

# Main image beautification function
def beautify_image(image, tint_color, contrast, brightness, saturation, apply_filter):
    img = image.convert("RGB")

    # Apply tint color if provided
    tint = parse_color(tint_color)
    if tint != "None":
        r, g, b = Image.new("RGB", img.size, tint).split()
        img = Image.blend(img, Image.merge("RGB", (r, g, b)), alpha=0.3)

    # Enhance image attributes
    img = ImageEnhance.Contrast(img).enhance(contrast)
    img = ImageEnhance.Brightness(img).enhance(brightness)
    img = ImageEnhance.Color(img).enhance(saturation)

    # Optional color filters
    if apply_filter != "None":
        img_cv = np.array(img)
        if apply_filter == "Sepia":
            kernel = np.array([[0.272, 0.534, 0.131],
                               [0.349, 0.686, 0.168],
                               [0.393, 0.769, 0.189]])
            img_cv = cv2.transform(img_cv, kernel)
            img_cv = np.clip(img_cv, 0, 255)
        else:
            try:
                cmap = getattr(cv2, f'COLORMAP_{apply_filter.upper()}')
                img_cv = cv2.applyColorMap(cv2.cvtColor(img_cv, cv2.COLOR_RGB2GRAY), cmap)
                img_cv = cv2.cvtColor(img_cv, cv2.COLOR_BGR2RGB)
            except AttributeError:
                pass  # invalid colormap, skip
        img = Image.fromarray(img_cv.astype(np.uint8))

    return img

# Gradio UI
demo = gr.Interface(
    fn=beautify_image,
    inputs=[
        gr.Image(type="pil", label="Upload Black & White Verse Image"),
        gr.ColorPicker(label="Tint Color", value="#FFD700"),
        gr.Slider(0.5, 2.0, 1.2, label="Contrast"),
        gr.Slider(0.5, 2.0, 1.1, label="Brightness"),
        gr.Slider(0.0, 2.0, 1.0, label="Saturation"),
        gr.Dropdown(["None", "Sepia", "JET", "OCEAN", "PINK", "AUTUMN"], label="Optional Filter"),
    ],
    outputs=gr.Image(label="Beautified Image"),
    title="🖌️ VerseGlow — Beautify Your B&W Verse Screenshots",
    description="Upload black-and-white verse screenshots and make them visually beautiful with tint, contrast, and filters.",
)

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