import gradio as gr
from huggingface_hub import InferenceClient
from transformers import AutoTokenizer, AutoModel
import torch
import numpy as np
import os

# --- 1. Récupérer le token Hugging Face depuis variable d'environnement ---
hf_token = os.environ.get("HF_TOKEN")
if hf_token is None:
    raise ValueError("Tu dois définir la variable d'environnement HF_TOKEN avec ton token Hugging Face.")

# --- 2. Charger SteelBERT pour embeddings ---
steelbert_tokenizer = AutoTokenizer.from_pretrained(
    "MGE-LLMs/SteelBERT", use_auth_token=hf_token
)
steelbert_model = AutoModel.from_pretrained(
    "MGE-LLMs/SteelBERT", use_auth_token=hf_token
).eval()

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
steelbert_model.to(device)

def embed(text):
    inputs = steelbert_tokenizer(text, return_tensors="pt", truncation=True, padding=True).to(device)
    with torch.no_grad():
        outputs = steelbert_model(**inputs, output_hidden_states=True)
    return outputs.hidden_states[-1][:,0,:].cpu().numpy()[0]

# --- 3. Base documentaire (exemple, à remplacer par tes documents techniques) ---
docs = {
    "doc1": "L’acier X42 a une résistance à la traction de 415 MPa.",
    "doc2": "L’acier inoxydable 304 est résistant à la corrosion et à l’oxydation."
}
doc_embeddings = {k: embed(v) for k,v in docs.items()}

def search_best_doc(question):
    q_emb = embed(question)
    def cosine(a,b): return np.dot(a,b)/(np.linalg.norm(a)*np.linalg.norm(b))
    best_doc = max(docs, key=lambda k: cosine(q_emb, doc_embeddings[k]))
    return docs[best_doc]

# --- 4. Fonction de réponse avec Mistral 7B Instruct ---
def respond(
    message,
    history: list[dict[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
    hf_token: gr.OAuthToken,
):
    client = InferenceClient(token=hf_token.token, model="mistralai/Mistral-7B-Instruct-v0.2")

    # Récupérer le contexte pertinent avec SteelBERT
    best_doc = search_best_doc(message)
    context = docs[best_doc]

    # Construire le prompt
    messages = [{"role": "system", "content": system_message}]
    messages.extend(history)
    messages.append({
        "role": "user",
        "content": f"Question: {message}\nContexte: {context}\nRéponds clairement en français :"
    })

    response = ""
    for message in client.chat_completion(
        messages,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    ):
        if len(message.choices) and message.choices[0].delta.content:
            token = message.choices[0].delta.content
            response += token
            yield response

# --- 5. Interface Gradio ---
chatbot = gr.ChatInterface(
    respond,
    type="messages",
    additional_inputs=[
        gr.Textbox(value="Tu es un assistant spécialisé en métallurgie et en acier.", label="System message"),
        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),
    ],
)

with gr.Blocks() as demo:
    with gr.Sidebar():
        gr.LoginButton()
    chatbot.render()

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