import gradio as gr
import random
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
from typing import Dict, List

class PromptGenerator:
    def __init__(self):
        # Carrega o modelo Mistral em português do Hugging Face
        print("Carregando modelo... (pode demorar alguns minutos na primeira vez)")
        self.model_name = "beomi/llama-2-ko-7b"  # Modelo base que funciona bem com português
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
        self.model = AutoModelForCausalLM.from_pretrained(
            self.model_name,
            torch_dtype=torch.float16,
            device_map="auto"
        )
        print("Modelo carregado!")

    def generate_with_model(self, prompt: str) -> str:
        """Gera texto usando o modelo local"""
        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
        
        outputs = self.model.generate(
            **inputs,
            max_new_tokens=512,
            temperature=0.7,
            top_p=0.9,
            do_sample=True,
            pad_token_id=self.tokenizer.eos_token_id
        )
        
        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)

# Templates de prompts por categoria
TEMPLATES = {
    "história": {
        "creative": "Crie uma história {gênero} sobre {tema} com {personagens} personagens. A história deve ter {comprimento} palavras e incluir {elementos}.",
        "analysis": "Analise a seguinte história considerando {aspectos}. Foque em {elementos} e forneça exemplos específicos do texto.",
        "continuation": "Continue a seguinte história mantendo o {tom} e desenvolvendo {elementos}."
    },
    "técnico": {
        "tutorial": "Crie um tutorial detalhado sobre {tema} para {público}. Inclua {elementos} e forneça {exemplos} exemplos práticos.",
        "documentation": "Documente {sistema} incluindo {aspectos}. Foque em {elementos} e forneça {exemplos} exemplos de uso.",
        "troubleshooting": "Crie um guia de solução de problemas para {sistema} cobrindo {problemas}."
    },
    "educacional": {
        "lesson": "Prepare uma aula sobre {tema} para {público}. Inclua {objetivos} e atividades práticas.",
        "exercise": "Crie exercícios sobre {tema} com {dificuldade} níveis de dificuldade.",
        "explanation": "Explique {conceito} para {público} usando {analogias} e exemplos práticos."
    }
}

# Inicializa o gerador de prompts
generator = PromptGenerator()

def generate_prompt(categoria: str, subcategoria: str, tema: str, elementos: str, 
                   público_alvo: str, tom: str, comprimento: int, 
                   usar_llm: bool = True) -> str:
    """Gera um prompt baseado nos parâmetros e usa o modelo para expandir"""
    
    if categoria not in TEMPLATES or subcategoria not in TEMPLATES[categoria]:
        return "Categoria ou subcategoria não encontrada"
    
    template = TEMPLATES[categoria][subcategoria]
    
    # Parâmetros base para o template
    params = {
        "tema": tema,
        "elementos": elementos,
        "público": público_alvo,
        "tom": tom,
        "comprimento": comprimento,
        "gênero": random.choice(["de aventura", "de mistério", "de fantasia", "de ficção científica"]),
        "personagens": random.randint(2, 5),
        "aspectos": "caracterização, desenvolvimento do enredo, temas principais",
        "exemplos": random.randint(3, 5),
        "sistema": tema,
        "problemas": "problemas comuns e casos específicos",
        "dificuldade": random.randint(3, 5),
        "conceito": tema,
        "analogias": "analogias cotidianas"
    }
    
    # Gera o prompt base
    base_prompt = template.format(**params)
    base_prompt += f"\n\nTom desejado: {tom}"
    
    if comprimento:
        base_prompt += f"\nComprimento aproximado: {comprimento} palavras"
    
    # Se solicitado, usa o modelo para expandir o prompt
    if usar_llm:
        instruction = f"""
        Você é um assistente especializado em criar prompts detalhados em português para LLMs.
        Expanda e melhore o seguinte prompt base, adicionando mais detalhes, exemplos e estrutura:

        {base_prompt}

        Crie uma versão mais completa e detalhada deste prompt, mantendo o objetivo original 
        mas adicionando mais contexto e especificidade.
        """
        return generator.generate_with_model(instruction)
    
    return base_prompt

def create_interface():
    """Cria a interface do Gradio"""
    
    with gr.Blocks(title="Gerador de Prompts para LLMs") as app:
        gr.Markdown("# 🤖 Gerador de Prompts Inteligente")
        gr.Markdown("Crie prompts estruturados e detalhados usando IA")
        
        with gr.Row():
            with gr.Column():
                categoria = gr.Dropdown(
                    choices=list(TEMPLATES.keys()),
                    label="Categoria",
                    value="história"
                )
                subcategoria = gr.Dropdown(
                    choices=list(TEMPLATES["história"].keys()),
                    label="Subcategoria",
                    value="creative"
                )
                tema = gr.Textbox(
                    label="Tema Principal",
                    placeholder="Ex: inteligência artificial, sustentabilidade"
                )
                
            with gr.Column():
                elementos = gr.Textbox(
                    label="Elementos Específicos",
                    placeholder="Ex: conflitos, conceitos, exemplos práticos"
                )
                público_alvo = gr.Textbox(
                    label="Público Alvo",
                    placeholder="Ex: iniciantes, profissionais, estudantes"
                )
                tom = gr.Dropdown(
                    choices=["formal", "informal", "técnico", "conversacional", "educativo"],
                    label="Tom",
                    value="formal"
                )
                comprimento = gr.Slider(
                    minimum=100,
                    maximum=2000,
                    step=100,
                    label="Comprimento (palavras)",
                    value=500
                )
                usar_llm = gr.Checkbox(
                    label="Usar IA para expandir prompt",
                    value=True
                )
        
        gerar_btn = gr.Button("Gerar Prompt")
        saida = gr.Textbox(label="Prompt Gerado", lines=10)
        
        def atualizar_subcategorias(categoria):
            return gr.Dropdown(choices=list(TEMPLATES[categoria].keys()))
            
        categoria.change(atualizar_subcategorias, inputs=[categoria], outputs=[subcategoria])
        gerar_btn.click(
            generate_prompt,
            inputs=[categoria, subcategoria, tema, elementos, público_alvo, tom, comprimento, usar_llm],
            outputs=[saida]
        )
    
    return app

# Criar e iniciar a aplicação
app = create_interface()

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