import os
import gradio as gr
from huggingface_hub import InferenceClient

# Coloque seu token Hugging Face aqui (ou use variável de ambiente)
HF_TOKEN = os.environ.get("HF_TOKEN") or "SEU_TOKEN_AQUI"

def responder(
    mensagem,
    historico: list[dict[str, str]],
    mensagem_do_sistema,
    max_tokens,
    temperatura,
    top_p,
    hf_token=None,  # Ignorado, usamos HF_TOKEN direto
):
    """
    Função que envia mensagens para o modelo Meta-LLaMA 3.1 8B Instruct usando token direto.
    """
    cliente = InferenceClient(token=HF_TOKEN, model="meta-llama/Meta-Llama-3.1-8B-Instruct")

    # Prepara mensagens no formato chat
    mensagens = [{"role": "system", "content": mensagem_do_sistema}]
    mensagens.extend(historico)
    mensagens.append({"role": "user", "content": mensagem})

    resposta = ""

    # Streaming da resposta token por token
    for trecho_da_mensagem in cliente.chat_completion(
        mensagens,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperatura,
        top_p=top_p,
    ):
        escolhas = trecho_da_mensagem.choices
        token = ""
        if len(escolhas) and escolhas[0].delta.content:
            token = escolhas[0].delta.content

        resposta += token
        yield resposta

# Cria interface de chat Gradio
chatbot = gr.ChatInterface(
    responder,
    type="messages",
    additional_inputs=[
        gr.Textbox(value="Você é um chatbot amigável.", label="Mensagem do sistema"),
        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Máximo de novos tokens"),
        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperatura"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (amostragem de núcleo)"),
    ],
)

# Layout simples, sem login
with gr.Blocks() as demo:
    chatbot.render()

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