import gradio as gr
from PIL import Image, ImageDraw, ImageFont
import torch
from diffusers import StableDiffusionImg2ImgPipeline
import tempfile
import os

# Carregar modelo de transformação de imagem
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
).to("cuda" if torch.cuda.is_available() else "cpu")

# Prompts por estilo
PROMPTS = {
    "Infantil": "a colorful, playful, cartoon-style kids room, minimal, bright colors, 3D illustration",
    "Adulto": "a modern, elegant, minimalist living room with warm tones and natural textures, stylish interior design"
}

def transformar_imagem(imagem_input, estilo, texto_usuario):
    if imagem_input is None:
        return None
    
    # Converter para RGB
    imagem_input = imagem_input.convert("RGB")

    prompt = PROMPTS.get(estilo, PROMPTS["Adulto"])

    imagem_saida = pipe(prompt=prompt, image=imagem_input, strength=0.75, guidance_scale=7.5).images[0]

    # Adicionar texto com estilo
    draw = ImageDraw.Draw(imagem_saida)
    font_path = "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf"  # fonte padrão Linux, pode mudar conforme o sistema
    font_size = 40
    font = ImageFont.truetype(font_path, font_size)

    # Centralizar o texto
    w, h = imagem_saida.size
    text_w, text_h = draw.textsize(texto_usuario, font=font)
    text_x = (w - text_w) / 2
    text_y = h - text_h - 30  # posição inferior com margem

    draw.text((text_x, text_y), texto_usuario, fill="white", font=font, stroke_width=2, stroke_fill="black")

    return imagem_saida

# Interface
with gr.Blocks() as demo:
    gr.Markdown("# 🖼️ Estilizador de Ambientes com IA")
    gr.Markdown("Transforme uma foto em um ambiente **infantil ou adulto** estilizado e com mensagem personalizada.")

    with gr.Row():
        imagem_input = gr.Image(label="1. Envie uma imagem do ambiente", type="pil")
        estilo = gr.Radio(choices=["Infantil", "Adulto"], label="2. Escolha o estilo do ambiente")

    texto_usuario = gr.Textbox(label="3. Digite o texto que deseja sobrepor na imagem (ex: frase, dedicatória)")
    gerar_botao = gr.Button("4. Gerar Imagem com Estilo + Texto")

    imagem_saida = gr.Image(label="Imagem Final", type="pil")
    download = gr.File(label="Download da Imagem", interactive=True)

    def gerar(imagem, estilo_escolhido, texto):
        imagem_final = transformar_imagem(imagem, estilo_escolhido, texto)

        # Salvar para download
        with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmpfile:
            imagem_final.save(tmpfile.name)
            return imagem_final, tmpfile.name

    gerar_botao.click(fn=gerar, inputs=[imagem_input, estilo, texto_usuario], outputs=[imagem_saida, download])

# Lançar app
demo.launch()