import gradio as gr
from huggingface_hub import InferenceClient

from openai import OpenAI
import chromadb
from sentence_transformers import SentenceTransformer
import os

"""
For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
"""
client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")

# Use a basic sentence transformer finetuned for text similarity tasks
model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

def embedding_func(text):
    return model.encode(text, convert_to_numpy=True)  

# Initialize ChromaDB (runs locally)
chroma_client = chromadb.PersistentClient(path="./chroma_db")  

# Create or load a collection with metadata support
collection = chroma_client.get_or_create_collection(
    name="oura_chunks",
    metadata={"hnsw:space": "cosine"}  
)

def query_openai(query, context):
    client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

    # Add instructions to the system message
    system_message = {"role": "system", 
        "content": "You are an expert assistant answering questions about the Oura Ring app.\n"
        "Use only the provided context for your responses.\n"
        "If you do not find relevant information, clearly state that the question is out of scope.\n"
        "Always cite the most relevant article, section, and link that are directly found in the provided context.\n"
        "Do not provide links that are not found in the provided context."}
    
    # Add the user's input
    user_message = {"role": "user", "content": query}

    # Submit system message for response
    completion = client.chat.completions.create(
        model="gpt-4",  
        messages=[
            system_message,  
            {"role": "assistant", "content": context},  
            user_message  
        ],
    )
    
    # Output the chatbot's response
    return completion.choices[0].message.content

# Query ChromaDB for relevant chunks
def retrieve_relevant_chunks(query, n_results=5):
    query_embedding = embedding_func(query).tolist()  
    results = collection.query(
        query_embeddings=[query_embedding], 
        n_results=n_results
    )
    return results

# Perform RAG (Retrieve + Generate) with OpenAI
def perform_rag(query, provider="openai"):
    # Retrieve relevant chunks from ChromaDB
    results = retrieve_relevant_chunks(query, n_results=5)
    
    # Combine the retrieved chunks for LLM input
    relevant_chunks = []
    for document, metadata, distance in zip(results["documents"][0],results["metadatas"][0],results["distances"][0]):
        # Apply relatively lenient filter to return most relevant chunks
        if distance > 0.75:
            continue
        # Structure chunks to include all relevant metadata
        relevant_chunks.append(f'\nArticle Title: {metadata["article_title"]}\nSection Title: {metadata["section_title"]}\nArticle Link: {metadata["link"]}\nRelevance Score: {100*round(1 - distance, 2)}\nArticle Content: {document}\n')
    if len(relevant_chunks) < 1:
        context = f'\nArticle Title: No Relevant Article\nSection Title: No Relevant Article\nArticle Link: No Relevant Article\nRelevance Score: No Relevant Article\nArticle Content: No Relevant Article\n'
    else:
        context = "\n".join(relevant_chunks)  

    # Query OpenAI gpt-4 for a response
    answer = query_openai(query, context)
    
    return answer

def ask_oura_assistant(message, history):
    response = perform_rag(message, "openai")  
    history.append((message, response))  
    return history

with gr.Blocks(theme="soft") as demo:
    gr.Markdown("### Ask me anything about the Oura Ring app!")
    gr.Markdown("# Oura Ring App Assistant")
    
    msg = gr.Textbox(label="Your question", placeholder="Type your question here...")
    submit_btn = gr.Button("Ask")
    
    chatbot = gr.Chatbot()
    
    submit_btn.click(ask_oura_assistant, inputs=[msg, chatbot], outputs=chatbot)


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