import torch
from sentence_transformers import SentenceTransformer

import gradio as gr
from huggingface_hub import InferenceClient

client = InferenceClient("Qwen/Qwen2.5-7B-Instruct")

with open("knowledgebase.txt", "r", encoding="utf-8") as file:
    knowledgebase_text = file.read()

def preprocess_text(text):

    cleaned_text = text.strip()

    chunks = cleaned_text.split("\n")

    # empty list where cleaned chunks will be stored
    cleaned_chunks = []

    # cleans chunks and adds to our list of cleaned chunks
    for chunk in chunks:
        strip_chunk = chunk.strip()
        if strip_chunk:
            cleaned_chunks.append(strip_chunk)

    return cleaned_chunks

cleaned_chunks = preprocess_text(knowledgebase_text)
    
model = SentenceTransformer('all-MiniLM-L6-v2')

# converts text chunks into vector embedding and stores as tensor for calculations
def create_embeddings(cleaned_chunks):
    
    chunk_embeddings = model.encode(cleaned_chunks, convert_to_tensor=True)

    return chunk_embeddings

chunk_embeddings = create_embeddings(cleaned_chunks)

def get_top_chunks (query, chunk_embeddings, cleaned_chunks):

    # converts query text into vector embedding
    query_embedding = model.encode(query, convert_to_tensor=True)

    query_embedding_normalized = query_embedding/query_embedding.norm()

    chunk_embeddings_normalized = chunk_embeddings/chunk_embeddings.norm(dim=1, keepdim=True)

    similarities = torch.matmul(chunk_embeddings_normalized, query_embedding_normalized)

    top_indices = torch.topk(similarities, k=3).indices

    top_chunks = []
    
    for i in top_indices:
        chunk = cleaned_chunks[i]
        top_chunks.append(chunk)

    return top_chunks

def respond(message, history):

    top_results = get_top_chunks(message, chunk_embeddings, cleaned_chunks)

    # allows LLM to interpret information better; without it would yield a Python list for the LLM to interpret which isn't as clean
    context = "\n".join(top_results)
    
    messages = [{"role": "system", "content": "You are a clairvoyant chatbot with vast knowledge on zodiac signs, but you also give the information very clearly so users can understand the message you're conveying."}]
    
    if history:
        messages.extend(history)
        
    messages.append({
    "role": "user",
    "content": f"""Context from the knowledge base:
{context}

User question:
{message}

Answer using the context above."""
})
    
    response = client.chat_completion(
        messages,
        max_tokens=300
    )
    
    return response.choices[0].message.content.strip()

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    chatbot = gr.ChatInterface(respond)

demo.launch()