"""
🎨 Image Colorizer
Versão otimizada para Hugging Face Spaces
"""

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

print("🚀 Iniciando Image Colorizer...")

MAX_IMAGE_SIZE = 1024

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

def validate_image(image):
    """Valida e prepara imagem"""
    try:
        if image is None:
            return None, "❌ Nenhuma imagem fornecida"
        
        # Converter para PIL Image se necessário
        if isinstance(image, np.ndarray):
            img = Image.fromarray(image.astype('uint8'))
        else:
            img = image
        
        # Verificar tamanho
        if max(img.size) > 4000:
            return None, "❌ Imagem muito grande (>4000px)"
        
        if min(img.size) < 32:
            return None, "❌ Imagem muito pequena (<32px)"
        
        log_message(f"✅ Imagem válida: {img.size}px, {img.mode}")
        return img, "ok"
        
    except Exception as e:
        log_message(f"❌ Erro na validação: {str(e)}")
        return None, str(e)

def resize_image(image, max_size):
    """Redimensiona mantendo aspect ratio"""
    if max(image.size) <= max_size:
        return image
    
    ratio = max_size / max(image.size)
    new_width = int(image.width * ratio)
    new_height = int(image.height * ratio)
    
    return image.resize((new_width, new_height), Image.Resampling.LANCZOS)

def apply_colorization(image, style="realistic", intensity=0.8):
    """Aplica colorização à imagem"""
    try:
        log_message(f"Aplicando colorização - Estilo: {style}, Intensidade: {intensity}")
        
        # Converter para RGB
        if image.mode != 'RGB':
            rgb_img = image.convert('RGB')
        else:
            rgb_img = image.copy()
        
        gray_img = rgb_img.convert('L')
        
        # Aplicar efeitos baseados no estilo
        if style == "realistic":
            result = rgb_img.copy()
            enhancer = ImageEnhance.Color(result)
            result = enhancer.enhance(1.0 + (intensity * 0.5))
            
            r, g, b = result.split()
            r = r.point(lambda x: min(255, int(x * (1.0 + intensity * 0.1))))
            b = b.point(lambda x: max(0, int(x * (1.0 - intensity * 0.05))))
            result = Image.merge('RGB', (r, g, b))
            
        elif style == "vibrant":
            result = rgb_img.copy()
            enhancer = ImageEnhance.Color(result)
            result = enhancer.enhance(1.0 + (intensity * 1.0))
            enhancer = ImageEnhance.Contrast(result)
            result = enhancer.enhance(1.0 + (intensity * 0.3))
            result = result.filter(ImageFilter.UnsharpMask(radius=1, percent=50, threshold=0))
            
        elif style == "vintage":
            result = rgb_img.copy()
            r, g, b = result.split()
            r = r.point(lambda x: min(255, int(x * 1.1)))
            g = g.point(lambda x: int(x * 0.9))
            b = b.point(lambda x: int(x * 0.8))
            result = Image.merge('RGB', (r, g, b))
            enhancer = ImageEnhance.Color(result)
            result = enhancer.enhance(0.7 + (intensity * 0.3))
            
        elif style == "cinematic":
            result = rgb_img.copy()
            r, g, b = result.split()
            r = r.point(lambda x: int(x * 0.9))
            g = g.point(lambda x: int(x * 1.0))
            b = b.point(lambda x: min(255, int(x * 1.1)))
            result = Image.merge('RGB', (r, g, b))
            enhancer = ImageEnhance.Contrast(result)
            result = enhancer.enhance(1.0 + (intensity * 0.4))
            
        else:  # balanced
            result = rgb_img.copy()
            enhancer = ImageEnhance.Color(result)
            result = enhancer.enhance(1.0 + (intensity * 0.6))
            enhancer = ImageEnhance.Contrast(result)
            result = enhancer.enhance(1.0 + (intensity * 0.2))
        
        # Misturar se imagem era grayscale
        if image.mode in ['L', 'LA', 'P']:
            gray_array = np.array(gray_img).astype(np.float32) / 255.0
            original_array = np.array(rgb_img).astype(np.float32)
            colorized_array = np.array(result).astype(np.float32)
            
            mixed = (1 - intensity) * original_array + intensity * colorized_array
            
            for i in range(3):
                mixed[:,:,i] = mixed[:,:,i] * (0.7 + 0.3 * gray_array)
            
            result = Image.fromarray(mixed.astype(np.uint8))
        
        log_message("✅ Colorização concluída")
        return result
        
    except Exception as e:
        log_message(f"❌ Erro na colorização: {str(e)}")
        return image

def create_comparison(original, colorized):
    """Cria comparação lado a lado"""
    try:
        if original is None or colorized is None:
            return None
        
        target_height = 400
        target_width_orig = int(original.width * (target_height / original.height))
        target_width_color = int(colorized.width * (target_height / colorized.height))
        
        original_resized = original.resize((target_width_orig, target_height), Image.Resampling.LANCZOS)
        colorized_resized = colorized.resize((target_width_color, target_height), Image.Resampling.LANCZOS)
        
        total_width = original_resized.width + colorized_resized.width + 20
        total_height = target_height + 60
        
        comparison = Image.new('RGB', (total_width, total_height), color=(240, 240, 240))
        draw = ImageDraw.Draw(comparison)
        
        draw.text((10, 10), "ORIGINAL", fill=(100, 100, 100))
        draw.text((original_resized.width + 30, 10), "COLORIZED", fill=(0, 150, 0))
        
        comparison.paste(original_resized, (0, 40))
        comparison.paste(colorized_resized, (original_resized.width + 20, 40))
        
        log_message("✅ Comparação criada")
        return comparison
        
    except Exception as e:
        log_message(f"❌ Erro na comparação: {str(e)}")
        return None

def save_image(image, prefix="colorized"):
    """Salva imagem para download"""
    try:
        if image is None:
            return None
        
        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".png", prefix=f"{prefix}_")
        image.save(temp_file.name, "PNG", optimize=True)
        
        log_message(f"Imagem salva: {temp_file.name}")
        return temp_file.name
        
    except Exception as e:
        log_message(f"❌ Erro ao salvar: {str(e)}")
        return None

# Interface Gradio
with gr.Blocks(title="🎨 Image Colorizer", theme=gr.themes.Soft()) as demo:
    
    gr.HTML("""
    <div style="text-align: center; padding: 20px; background: linear-gradient(135deg, #ff7e5f 0%, #feb47b 100%); border-radius: 10px; color: white; margin-bottom: 20px;">
        <h1 style="margin: 0;">🎨 Image Colorizer</h1>
        <p style="margin: 5px 0 0 0;">Colorize fotos preto e branco automaticamente</p>
    </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 📤 Upload da Foto")
            image_input = gr.Image(type="pil", label="Clique ou arraste uma imagem")
            
            gr.Markdown("### ⚙️ Configurações")
            gr.Markdown("**Estilo de colorização:**")
            
            with gr.Row():
                style_realistic = gr.Button("Realista", size="sm")
                style_vibrant = gr.Button("Vibrante", size="sm")
                style_vintage = gr.Button("Vintage", size="sm")
            
            style_selected = gr.Textbox(value="realistic", visible=False)
            
            gr.Markdown("**Intensidade das cores:**")
            intensity_slider = gr.Slider(
                minimum=0.1, maximum=1.0, value=0.7, step=0.1,
                label="", info="0.1 = Suave | 1.0 = Intenso"
            )
            
            colorize_btn = gr.Button("🎨 Colorizar Imagem", variant="primary", size="lg")
        
        with gr.Column(scale=1):
            gr.Markdown("### 📊 Resultado")
            status_output = gr.Markdown("**Status:** Aguardando imagem...")
            
            with gr.Tabs():
                with gr.TabItem("🔄 Comparação"):
                    comparison_output = gr.Image(type="pil", label="Antes e Depois")
                
                with gr.TabItem("📷 Original"):
                    original_output = gr.Image(type="pil", label="Original")
                
                with gr.TabItem("🌈 Colorizada"):
                    colorized_output = gr.Image(type="pil", label="Colorizada")
            
            gr.Markdown("### 💾 Download")
            with gr.Row():
                download_colorized = gr.Button("📥 Baixar Colorizada")
                download_comparison = gr.Button("📊 Baixar Comparação")
            
            download_file = gr.File(label="Arquivo para download", interactive=False)
    
    with gr.Accordion("📖 Como usar", open=False):
        gr.Markdown("""
        ## Instruções:
        1. **Carregue** uma foto em preto e branco
        2. **Escolha** o estilo de colorização
        3. **Ajuste** a intensidade das cores
        4. **Clique** em "Colorizar Imagem"
        5. **Compare** os resultados e baixe
        
        ## Estilos:
        - **Realista**: Tons naturais para retratos
        - **Vibrante**: Cores vivas para paisagens
        - **Vintage**: Estilo antigo e nostálgico
        """)
    
    gr.HTML("""
    <div style="text-align: center; margin-top: 20px; padding: 15px; background: #f8f9fa; border-radius: 8px;">
        <p style="margin: 0; color: #666;">🎨 Image Colorizer - Photoshop AI Ecosystem</p>
    </div>
    """)
    
    # Event handlers
    style_realistic.click(fn=lambda: "realistic", outputs=[style_selected])
    style_vibrant.click(fn=lambda: "vibrant", outputs=[style_selected])
    style_vintage.click(fn=lambda: "vintage", outputs=[style_selected])
    
    def process_colorization(image, style, intensity):
        if image is None:
            return None, None, None, "❌ Por favor, carregue uma imagem primeiro"
        
        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)
        
        colorized = apply_colorization(valid_img, style, intensity)
        comparison = create_comparison(valid_img, colorized)
        
        status = f"""
        ✅ **Colorização aplicada com sucesso!**
        
        **Detalhes:**
        • Estilo: {style.title()}
        • Intensidade: {intensity*100:.0f}%
        • Tamanho: {valid_img.size[0]}×{valid_img.size[1]}px
        """
        
        return valid_img, colorized, comparison, status
    
    colorize_btn.click(
        fn=process_colorization,
        inputs=[image_input, style_selected, intensity_slider],
        outputs=[original_output, colorized_output, comparison_output, status_output]
    )
    
    download_colorized.click(fn=save_image, inputs=[colorized_output], outputs=[download_file])
    download_comparison.click(fn=save_image, inputs=[comparison_output], outputs=[download_file])
    
    def clear_on_upload(image):
        if image is None:
            return None, None, None, "**Status:** Aguardando imagem..."
        
        valid_img, msg = validate_image(image)
        if valid_img is None:
            return None, None, None, f"❌ {msg}"
        
        status = f"✅ **Imagem carregada!** ({valid_img.size[0]}×{valid_img.size[1]}px)\n\nEscolha um estilo e clique em Colorizar."
        return None, None, None, status
    
    image_input.change(
        fn=clear_on_upload,
        inputs=[image_input],
        outputs=[original_output, colorized_output, comparison_output, status_output]
    )

if __name__ == "__main__":
    log_message("✅ Iniciando aplicação...")
    demo.launch()