import gradio as gr
import os
import json
import requests

# Configuración de la API de Groq
GROQ_API_KEY = os.environ.get("GROQ_API_KEY")
GROQ_API_URL = "https://api.groq.com/openai/v1/chat/completions"

# Configuraciones predefinidas (no visibles para el usuario)
SYSTEM_MESSAGE = "Eres Barcenas Economía, un asistente especializado en proporcionar información sobre la economía mexicana. Responde con datos relevantes y actualizados sobre temas económicos de México."
MODEL_NAME = "llama-3.3-70b-specdec"
MAX_TOKENS = 1024
TEMPERATURE = 0.7
TOP_P = 0.95

def respond(message, history: list[tuple[str, str]]):
    # Preparar los mensajes en el formato que Groq espera
    messages = [{"role": "system", "content": SYSTEM_MESSAGE}]
    
    for user_msg, assistant_msg in history:
        if user_msg:
            messages.append({"role": "user", "content": user_msg})
        if assistant_msg:
            messages.append({"role": "assistant", "content": assistant_msg})
    
    messages.append({"role": "user", "content": message})
    
    # Configurar los headers para la API de Groq
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {GROQ_API_KEY}"
    }
    
    # Configurar el payload para la solicitud
    payload = {
        "model": MODEL_NAME,
        "messages": messages,
        "max_tokens": MAX_TOKENS,
        "temperature": TEMPERATURE,
        "top_p": TOP_P,
        "stream": True
    }
    
    # Realizar la solicitud a la API de Groq con streaming
    response = requests.post(
        GROQ_API_URL,
        headers=headers,
        json=payload,
        stream=True
    )
    
    # Manejar la respuesta en streaming
    accumulated_response = ""
    for line in response.iter_lines():
        if line:
            # Las líneas comienzan con "data: " en SSE
            line_text = line.decode('utf-8')
            if line_text.startswith("data: "):
                data_str = line_text[6:]  # Eliminar el prefijo "data: "
                
                # El último mensaje es "data: [DONE]"
                if data_str == "[DONE]":
                    break
                
                try:
                    data = json.loads(data_str)
                    if 'choices' in data and len(data['choices']) > 0:
                        delta = data['choices'][0].get('delta', {})
                        if 'content' in delta and delta['content']:
                            token = delta['content']
                            accumulated_response += token
                            yield accumulated_response
                except json.JSONDecodeError:
                    continue
    
    if not accumulated_response:
        yield "Lo siento, ocurrió un error al procesar tu solicitud."

# Interfaz de Gradio simplificada (sin opciones adicionales visibles)
demo = gr.ChatInterface(
    respond,
    examples=[["¡Bienvenido a Barcenas Economía! Estoy aquí para ayudarte con información sobre la economía mexicana. ¿En qué puedo asistirte hoy?"]]
)

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