import os
import gradio as gr
from huggingface_hub import InferenceClient

# Cliente com token seguro
client = InferenceClient(
    model="mistralai/Mistral-7B-Instruct-v0.3",
    token=os.getenv("HF_TOKEN")
)

# Função de resposta no estilo chat
def responder(mensagem, historico):
    mensagens = []
    if historico is None:
        historico = []

    for entrada in historico:
        if isinstance(entrada, list) and len(entrada) == 2:
            mensagens.append({"role": "user", "content": entrada[0]})
            mensagens.append({"role": "assistant", "content": entrada[1]})

    mensagens.append({"role": "user", "content": mensagem})
    resposta = ""

    try:
        resposta_stream = client.chat_completion(
            messages=mensagens,
            temperature=0.5,
            max_tokens=500,
            stream=True,
        )

        for parte in resposta_stream:
            token = parte.get("choices", [{}])[0].get("delta", {}).get("content", "")
            if token:
                resposta += token
                yield resposta

    except Exception as e:
        print(f"[ERRO] Falha ao consultar modelo: {e}")
        yield "Ocorreu um erro ao gerar a resposta. Verifique o modelo ou o token."

    if not resposta.strip():
        yield "Nenhuma resposta gerada. Tente novamente."

# Interface Gradio
demo = gr.ChatInterface(
    responder,
    title="Benjamin – Assistente Virtual da CEaD - IBC",
    description="Tire dúvidas com minha inteligência artificial (minha base de dados vai até 2021)",
    textbox=gr.Textbox(placeholder="Digite uma mensagem e tecle Enter"),
    type="messages"
)

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