from sentence_transformers import SentenceTransformer
import torch
import gradio as gr
from huggingface_hub import InferenceClient

# LOAD DATA
# ===================
# information from dataset was gathered from the spruce
# plant knowledge base used for retrieval (RAG source)
with open("plants.txt", "r", encoding="utf-8") as file:
    plants_text = file.read()

# PREPROCESS TEXT
# ===================
# splits raw dataset into clean, searchable chunks
def preprocess_text(text):
    # strip extra whitespace from the beginning and the end of the text
    cleaned_text = text.strip()
    # split the cleaned_text by every newline character (\n)
    chunks = cleaned_text.split("\n")
    # clean each chunk and store it in cleaned_chunks
    cleaned_chunks = [chunk.strip() for chunk in chunks if chunk.strip() != ""]
    # return the cleaned_chunks
    return cleaned_chunks

# EMBEDDING MODEL
# ===================
# load the pre-trained embedding model that converts text to vectors
model = SentenceTransformer('all-MiniLM-L6-v2')

def create_embeddings(text_chunks):
  # convert each text chunk into a vector embedding and store as a tensor
  chunk_embeddings = model.encode(text_chunks, convert_to_tensor=True)
  # return the chunk_embeddings
  return chunk_embeddings 

# RETRIEVAL FUNCTION
# ===================
# finds the most relevant text chunks based on semantic similarity
def get_top_chunks(query, chunk_embeddings, text_chunks):
  # convert the query text into a vector embedding
  query_embedding = model.encode(query, convert_to_tensor=True)
    
  # normalize the query embedding to unit length for accurate similarity comparison
  query_embedding_normalized = query_embedding / query_embedding.norm()
    
  # normalize all chunk embeddings to unit length for consistent comparison
  chunk_embeddings_normalized = chunk_embeddings / chunk_embeddings.norm(dim=1, keepdim=True)
    
  # calculate cosine similarity between query and all chunks using matrix multiplication
  similarities = torch.matmul(chunk_embeddings_normalized, query_embedding_normalized)
    
  # find the indices of the 3 chunks with highest similarity scores
  top_indices = torch.topk(similarities, k=3).indices
    
  # store the most relevant chunks
  top_chunks = [text_chunks[i.item()] for i in top_indices]
    
  # return the list of most relevant chunks
  return top_chunks

# PREP DATA + EMBEDDINGS
# ===================
# preprocess dataset and create embeddings once at startup
cleaned_chunks = preprocess_text(plants_text)
chunk_embeddings = create_embeddings(cleaned_chunks)

# HUGGING FACE MODEL CLIENT
# ===================
client = InferenceClient("meta-llama/Llama-3.1-8B-Instruct")

# CHAT FUNCTION (RAG PIPELINE)
# ===================
def respond(message, history):
    # retrieve relevant knowledge from dataset
    top_chunks = get_top_chunks(message, chunk_embeddings, cleaned_chunks)
    # convert retrieved chunks into context string
    context = "\n\n".join(top_chunks)
    
    # build system prompt (behavior + rules)
    messages = [{
        "role": "system", 
        "content": (
            "You are a helpful indoor plant care assistant.\n"
            "Introduce yourself as PlantPal AI and what you offer.\n"
            "Only use context when it's relevant.\n"
            "Use the provided context to answer the user's question.\n"
            "Ask the user what plant they have if needed.\n"
            "Ask user if they have any questions or are having problems if needed.\n"
            "Explain clearly and naturally in your own words.\n"
            "Do NOT copy the context word-for-word.\n"
            "Be friendly, inquisitive, and conversational.\n"
            f"Context:\n{context}"
        )
    }]
    
    # add conversation history 
    if history is not None:
        for msg in history or []:
            if isinstance(msg, (list, tuple)) and len(msg) >= 2:
                messages.append({"role": "user", "content": msg[0]})
                messages.append({"role": "assistant", "content": msg[1]})
            
    # add current user question with retrieved context
    messages.append({
        "role": "user", 
        "content": f"""
    Use the context below to answer the question.
Context:
{context}
Question:
{message}
Answer directly and clearly.
"""
    })
    
    # generate streaming response from LLM
    response = ""
    
    stream = client.chat_completion(
        messages,
        max_tokens=500, 
        temperature=0.5,
        top_p=0.7,
        stream=True
    )

    for chunk in stream:
        token = chunk.choices[0].delta.content
        if token: 
            response += token
            yield response

# GRADIO UI
# ===================
chatbot = gr.ChatInterface(respond)
chatbot.launch(debug=True)