import gradio as gr
from huggingface_hub import InferenceClient
import os

# Initialize client with token from environment
client = InferenceClient(token=os.environ.get("HF_TOKEN"))

def chat(message, history):
    """
    Cascading chat function with multiple AI models
    
    Args:
        message: Current user message (string)
        history: Previous conversation
    
    Returns:
        Combined response from cascaded models (streaming)
    """
    # Build messages in chat format
    messages = []
    
    # Add conversation history
    if history:
        for msg in history:
            if isinstance(msg, dict):
                messages.append(msg)
            elif isinstance(msg, (list, tuple)) and len(msg) == 2:
                messages.append({"role": "user", "content": msg[0]})
                messages.append({"role": "assistant", "content": msg[1]})
    
    # Add current message
    messages.append({"role": "user", "content": message})
    
    try:
        # ===== MODELO 1: Llama 3.2 (Processamento Inicial) =====
        yield "🔄 **[Modelo 1/3]** Processando com Llama 3.2...\n\n"
        
        response1 = client.chat_completion(
            model="meta-llama/Llama-3.2-3B-Instruct",
            messages=messages,
            max_tokens=300,
            stream=False
        )
        
        first_response = response1.choices[0].message.content
        
        yield f"✅ **[Modelo 1: Llama 3.2 3B]**\n{first_response}\n\n---\n\n🔄 **[Modelo 2/3]** Refinando com Qwen 2.5...\n\n"
        
        # ===== MODELO 2: Qwen 2.5 (Refinamento/Análise) =====
        # Criar prompt para o segundo modelo analisar/refinar a resposta
        refinement_prompt = f"""Analise e refine a seguinte resposta, tornando-a mais clara e completa:

Pergunta original: {message}

Resposta para refinar:
{first_response}

Forneça uma versão melhorada e mais detalhada:"""
        
        response2 = client.chat_completion(
            model="Qwen/Qwen2.5-72B-Instruct",
            messages=[{"role": "user", "content": refinement_prompt}],
            max_tokens=400,
            stream=False
        )
        
        second_response = response2.choices[0].message.content
        
        yield f"✅ **[Modelo 1: Llama 3.2 3B]**\n{first_response}\n\n---\n\n✅ **[Modelo 2: Qwen 2.5 72B]**\n{second_response}\n\n---\n\n🔄 **[Modelo 3/3]** Gerando síntese final com Gemma 2...\n\n"
        
        # ===== MODELO 3: Gemma 2 (Síntese Final) =====
        # Criar prompt para o terceiro modelo fazer síntese das respostas anteriores
        synthesis_prompt = f"""Com base nas duas respostas abaixo, crie uma síntese final clara e completa:

Pergunta original: {message}

Resposta 1 (Llama 3.2):
{first_response}

Resposta 2 (Qwen 2.5):
{second_response}

Crie uma resposta final que combine os melhores pontos de ambas, de forma clara e estruturada:"""
        
        response3 = client.chat_completion(
            model="google/gemma-2-9b-it",
            messages=[{"role": "user", "content": synthesis_prompt}],
            max_tokens=500,
            stream=False
        )
        
        third_response = response3.choices[0].message.content
        
        # Resposta final combinada com os 3 modelos
        final_response = f"""✅ **[Modelo 1: Llama 3.2 3B Instruct]**
{first_response}

---

✅ **[Modelo 2: Qwen 2.5 72B Instruct]**
{second_response}

---

✅ **[Modelo 3: Gemma 2 9B IT]**
{third_response}

---

💡 **Modelos utilizados na cascata:**
1. **meta-llama/Llama-3.2-3B-Instruct** - Resposta inicial rápida
2. **Qwen/Qwen2.5-72B-Instruct** - Refinamento e análise detalhada
3. **google/gemma-2-9b-it** - Síntese final combinando ambas respostas"""
        
        yield final_response
                
    except Exception as e:
        yield f"❌ Erro ao processar: {str(e)}"

# Create the Gradio interface
demo = gr.ChatInterface(
    fn=chat,
    title="🤖 Chatbot em Cascata - 3 Modelos de IA",
    description="""**Arquitetura em Cascata com 3 Modelos:**

1️⃣ **Llama 3.2 3B Instruct** (Meta) - Processa a pergunta inicial
2️⃣ **Qwen 2.5 72B Instruct** (Alibaba) - Refina e expande a resposta
3️⃣ **Gemma 2 9B IT** (Google) - Sintetiza uma resposta final combinada

Você verá as respostas dos 3 modelos trabalhando em cascata!""",
    examples=[
        "Explique o que é machine learning",
        "O que é inteligência artificial?",
        "Como funciona uma rede neural?"
    ]
)

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