import gradio as gr
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
from transformers import MarianMTModel, MarianTokenizer
from diffusers import StableDiffusionPipeline
import torch
import numpy as np
from PIL import Image
import spacy

class AIServices:
    def __init__(self):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.image_generator = None
        self.translator = None
        self.sentiment_analyzer = None
        self.summarizer = None
        self.ner_model = None
        self.classifier = None

    def load_image_generator(self):
        if self.image_generator is None:
            model_id = "CompVis/stable-diffusion-v1-4"
            self.image_generator = StableDiffusionPipeline.from_pretrained(
                model_id,
                torch_dtype=torch.float32
            ).to(self.device)
        return self.image_generator

    def generate_image(self, prompt, num_images=1):
        try:
            generator = self.load_image_generator()
            images = generator(
                prompt,
                num_images_per_prompt=num_images,
                guidance_scale=7.5
            ).images
            return images[0] if num_images == 1 else images
        except Exception as e:
            return f"Erro na geração de imagem: {str(e)}"

    def translate(self, text, src_lang, tgt_lang):
        if self.translator is None:
            model_name = f'Helsinki-NLP/opus-mt-{src_lang}-{tgt_lang}'
            self.translator = pipeline('translation', model=model_name)
        try:
            result = self.translator(text)[0]['translation_text']
            return result
        except Exception as e:
            return f"Erro na tradução: {str(e)}"

    def analyze_sentiment(self, text):
        if self.sentiment_analyzer is None:
            self.sentiment_analyzer = pipeline(
                'sentiment-analysis',
                model='nlptown/bert-base-multilingual-uncased-sentiment'
            )
        try:
            result = self.sentiment_analyzer(text)[0]
            return f"Sentimento: {result['label']}, Confiança: {result['score']:.2f}"
        except Exception as e:
            return f"Erro na análise: {str(e)}"

    def summarize_text(self, text, max_length=150, min_length=50):
        if self.summarizer is None:
            self.summarizer = pipeline('summarization', model='facebook/bart-large-cnn')
        try:
            summary = self.summarizer(text, max_length=max_length, min_length=min_length)[0]
            return summary['summary_text']
        except Exception as e:
            return f"Erro ao resumir: {str(e)}"

    def extract_entities(self, text):
        try:
            if self.ner_model is None:
                self.ner_model = spacy.load("pt_core_news_sm")
            
            doc = self.ner_model(text)
            entities = []
            
            for ent in doc.ents:
                entities.append({
                    'text': ent.text,
                    'label': ent.label_,
                    'start': ent.start_char,
                    'end': ent.end_char
                })
            
            # Formatando a saída
            result = ""
            for e in entities:
                result += f"📌 {e['text']} ({e['label']})\n"
            
            return result if result else "Nenhuma entidade encontrada."
        except Exception as e:
            return f"Erro na extração de entidades: {str(e)}"

    def classify_text(self, text):
        if self.classifier is None:
            self.classifier = pipeline(
                "zero-shot-classification",
                model="facebook/bart-large-mnli"
            )
        try:
            # Categorias predefinidas
            candidate_labels = [
                "negócios", "tecnologia", "política", 
                "entretenimento", "esportes", "ciência",
                "saúde", "educação"
            ]
            
            result = self.classifier(text, candidate_labels, multi_label=True)
            
            # Formatando a saída
            output = "🏷️ Classificação:\n"
            for label, score in zip(result['labels'], result['scores']):
                if score > 0.1:  # Mostra apenas scores relevantes
                    output += f"{label}: {score:.1%}\n"
            
            return output
        except Exception as e:
            return f"Erro na classificação: {str(e)}"

# Instância global dos serviços
services = AIServices()

# Interface Gradio
with gr.Blocks(title="Hub de Serviços de IA") as demo:
    gr.Markdown("""
    # 🤖 Hub de Serviços de IA
    
    Esta aplicação oferece diversos serviços de IA para processamento de texto e imagem.
    """)
    
    # 1. Geração de Imagem
    with gr.Tab("Geração de Imagem"):
        gr.Markdown("### 🎨 Gerador de Imagens com Stable Diffusion")
        with gr.Row():
            img_prompt = gr.Textbox(
                label="Descrição da imagem",
                placeholder="Descreva a imagem que deseja gerar...",
                lines=3
            )
            img_output = gr.Image(label="Imagem Gerada")
        with gr.Row():
            img_num = gr.Slider(
                minimum=1,
                maximum=4,
                value=1,
                step=1,
                label="Número de imagens"
            )
        img_button = gr.Button("🎨 Gerar Imagem")
        img_button.click(
            services.generate_image,
            inputs=[img_prompt, img_num],
            outputs=img_output
        )
    
    # 2. Tradução
    with gr.Tab("Tradutor"):
        gr.Markdown("### 🌐 Tradutor Multilíngue")
        with gr.Row():
            trans_input = gr.Textbox(
                label="Texto para traduzir",
                placeholder="Digite o texto aqui...",
                lines=3
            )
            trans_output = gr.Textbox(
                label="Tradução",
                lines=3
            )
        with gr.Row():
            src_lang = gr.Dropdown(
                choices=["en", "pt", "es", "fr", "de"],
                value="en",
                label="Idioma de origem"
            )
            tgt_lang = gr.Dropdown(
                choices=["pt", "en", "es", "fr", "de"],
                value="pt",
                label="Idioma de destino"
            )
        trans_button = gr.Button("🔄 Traduzir")
        trans_button.click(
            services.translate,
            inputs=[trans_input, src_lang, tgt_lang],
            outputs=trans_output
        )
    
    # 3. Análise de Sentimentos
    with gr.Tab("Análise de Sentimentos"):
        gr.Markdown("### 😊 Análise de Sentimentos")
        with gr.Row():
            sent_input = gr.Textbox(
                label="Texto para análise",
                placeholder="Digite o texto para analisar o sentimento...",
                lines=3
            )
            sent_output = gr.Textbox(
                label="Resultado da análise",
                lines=2
            )
        sent_button = gr.Button("💭 Analisar Sentimento")
        sent_button.click(
            services.analyze_sentiment,
            inputs=sent_input,
            outputs=sent_output
        )

    # 4. Resumo de Texto
    with gr.Tab("Resumo"):
        gr.Markdown("### 📝 Resumo Automático de Texto")
        with gr.Row():
            sum_input = gr.Textbox(
                label="Texto para resumir",
                placeholder="Cole aqui o texto que deseja resumir...",
                lines=8
            )
            sum_output = gr.Textbox(
                label="Resumo",
                lines=4
            )
        with gr.Row():
            max_len = gr.Slider(
                minimum=50,
                maximum=500,
                value=150,
                step=10,
                label="Tamanho máximo do resumo"
            )
            min_len = gr.Slider(
                minimum=30,
                maximum=200,
                value=50,
                step=10,
                label="Tamanho mínimo do resumo"
            )
        sum_button = gr.Button("📚 Gerar Resumo")
        sum_button.click(
            services.summarize_text,
            inputs=[sum_input, max_len, min_len],
            outputs=sum_output
        )

    # 5. Extração de Entidades
    with gr.Tab("Entidades"):
        gr.Markdown("### 🔍 Reconhecimento de Entidades")
        with gr.Row():
            ner_input = gr.Textbox(
                label="Texto para análise",
                placeholder="Digite o texto para extrair entidades...",
                lines=5
            )
            ner_output = gr.Textbox(
                label="Entidades encontradas",
                lines=5
            )
        ner_button = gr.Button("🎯 Extrair Entidades")
        ner_button.click(
            services.extract_entities,
            inputs=ner_input,
            outputs=ner_output
        )

    # 6. Classificação de Texto
    with gr.Tab("Classificação"):
        gr.Markdown("### 📋 Classificação de Texto")
        with gr.Row():
            class_input = gr.Textbox(
                label="Texto para classificar",
                placeholder="Digite o texto para classificar...",
                lines=5
            )
            class_output = gr.Textbox(
                label="Classificação",
                lines=5
            )
        class_button = gr.Button("🏷️ Classificar Texto")
        class_button.click(
            services.classify_text,
            inputs=class_input,
            outputs=class_output
        )

    gr.Markdown("""
    ### 📝 Notas:
    - A geração de imagens requer GPU para melhor performance
    - Os modelos são carregados sob demanda para economizar memória
    - Primeira execução pode ser mais lenta devido ao download dos modelos
    - Todos os modelos são open source
    """)

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