"""
🎨 Style Transfer
Versão otimizada para Hugging Face Spaces
"""

import gradio as gr
from PIL import Image, ImageFilter, ImageOps, ImageEnhance, ImageDraw
import numpy as np
import tempfile
import time

print("🚀 Iniciando Style Transfer...")

MAX_IMAGE_SIZE = 1024
STYLES = {
    "van_gogh": "🎨 Van Gogh",
    "monet": "🌅 Monet",
    "picasso": "🧑‍🎨 Picasso",
    "warhol": "👑 Warhol",
    "hokusai": "🌊 Hokusai",
    "starry_night": "✨ Noite Estrelada",
    "oil_painting": "🖼️ Pintura a Óleo",
    "watercolor": "🎨 Aquarela",
    "sketch": "✏️ Desenho a Lápis",
    "comic": "💥 Estilo HQ"
}

def log_message(message):
    timestamp = time.strftime("%H:%M:%S")
    print(f"[{timestamp}] {message}")

def validate_image(image):
    try:
        if image is None:
            return None, "❌ Nenhuma imagem fornecida"
        
        if isinstance(image, np.ndarray):
            img = Image.fromarray(image.astype('uint8'))
        else:
            img = image
        
        if max(img.size) > 4000:
            return None, "❌ Imagem muito grande"
        if min(img.size) < 32:
            return None, "❌ Imagem muito pequena"
        
        log_message(f"✅ Imagem válida: {img.size}px")
        return img, "ok"
    except Exception as e:
        return None, str(e)

def resize_image(image, max_size):
    if max(image.size) <= max_size:
        return image
    ratio = max_size / max(image.size)
    new_size = (int(image.width * ratio), int(image.height * ratio))
    return image.resize(new_size, Image.Resampling.LANCZOS)

def apply_van_gogh_style(image, intensity=0.8):
    try:
        img = image.convert('RGB')
        enhancer = ImageEnhance.Color(img)
        img = enhancer.enhance(1.0 + intensity * 0.5)
        enhancer = ImageEnhance.Contrast(img)
        img = enhancer.enhance(1.0 + intensity * 0.3)
        
        result = img.copy()
        for angle in [0, 45, 90, 135]:
            motion = img.filter(ImageFilter.GaussianBlur(radius=1))
            motion = motion.rotate(angle, expand=False)
            result = Image.blend(result, motion, alpha=0.1)
        
        return result
    except Exception as e:
        log_message(f"Erro Van Gogh: {e}")
        return image

def apply_monet_style(image, intensity=0.8):
    try:
        img = image.convert('RGB')
        img = img.filter(ImageFilter.GaussianBlur(radius=1 + intensity))
        enhancer = ImageEnhance.Brightness(img)
        img = enhancer.enhance(1.0 + intensity * 0.2)
        enhancer = ImageEnhance.Color(img)
        return enhancer.enhance(0.8 + intensity * 0.3)
    except Exception as e:
        return image

def apply_picasso_style(image, intensity=0.8):
    try:
        img = image.convert('RGB')
        width, height = img.size
        result = Image.new('RGB', (width, height))
        draw = ImageDraw.Draw(result)
        
        block_size = max(10, int(50 * (1 - intensity * 0.5)))
        
        for y in range(0, height, block_size):
            for x in range(0, width, block_size):
                block = img.crop((x, y, min(x + block_size, width), min(y + block_size, height)))
                avg_color = block.resize((1, 1)).getpixel((0, 0))
                
                shape = (x + y) % 3
                if shape == 0:
                    points = [(x, y), (x + block_size, y), (x + block_size // 2, y + block_size)]
                    draw.polygon(points, fill=avg_color)
                elif shape == 1:
                    offset = block_size // 4
                    points = [(x + offset, y), (x + block_size, y), 
                             (x + block_size - offset, y + block_size), (x, y + block_size)]
                    draw.polygon(points, fill=avg_color)
                else:
                    draw.ellipse([x, y, x + block_size, y + block_size], fill=avg_color)
        
        return Image.blend(result, img, alpha=0.3)
    except Exception as e:
        return image

def apply_warhol_style(image, intensity=0.8):
    try:
        img = image.convert('RGB')
        gray = img.convert('L')
        gray = gray.point(lambda x: 0 if x < 64 else 85 if x < 128 else 170 if x < 192 else 255)
        
        width, height = img.size
        result = Image.new('RGB', (width, height))
        
        palettes = [
            [(255, 0, 0), (255, 255, 0), (0, 255, 255), (255, 255, 255)],
            [(0, 255, 0), (255, 0, 255), (255, 255, 0), (0, 0, 0)],
            [(0, 0, 255), (255, 165, 0), (255, 192, 203), (255, 255, 255)],
            [(148, 0, 211), (255, 20, 147), (0, 255, 127), (255, 255, 0)]
        ]
        
        palette = palettes[int(intensity * 3) % 4]
        
        for y in range(height):
            for x in range(width):
                pixel = gray.getpixel((x, y))
                color = palette[0 if pixel < 64 else 1 if pixel < 128 else 2 if pixel < 192 else 3]
                result.putpixel((x, y), color)
        
        enhancer = ImageEnhance.Contrast(result)
        return enhancer.enhance(1.5)
    except Exception as e:
        return image

def apply_hokusai_style(image, intensity=0.8):
    try:
        img = image.convert('RGB')
        img = img.convert('P', palette=Image.ADAPTIVE, colors=32).convert('RGB')
        img = img.filter(ImageFilter.SMOOTH_MORE)
        
        r, g, b = img.split()
        b = b.point(lambda x: min(255, int(x * 1.3)))
        r = r.point(lambda x: int(x * 0.8))
        return Image.merge('RGB', (r, g, b))
    except Exception as e:
        return image

def apply_starry_night_style(image, intensity=0.8):
    try:
        img = image.convert('RGB')
        enhancer = ImageEnhance.Brightness(img)
        img = enhancer.enhance(0.7)
        enhancer = ImageEnhance.Contrast(img)
        img = enhancer.enhance(1.3)
        
        r, g, b = img.split()
        b = b.point(lambda x: min(255, int(x * 1.2)))
        r = r.point(lambda x: int(x * 0.9))
        return Image.merge('RGB', (r, g, b))
    except Exception as e:
        return image

def apply_oil_painting_style(image, intensity=0.8):
    try:
        img = image.convert('RGB')
        img = img.filter(ImageFilter.GaussianBlur(radius=0.5))
        enhancer = ImageEnhance.Color(img)
        img = enhancer.enhance(1.0 + intensity * 0.4)
        enhancer = ImageEnhance.Contrast(img)
        return enhancer.enhance(1.0 + intensity * 0.3)
    except Exception as e:
        return image

def apply_watercolor_style(image, intensity=0.8):
    try:
        img = image.convert('RGB')
        img = img.filter(ImageFilter.GaussianBlur(radius=1 + intensity))
        enhancer = ImageEnhance.Color(img)
        img = enhancer.enhance(0.8 + intensity * 0.2)
        enhancer = ImageEnhance.Brightness(img)
        return enhancer.enhance(1.1)
    except Exception as e:
        return image

def apply_sketch_style(image, intensity=0.8):
    try:
        img = image.convert('RGB')
        gray = img.convert('L')
        inverted = ImageOps.invert(gray)
        blurred = inverted.filter(ImageFilter.GaussianBlur(radius=2 + intensity))
        
        result = Image.new('L', img.size)
        for y in range(img.height):
            for x in range(img.width):
                a = gray.getpixel((x, y))
                b = blurred.getpixel((x, y))
                result.putpixel((x, y), min(255, (a * 255) // (255 - b) if b < 255 else 255))
        
        return result.convert('RGB')
    except Exception as e:
        return image

def apply_comic_style(image, intensity=0.8):
    try:
        img = image.convert('RGB')
        enhancer = ImageEnhance.Contrast(img)
        img = enhancer.enhance(2.0)
        enhancer = ImageEnhance.Color(img)
        img = enhancer.enhance(1.5)
        
        edges = img.filter(ImageFilter.FIND_EDGES).convert('L')
        edges = edges.point(lambda x: 0 if x < 50 else 255)
        
        return Image.composite(Image.new('RGB', img.size, (0, 0, 0)), img, edges)
    except Exception as e:
        return image

def apply_style(image, style_name, intensity=0.8):
    styles = {
        "van_gogh": apply_van_gogh_style,
        "monet": apply_monet_style,
        "picasso": apply_picasso_style,
        "warhol": apply_warhol_style,
        "hokusai": apply_hokusai_style,
        "starry_night": apply_starry_night_style,
        "oil_painting": apply_oil_painting_style,
        "watercolor": apply_watercolor_style,
        "sketch": apply_sketch_style,
        "comic": apply_comic_style
    }
    return styles.get(style_name, lambda i, _: i)(image, intensity)

def create_comparison(original, styled, style_name):
    try:
        if original is None or styled is None:
            return None
        
        target_height = 400
        orig_w = int(original.width * (target_height / original.height))
        styled_w = int(styled.width * (target_height / styled.height))
        
        orig_r = original.resize((orig_w, target_height), Image.Resampling.LANCZOS)
        styled_r = styled.resize((styled_w, target_height), Image.Resampling.LANCZOS)
        
        total_w = orig_r.width + styled_r.width + 20
        comp = Image.new('RGB', (total_w, target_height + 60), (240, 240, 240))
        draw = ImageDraw.Draw(comp)
        
        draw.text((10, 10), "ORIGINAL", fill=(100, 100, 100))
        draw.text((orig_r.width + 30, 10), f"ESTILO: {style_name.upper()}", fill=(0, 100, 200))
        
        comp.paste(orig_r, (0, 40))
        comp.paste(styled_r, (orig_r.width + 20, 40))
        
        return comp
    except Exception as e:
        return None

def save_image(image):
    try:
        if image is None:
            return None
        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
        image.save(temp_file.name, "PNG", optimize=True)
        return temp_file.name
    except Exception as e:
        return None

# Interface Gradio
with gr.Blocks(title="🎨 Style Transfer", theme=gr.themes.Soft()) as demo:
    
    gr.HTML("""
    <div style="text-align: center; padding: 20px; background: linear-gradient(135deg, #8A2BE2 0%, #4B0082 100%); 
                border-radius: 10px; color: white; margin-bottom: 20px;">
        <h1 style="margin: 0;">🎨 Style Transfer</h1>
        <p style="margin: 5px 0;">Transforme suas fotos em obras de arte</p>
    </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 📤 Upload")
            image_input = gr.Image(type="pil", label="Sua imagem")
            
            gr.Markdown("### 🎨 Estilo")
            with gr.Row():
                s1 = gr.Button("🎨 Van Gogh", size="sm")
                s2 = gr.Button("🌅 Monet", size="sm")
            with gr.Row():
                s3 = gr.Button("🧑‍🎨 Picasso", size="sm")
                s4 = gr.Button("👑 Warhol", size="sm")
            with gr.Row():
                s5 = gr.Button("🌊 Hokusai", size="sm")
                s6 = gr.Button("✨ Noite Estrelada", size="sm")
            with gr.Row():
                s7 = gr.Button("🖼️ Óleo", size="sm")
                s8 = gr.Button("🎨 Aquarela", size="sm")
            with gr.Row():
                s9 = gr.Button("✏️ Lápis", size="sm")
                s10 = gr.Button("💥 HQ", size="sm")
            
            style_selected = gr.Textbox(value="van_gogh", visible=False)
            
            gr.Markdown("### ⚙️ Intensidade")
            intensity = gr.Slider(0.1, 1.0, 0.8, step=0.1, label="")
            
            apply_btn = gr.Button("✨ Aplicar Estilo", variant="primary", size="lg")
        
        with gr.Column(scale=1):
            status = gr.Markdown("**Status:** Aguardando imagem...")
            
            with gr.Tabs():
                with gr.TabItem("🔄 Comparação"):
                    comp_out = gr.Image(type="pil", label="Antes/Depois")
                with gr.TabItem("📷 Original"):
                    orig_out = gr.Image(type="pil", label="Original")
                with gr.TabItem("🎨 Estilizada"):
                    styled_out = gr.Image(type="pil", label="Estilizada")
            
            gr.Markdown("### 💾 Download")
            with gr.Row():
                dl_styled = gr.Button("📥 Baixar Estilizada")
                dl_comp = gr.Button("📊 Baixar Comparação")
            
            dl_file = gr.File(label="Arquivo", interactive=False)
    
    with gr.Accordion("📖 Guia", open=False):
        gr.Markdown("""
        ## Estilos Disponíveis:
        - **Van Gogh**: Pinceladas expressivas
        - **Monet**: Impressionismo suave
        - **Picasso**: Cubismo geométrico
        - **Warhol**: Pop Art vibrante
        - **Hokusai**: Ukiyo-e japonês
        - **Noite Estrelada**: Tons azulados
        - **Pintura a Óleo**: Textura clássica
        - **Aquarela**: Suave e transparente
        - **Desenho a Lápis**: Contornos
        - **Estilo HQ**: Bordas e cores planas
        """)
    
    # Eventos
    s1.click(fn=lambda: "van_gogh", outputs=[style_selected])
    s2.click(fn=lambda: "monet", outputs=[style_selected])
    s3.click(fn=lambda: "picasso", outputs=[style_selected])
    s4.click(fn=lambda: "warhol", outputs=[style_selected])
    s5.click(fn=lambda: "hokusai", outputs=[style_selected])
    s6.click(fn=lambda: "starry_night", outputs=[style_selected])
    s7.click(fn=lambda: "oil_painting", outputs=[style_selected])
    s8.click(fn=lambda: "watercolor", outputs=[style_selected])
    s9.click(fn=lambda: "sketch", outputs=[style_selected])
    s10.click(fn=lambda: "comic", outputs=[style_selected])
    
    def process(image, style, inten):
        if image is None:
            return None, None, None, "❌ Carregue uma imagem"
        
        valid_img, msg = validate_image(image)
        if valid_img is None:
            return None, None, None, f"❌ {msg}"
        
        if max(valid_img.size) > MAX_IMAGE_SIZE:
            valid_img = resize_image(valid_img, MAX_IMAGE_SIZE)
        
        styled = apply_style(valid_img, style, inten)
        comp = create_comparison(valid_img, styled, style)
        
        st = f"✅ Estilo {STYLES.get(style, style)} aplicado! Intensidade: {int(inten*100)}%"
        return valid_img, styled, comp, st
    
    apply_btn.click(
        fn=process,
        inputs=[image_input, style_selected, intensity],
        outputs=[orig_out, styled_out, comp_out, status]
    )
    
    dl_styled.click(fn=save_image, inputs=[styled_out], outputs=[dl_file])
    dl_comp.click(fn=save_image, inputs=[comp_out], outputs=[dl_file])
    
    def on_upload(image):
        if image is None:
            return None, None, None, "**Status:** Aguardando..."
        valid, msg = validate_image(image)
        if valid is None:
            return None, None, None, f"❌ {msg}"
        return None, None, None, f"✅ Imagem carregada! Escolha um estilo."
    
    image_input.change(fn=on_upload, inputs=[image_input], 
                       outputs=[orig_out, styled_out, comp_out, status])

if __name__ == "__main__":
    log_message("✅ Iniciando...")
    demo.launch()