import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import logging

# Configuração básica de logging
logging.basicConfig(level=logging.INFO)

# Templates simplificados
TEMPLATES = {
    "história": [
        "Crie uma história sobre {tema}",
        "Escreva um conto sobre {tema}",
        "Desenvolva uma narrativa sobre {tema}"
    ],
    "tutorial": [
        "Crie um tutorial sobre {tema}",
        "Explique como fazer {tema}",
        "Desenvolva um guia sobre {tema}"
    ],
    "explicação": [
        "Explique {tema} de forma simples",
        "Descreva como funciona {tema}",
        "Faça uma análise de {tema}"
    ]
}

class SimplePromptGenerator:
    def __init__(self):
        self.model_name = "pierreguillou/gpt2-small-portuguese"  # Modelo menor em português
        self.load_model()
        
    def load_model(self):
        try:
            logging.info("Carregando modelo...")
            self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
            self.model = AutoModelForCausalLM.from_pretrained(
                self.model_name,
                torch_dtype=torch.float16,
                device_map="auto",
                low_cpu_mem_usage=True
            )
            logging.info("Modelo carregado!")
        except Exception as e:
            logging.error(f"Erro ao carregar modelo: {e}")
            raise

    def generate(self, prompt: str, max_length: int = 100) -> str:
        try:
            inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
            outputs = self.model.generate(
                **inputs,
                max_new_tokens=max_length,
                temperature=0.7,
                do_sample=True,
                pad_token_id=self.tokenizer.eos_token_id
            )
            return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        except Exception as e:
            logging.error(f"Erro na geração: {e}")
            return f"Erro: {str(e)}"

def create_interface():
    generator = SimplePromptGenerator()
    
    def gerar_prompt(categoria, tema, comprimento):
        if not tema:
            return "Por favor, insira um tema."
            
        template = TEMPLATES[categoria][0].format(tema=tema)
        prompt = f"Gere um texto em português:\n{template}"
        
        return generator.generate(prompt, max_length=comprimento)

    with gr.Blocks(theme=gr.themes.Soft()) as app:
        gr.Markdown(
            """
            # 🤖 Gerador de Prompts
            ### Crie textos em português usando IA
            """
        )
        
        with gr.Row():
            with gr.Column():
                categoria = gr.Dropdown(
                    choices=list(TEMPLATES.keys()),
                    label="📚 Categoria",
                    value="história"
                )
                tema = gr.Textbox(
                    label="💡 Tema",
                    placeholder="Digite o tema aqui..."
                )
                comprimento = gr.Slider(
                    minimum=50,
                    maximum=200,
                    value=100,
                    step=10,
                    label="📏 Comprimento (palavras)"
                )
                gerar_btn = gr.Button("🚀 Gerar", variant="primary")
        
        saida = gr.Textbox(
            label="📝 Texto Gerado",
            lines=10
        )
        
        gerar_btn.click(
            gerar_prompt,
            inputs=[categoria, tema, comprimento],
            outputs=saida
        )
    
    return app

if __name__ == "__main__":
    app = create_interface()
    app.queue()
    app.launch(server_name="0.0.0.0", server_port=7860)