import gradio as gr
import random
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
from typing import Dict, List, Tuple, Optional
from dataclasses import dataclass
import json
import os
from datetime import datetime
import logging

# Configuração de logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)

@dataclass
class PromptConfig:
    """Configuração para geração de prompts"""
    temperatura: float = 0.7
    top_p: float = 0.9
    max_tokens: int = 512
    repetition_penalty: float = 1.1

# Templates corrigidos com estrutura adequada
TEMPLATES = {
    "história": {
        "creative": {
            "template": "Crie uma história {gênero} sobre {tema} com {personagens} personagens. A história deve ter {comprimento} palavras e incluir {elementos}.",
            "sugestoes": ["Narrativa épica", "História de mistério", "Conto de fantasia"]
        },
        "analysis": {
            "template": "Analise a seguinte história considerando {aspectos}. Foque em {elementos} e forneça exemplos específicos do texto.",
            "sugestoes": ["Análise de personagens", "Análise de trama", "Análise temática"]
        },
        "continuation": {
            "template": "Continue a seguinte história mantendo o {tom} e desenvolvendo {elementos}.",
            "sugestoes": ["Desenvolvimento de personagem", "Resolução de conflito", "Expansão do universo"]
        }
    },
    "técnico": {
        "tutorial": {
            "template": "Crie um tutorial detalhado sobre {tema} para {público}. Inclua {elementos} e forneça {exemplos} exemplos práticos.",
            "sugestoes": ["Tutorial passo a passo", "Guia de início rápido", "Manual detalhado"]
        },
        "documentation": {
            "template": "Documente {sistema} incluindo {aspectos}. Foque em {elementos} e forneça {exemplos} exemplos de uso.",
            "sugestoes": ["Documentação técnica", "Manual do usuário", "Especificação de API"]
        },
        "troubleshooting": {
            "template": "Crie um guia de solução de problemas para {sistema} cobrindo {problemas}.",
            "sugestoes": ["Resolução de erros comuns", "Guia de manutenção", "FAQ técnico"]
        }
    },
    "educacional": {
        "lesson": {
            "template": "Prepare uma aula sobre {tema} para {público}. Inclua {objetivos} e atividades práticas.",
            "sugestoes": ["Plano de aula", "Roteiro de workshop", "Material didático"]
        },
        "exercise": {
            "template": "Crie exercícios sobre {tema} com {dificuldade} níveis de dificuldade.",
            "sugestoes": ["Exercícios práticos", "Desafios progressivos", "Questões de revisão"]
        },
        "explanation": {
            "template": "Explique {conceito} para {público} usando {analogias} e exemplos práticos.",
            "sugestoes": ["Explicação detalhada", "Guia simplificado", "Tutorial conceitual"]
        }
    }
}

# Mudança principal na classe do gerador para usar um modelo mais leve
class EnhancedPromptGenerator:
    def __init__(self, model_name: str = "neuralmind/bert-base-portuguese-cased"):
        self.model_name = model_name
        self.load_model()
        
    def load_model(self):
        """Carrega o modelo com otimizações para Spaces"""
        try:
            logging.info("Carregando modelo... (pode demorar alguns minutos na primeira vez)")
            self.tokenizer = AutoTokenizer.from_pretrained(
                self.model_name,
                use_fast=True
            )
            self.model = AutoModelForCausalLM.from_pretrained(
                "microsoft/phi-2",  # Modelo mais leve e eficiente
                torch_dtype=torch.float16,
                device_map="auto",
                low_cpu_mem_usage=True
            )
            logging.info("Modelo carregado com sucesso!")
        except Exception as e:
            logging.error(f"Erro ao carregar o modelo: {e}")
            raise

    def generate_with_model(self, prompt: str, config: PromptConfig = PromptConfig()) -> str:
        """Gera texto usando o modelo com configurações otimizadas"""
        try:
            # Adiciona contexto em português
            enhanced_prompt = f"""
            Por favor, gere um texto em português de acordo com as instruções:
            {prompt}
            """
            inputs = self.tokenizer(enhanced_prompt, return_tensors="pt").to(self.model.device)
            
            outputs = self.model.generate(
                **inputs,
                max_new_tokens=config.max_tokens,
                temperature=config.temperatura,
                top_p=config.top_p,
                repetition_penalty=config.repetition_penalty,
                do_sample=True,
                pad_token_id=self.tokenizer.eos_token_id,
                num_beams=4,  # Melhor qualidade para português
                early_stopping=True
            )
            
            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 na geração: {str(e)}"
            
generator = EnhancedPromptGenerator()

def get_suggestions(categoria: str, subcategoria: str) -> List[str]:
    """Retorna sugestões baseadas na categoria e subcategoria"""
    try:
        if categoria in TEMPLATES and subcategoria in TEMPLATES[categoria]:
            return TEMPLATES[categoria][subcategoria].get("sugestoes", [])
        return []
    except Exception as e:
        logging.error(f"Erro ao obter sugestões: {e}")
        return []

def generate_prompt(
    categoria: str,
    subcategoria: str,
    tema: str,
    elementos: str,
    público_alvo: str,
    tom: str,
    comprimento: int,
    temperatura: float,
    usar_llm: bool = True
) -> Tuple[str, dict]:
    """Função aprimorada de geração de prompts"""
    try:
        if not tema or not elementos:
            return "Por favor, preencha o tema e os elementos.", {"erro": "Campos obrigatórios não preenchidos"}

        if categoria not in TEMPLATES or subcategoria not in TEMPLATES[categoria]:
            return "Categoria ou subcategoria não encontrada", {"erro": "Categoria inválida"}
        
        template = TEMPLATES[categoria][subcategoria]["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"
        }
        
        base_prompt = template.format(**params)
        base_prompt += f"\n\nTom desejado: {tom}"
        
        if comprimento:
            base_prompt += f"\nComprimento aproximado: {comprimento} palavras"
        
        if usar_llm:
            config = PromptConfig(temperatura=temperatura)
            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}

            Considere:
            - Público-alvo: {público_alvo}
            - Elementos específicos: {elementos}
            - Tom desejado: {tom}

            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, config), {"status": "success"}
        
        return base_prompt, {"status": "success"}
    except Exception as e:
        logging.error(f"Erro na geração do prompt: {e}")
        return f"Erro na geração do prompt: {str(e)}", {"erro": str(e)}

def create_interface():
    """Interface aprimorada com mais recursos"""
    with gr.Blocks(
        title="Gerador de Prompts Inteligente",
        theme=gr.themes.Soft(
            primary_hue="indigo",
            secondary_hue="blue",
        ),
    ) as app:
        gr.Markdown(
            """
            # 🤖 Gerador de Prompts Inteligente v2.0
            ### Crie prompts estruturados e detalhados usando IA avançada
            """
        )
        
        with gr.Tabs():
            with gr.TabItem("✨ Gerador de Prompts"):
                with gr.Group():
                    with gr.Row():
                        with gr.Column(scale=1):
                            categoria = gr.Dropdown(
                                choices=list(TEMPLATES.keys()),
                                label="📚 Categoria",
                                value="história",
                                container=False,
                            )
                            subcategoria = gr.Dropdown(
                                choices=list(TEMPLATES["história"].keys()),
                                label="🔍 Subcategoria",
                                value="creative",
                                container=False,
                            )
                            sugestoes = gr.Dropdown(
                                choices=get_suggestions("história", "creative"),
                                label="💫 Sugestões",
                                container=False,
                            )
                        
                        with gr.Column(scale=2):
                            tema = gr.Textbox(
                                label="💡 Tema Principal",
                                placeholder="Ex: inteligência artificial, sustentabilidade",
                                container=False,
                            )
                            elementos = gr.Textbox(
                                label="🔮 Elementos Específicos",
                                placeholder="Ex: conflitos, conceitos, exemplos práticos",
                                container=False,
                            )
                    
                    with gr.Row():
                        with gr.Column(scale=1):
                            público_alvo = gr.Textbox(
                                label="👥 Público Alvo",
                                placeholder="Ex: iniciantes, profissionais, estudantes",
                                container=False,
                            )
                            tom = gr.Dropdown(
                                choices=["formal", "informal", "técnico", "conversacional", "educativo"],
                                label="🎭 Tom",
                                value="formal",
                                container=False,
                            )
                        
                        with gr.Column(scale=2):
                            comprimento = gr.Slider(
                                minimum=100,
                                maximum=2000,
                                step=100,
                                label="📏 Comprimento (palavras)",
                                value=500,
                                container=False,
                            )
                            temperatura = gr.Slider(
                                minimum=0.1,
                                maximum=1.0,
                                step=0.1,
                                label="🌡️ Temperatura",
                                value=0.7,
                                container=False,
                            )
                    
                    usar_llm = gr.Checkbox(
                        label="🤖 Usar IA para expandir prompt",
                        value=True,
                        container=False,
                    )
                
                with gr.Row():
                    gerar_btn = gr.Button(
                        "🚀 Gerar Prompt",
                        variant="primary",
                        scale=1,
                    )
                
                with gr.Row():
                    saida = gr.TextArea(
                        label="📝 Prompt Gerado",
                        lines=10,
                        container=False,
                    )
                    status = gr.JSON(
                        label="📊 Status",
                        container=False,
                    )
        
        def atualizar_interface(categoria, subcategoria):
            try:
                sugestoes_list = get_suggestions(categoria, subcategoria)
                return [
                    gr.Dropdown(choices=list(TEMPLATES[categoria].keys())),
                    gr.Dropdown(choices=sugestoes_list)
                ]
            except Exception as e:
                logging.error(f"Erro ao atualizar interface: {e}")
                return [gr.Dropdown(), gr.Dropdown()]
            
        categoria.change(
            atualizar_interface,
            inputs=[categoria, subcategoria],
            outputs=[subcategoria, sugestoes]
        )
        
        gerar_btn.click(
            generate_prompt,
            inputs=[
                categoria, subcategoria, tema, elementos,
                público_alvo, tom, comprimento, temperatura, usar_llm
            ],
            outputs=[saida, status]
        )
    
    return app

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