import gradio as gr
import os
from huggingface_hub import InferenceClient
import torch
from sentence_transformers import SentenceTransformer
hf_token = os.getenv("HF_TOKEN")
client = InferenceClient("Qwen/Qwen2.5-7B-Instruct",token = hf_token)

#knowledge base (emulating step 2 in the lesson 12 collab)
#load ai generated txt for meal prep info centered on sustainability, budget, and efficiency for students 
def preprocess_text():
    with open("meal_plan.txt", "r", encoding="utf-8") as file:
      text = file.read()
    cleaned_text = text.strip() #strip whitespace from begin & end
    chunks = cleaned_text.split("\n") 
    cleaned_chunks = [chunk.strip() for chunk in chunks if chunk.strip()] #list comprehension
    return cleaned_chunks

cleaned_chunks = preprocess_text()

#making the vector embeddings
model = SentenceTransformer('all-MiniLM-L6-v2')
chunk_embeddings = model.encode(cleaned_chunks, convert_to_tensor=True)

def get_top_chunks(query, chunk_embeddings, text_chunks):
    #convert text to vector embedding
    query_embedding = model.encode(query, convert_to_tensor=True)
    #normalize query embedding
    query_embedding_normalized = query_embedding / query_embedding.norm()
    #normalize all chunk embeddings
    chunk_embeddings_normalized = chunk_embeddings / chunk_embeddings.norm(dim=1, keepdim=True)
    #cosine similarity calcuation
    similarities = torch.matmul(chunk_embeddings_normalized, query_embedding_normalized) 
    #returning list of top most relevant chunks
    top_indices = torch.topk(similarities, k=3).indices
    top_chunks = []
    for i in top_indices:
        chunk = text_chunks[i]
        top_chunks.append(chunk)
    return top_chunks


def respond(message, history):
    relevant = get_top_chunks(message, chunk_embeddings, cleaned_chunks)
    context = "\n".join(relevant)
    SYSTEM_MESSAGE = ("You are a friendly chatbot named PrepBot, a helpful assistant that helps students with healthy and easy meal preparations. Refer to these rules to help the user:", context, "Keep the advice practical, realistic, and under 100 words." )
    messages = [{"role": "system","content":SYSTEM_MESSAGE}]

    response = ""
    for message_chunk in client.chat_completion(
        messages,
        max_tokens=256, 
        temperature = 0.3,
        stream=True, 
    ):
        token = message_chunk.choices[0].delta.content
        response += token
        yield response

#bot_theme = gr.themes.Soft(primary_hue = "green",secondary_hue = "teal")

chatbot = gr.ChatInterface(respond,title = "🥗 PrepBot: Sustainable Student Meal Planning")
chatbot.launch(debug=True)