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

@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

class EnhancedPromptGenerator:
    def __init__(self, model_name: str = "beomi/llama-2-ko-7b"):
        self.model_name = model_name
        self.history_file = "prompt_history.json"
        self.load_model()
        self.load_history()
        
    def load_model(self):
        """Carrega o modelo com tratamento de erro"""
        try:
            print("Carregando modelo... (pode demorar alguns minutos na primeira vez)")
            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
            )
            print("Modelo carregado com sucesso!")
        except Exception as e:
            print(f"Erro ao carregar o modelo: {e}")
            raise

    def load_history(self):
        """Carrega histórico de prompts gerados"""
        self.history = []
        if os.path.exists(self.history_file):
            try:
                with open(self.history_file, 'r', encoding='utf-8') as f:
                    self.history = json.load(f)
            except:
                print("Erro ao carregar histórico. Criando novo arquivo.")

    def save_history(self):
        """Salva histórico de prompts"""
        try:
            with open(self.history_file, 'w', encoding='utf-8') as f:
                json.dump(self.history, f, ensure_ascii=False, indent=2)
        except Exception as e:
            print(f"Erro ao salvar histórico: {e}")

    def analyze_prompt_quality(self, prompt: str) -> Dict[str, float]:
        """Analisa a qualidade do prompt gerado"""
        analysis = {
            "complexidade": len(prompt.split()) / 100,  # Normalizado para 0-1
            "especificidade": len(set(prompt.split())) / len(prompt.split()),
            "clareza": 1 - (prompt.count(',') + prompt.count(';')) / len(prompt.split())
        }
        return analysis

    def generate_with_model(self, prompt: str, config: PromptConfig = PromptConfig()) -> Tuple[str, Dict[str, float]]:
        """Gera texto usando o modelo com configurações avançadas"""
        try:
            inputs = self.tokenizer(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
            )
            
            generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
            quality_metrics = self.analyze_prompt_quality(generated_text)
            
            # Salva no histórico
            self.history.append({
                "timestamp": datetime.now().isoformat(),
                "prompt": prompt,
                "generated": generated_text,
                "metrics": quality_metrics
            })
            self.save_history()
            
            return generated_text, quality_metrics
            
        except Exception as e:
            print(f"Erro na geração: {e}")
            return str(e), {"erro": 1.0}

# Templates expandidos com mais opções e detalhes
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}.",
            "sugestões": ["aventura épica", "drama pessoal", "mistério sobrenatural"]
        },
        "analysis": {
            "template": "Analise a seguinte história considerando {aspectos}. Foque em {elementos} e forneça exemplos específicos do texto.",
            "sugestões": ["análise profunda", "crítica construtiva", "interpretação temática"]
        },
        "continuation": {
            "template": "Continue a seguinte história mantendo o {tom} e desenvolvendo {elementos}.",
            "sugestões": ["expansão do universo", "desenvolvimento de personagem", "resolução de conflito"]
        }
    },
    "técnico": {
        "tutorial": {
            "template": "Crie um tutorial detalhado sobre {tema} para {público}. Inclua {elementos} e forneça {exemplos} exemplos práticos.",
            "sugestões": ["passo a passo", "guia completo", "referência rápida"]
        },
        "documentation": {
            "template": "Documente {sistema} incluindo {aspectos}. Foque em {elementos} e forneça {exemplos} exemplos de uso.",
            "sugestões": ["documentação técnica", "manual do usuário", "guia de referência"]
        },
        "troubleshooting": {
            "template": "Crie um guia de solução de problemas para {sistema} cobrindo {problemas}.",
            "sugestões": ["resolução de problemas", "diagnóstico", "manutenção"]
        }
    },
    "educacional": {
        "lesson": {
            "template": "Prepare uma aula sobre {tema} para {público}. Inclua {objetivos} e atividades práticas.",
            "sugestões": ["aula interativa", "workshop prático", "seminário"]
        },
        "exercise": {
            "template": "Crie exercícios sobre {tema} com {dificuldade} níveis de dificuldade.",
            "sugestões": ["exercícios práticos", "desafios", "problemas"]
        },
        "explanation": {
            "template": "Explique {conceito} para {público} usando {analogias} e exemplos práticos.",
            "sugestões": ["explicação detalhada", "conceitos básicos", "aprofundamento"]
        }
    }
}

generator = EnhancedPromptGenerator()

def get_suggestions(categoria: str, subcategoria: str) -> List[str]:
    """Retorna sugestões baseadas na categoria e subcategoria"""
    return TEMPLATES[categoria][subcategoria].get("sugestões", [])

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[str, float]]:
    """Função aprimorada de geração de prompts"""
    if not tema or not elementos:
        return "Por favor, preencha o tema e os elementos.", {"erro": 1.0}

    if categoria not in TEMPLATES or subcategoria not in TEMPLATES[categoria]:
        return "Categoria ou subcategoria não encontrada", {"erro": 1.0}
    
    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)
    
    return base_prompt, generator.analyze_prompt_quality(base_prompt)

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=[],
                                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):
                            with gr.Row():
                                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,
                    )
                    metricas = gr.JSON(
                        label="📊 Métricas de Qualidade",
                        container=False,
                    )
            
            with gr.TabItem("📈 Histórico"):
                historico = gr.DataFrame(
                    headers=["Data", "Prompt Original", "Texto Gerado", "Métricas"],
                    label="Histórico de Prompts Gerados"
                )
            
            with gr.TabItem("ℹ️ Sobre"):
                gr.Markdown(
                    """
                    ### Sobre o Gerador de Prompts v2.0
                    
                    Esta é uma versão aprimorada do gerador de prompts que inclui:
                    
                    - 📊 Análise de qualidade dos prompts gerados
                    - 📈 Histórico de gerações
                    - 🎯 Sugestões contextuais
                    - 🌡️ Controle de temperatura
                    - 🔄 Sistema de feedback
                    
                    O gerador usa um modelo de linguagem avançado para criar prompts
                    detalhados e específicos para suas necessidades.
                    """
                )
        
        def atualizar_interface(categoria, subcategoria):
            sugestoes_list = get_suggestions(categoria, subcategoria)
            return [
                gr.Dropdown(choices=list(TEMPLATES[categoria].keys())),
                gr.Dropdown(choices=sugestoes_list)
            ]
        
        def atualizar_historico():
            if generator.history:
                dados = [(
                    h["timestamp"],
                    h["prompt"],
                    h["generated"],
                    str(h["metrics"])
                ) for h in generator.history]
                return gr.DataFrame(value=dados)
            return gr.DataFrame()
            
        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, metricas]
        ).then(
            atualizar_historico,outputs=[historico]
        )
        
        # Adiciona feedback visual de carregamento
        gerar_btn.click(
            lambda: gr.update(interactive=False),
            outputs=[gerar_btn]
        ).then(
            lambda: gr.update(interactive=True),
            outputs=[gerar_btn]
        )
        
        # Função para exportar histórico
        def exportar_historico():
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"historico_prompts_{timestamp}.json"
            try:
                with open(filename, 'w', encoding='utf-8') as f:
                    json.dump(generator.history, f, ensure_ascii=False, indent=2)
                return f"Histórico exportado com sucesso para {filename}"
            except Exception as e:
                return f"Erro ao exportar histórico: {e}"
        
        # Adiciona botão de exportação na aba de histórico
        with gr.TabItem("📈 Histórico"):
            with gr.Row():
                exportar_btn = gr.Button("📥 Exportar Histórico")
                status_exportacao = gr.Textbox(label="Status da Exportação")
            
            exportar_btn.click(
                exportar_historico,
                outputs=[status_exportacao]
            )
        
        # Adiciona sistema de templates favoritos
        class TemplateManager:
            def __init__(self, filename="templates_favoritos.json"):
                self.filename = filename
                self.templates = self.carregar_templates()
            
            def carregar_templates(self):
                if os.path.exists(self.filename):
                    try:
                        with open(self.filename, 'r', encoding='utf-8') as f:
                            return json.load(f)
                    except:
                        return {}
                return {}
            
            def salvar_templates(self):
                with open(self.filename, 'w', encoding='utf-8') as f:
                    json.dump(self.templates, f, ensure_ascii=False, indent=2)
            
            def adicionar_template(self, nome: str, template: dict):
                self.templates[nome] = template
                self.salvar_templates()
            
            def remover_template(self, nome: str):
                if nome in self.templates:
                    del self.templates[nome]
                    self.salvar_templates()
        
        template_manager = TemplateManager()
        
        # Adiciona aba de templates favoritos
        with gr.TabItem("⭐ Templates Favoritos"):
            with gr.Row():
                with gr.Column():
                    nome_template = gr.Textbox(label="Nome do Template")
                    salvar_template_btn = gr.Button("💾 Salvar Template Atual")
                    templates_salvos = gr.Dropdown(
                        choices=list(template_manager.templates.keys()),
                        label="Templates Salvos"
                    )
                    carregar_template_btn = gr.Button("📂 Carregar Template")
                    remover_template_btn = gr.Button("🗑️ Remover Template")
            
            def salvar_template_atual(nome, categoria, subcategoria, tema, elementos, público_alvo, tom, comprimento):
                if not nome:
                    return "Por favor, forneça um nome para o template"
                
                template = {
                    "categoria": categoria,
                    "subcategoria": subcategoria,
                    "tema": tema,
                    "elementos": elementos,
                    "público_alvo": público_alvo,
                    "tom": tom,
                    "comprimento": comprimento
                }
                
                template_manager.adicionar_template(nome, template)
                return f"Template '{nome}' salvo com sucesso!"
            
            def carregar_template(nome):
                if nome not in template_manager.templates:
                    return [gr.update() for _ in range(7)]
                
                template = template_manager.templates[nome]
                return [
                    gr.update(value=template["categoria"]),
                    gr.update(value=template["subcategoria"]),
                    gr.update(value=template["tema"]),
                    gr.update(value=template["elementos"]),
                    gr.update(value=template["público_alvo"]),
                    gr.update(value=template["tom"]),
                    gr.update(value=template["comprimento"])
                ]
            
            def remover_template(nome):
                if nome:
                    template_manager.remover_template(nome)
                    return [
                        gr.update(choices=list(template_manager.templates.keys())),
                        f"Template '{nome}' removido com sucesso!"
                    ]
                return [gr.update(), "Selecione um template para remover"]
            
            salvar_template_btn.click(
                salvar_template_atual,
                inputs=[
                    nome_template, categoria, subcategoria,
                    tema, elementos, público_alvo, tom, comprimento
                ],
                outputs=[gr.Textbox(label="Status")]
            )
            
            carregar_template_btn.click(
                carregar_template,
                inputs=[templates_salvos],
                outputs=[categoria, subcategoria, tema, elementos, público_alvo, tom, comprimento]
            )
            
            remover_template_btn.click(
                remover_template,
                inputs=[templates_salvos],
                outputs=[templates_salvos, gr.Textbox(label="Status")]
            )
        
        # Adiciona sistema de feedback
        with gr.TabItem("📝 Feedback"):
            with gr.Row():
                feedback_text = gr.TextArea(
                    label="Seu feedback sobre o prompt gerado",
                    placeholder="Digite seu feedback aqui..."
                )
                rating = gr.Slider(
                    minimum=1,
                    maximum=5,
                    step=1,
                    label="Avaliação (1-5 estrelas)"
                )
            
            def salvar_feedback(texto, nota):
                if not texto:
                    return "Por favor, forneça um feedback"
                
                feedback = {
                    "texto": texto,
                    "avaliacao": nota,
                    "timestamp": datetime.now().isoformat()
                }
                
                try:
                    with open("feedbacks.json", "a+") as f:
                        json.dump(feedback, f)
                        f.write("\n")
                    return "Feedback salvo com sucesso! Obrigado pela sua avaliação."
                except Exception as e:
                    return f"Erro ao salvar feedback: {e}"
            
            feedback_btn = gr.Button("✉️ Enviar Feedback")
            feedback_status = gr.Textbox(label="Status do Feedback")
            
            feedback_btn.click(
                salvar_feedback,
                inputs=[feedback_text, rating],
                outputs=[feedback_status]
            )
    
    return app

# No final do código, substitua o bloco de launch por:

if __name__ == "__main__":
    app = create_interface()
    app.queue()  # Habilita queue de forma compatível
    app.launch(
        share=True,
        server_name="0.0.0.0",
        server_port=7860
    )