import gradio as gr
import requests
import json
import os  # Para acceder a variables de entorno

# Cargar clave API desde Secret configurado en Hugging Face Spaces
API_KEY = os.getenv("SAMBANOVA_API_KEY")
API_URL = "https://api.sambanova.ai/v1/chat/completions"

headers = {
    "Authorization": f"Bearer {API_KEY}",
    "Content-Type": "application/json",
}

def query_deepseek_stream(user_message, chat_history):
    messages = [{"role": "system", "content": "You are a helpful assistant"}]
    for user_msg, ai_msg in chat_history:
        messages.append({"role": "user", "content": user_msg})
        messages.append({"role": "assistant", "content": ai_msg})
    messages.append({"role": "user", "content": user_message})

    json_data = {
        "model": "DeepSeek-V3.1",
        "messages": messages,
        "stream": True,
    }

    try:
        response = requests.post(API_URL, headers=headers, json=json_data, stream=True)
        response.raise_for_status()

        collected_text = ""
        updated_history = chat_history.copy()

        updated_history.append((user_message, ""))

        for line in response.iter_lines(decode_unicode=True):
            if line:
                if line.startswith("data: "):
                    json_str = line[len("data: "):]
                    if json_str == "[DONE]":
                        break
                    try:
                        data = json.loads(json_str)
                        delta = data.get("choices", [{}])[0].get("delta", {})
                        text_fragment = delta.get("content", "")

                        collected_text += text_fragment

                        if updated_history:
                            updated_history[-1] = (updated_history[-1][0], collected_text)

                        yield updated_history, updated_history
                    except json.JSONDecodeError:
                        continue
    except requests.exceptions.RequestException as err:
        error_msg = f"Error: {err}"
        updated_history = chat_history.copy()
        if not updated_history or updated_history[-1][0] != user_message:
            updated_history.append((user_message, error_msg))
        else:
            updated_history[-1] = (updated_history[-1][0], error_msg)
        yield updated_history, updated_history


with gr.Blocks() as demo:
    chat_history = gr.State([])
    chatbot = gr.Chatbot()
    msg = gr.Textbox(label="Escribe tu mensaje")
    btn = gr.Button("Enviar")

    btn.click(query_deepseek_stream, inputs=[msg, chat_history], outputs=[chatbot, chat_history])
    msg.submit(query_deepseek_stream, inputs=[msg, chat_history], outputs=[chatbot, chat_history])

demo.launch()