import streamlit as st
from pymongo import MongoClient
from gradio_client import Client
import time
import os

# Connect to MongoDB
password = os.getenv("password")
uri = f"mongodb+srv://E:{password}@cluster0.rvt8psd.mongodb.net/?retryWrites=true&w=majority&appName=Cluster0"
client = MongoClient(uri)
db = client["instantqa"]
collection = db["instantqa"]

client_endpoint = "https://skier8402-mistral-super-fast.hf.space/"

def ask_question(question):
    # Check if the question exists in the database
    query = {"question": question}
    result = collection.find_one(query)

    if result:
        return result["answer"]
    else:
        client = Client(client_endpoint)
        result = client.predict(
            question,
            0.9,  # Temperature
            2000,  # Max new tokens
            0.9,  # Top-p (nucleus sampling)
            1.2,  # Repetition penalty
            api_name="/chat"
        )
        answer = result
        answer = answer[:-4]
        
        # Store the question-answer pair in MongoDB
        data = {"question": question, "answer": answer}
        collection.insert_one(data)

        return answer

def typewriter(text: str, speed: int):
    tokens = text.split()
    container = st.empty()
    for index in range(len(tokens) + 1):
        curr_full_text = " ".join(tokens[:index])
        container.markdown(curr_full_text)
        time.sleep(1 / speed)
        # JavaScript to scroll to the bottom of the page while typing
        st.write(f"<script>window.scrollTo(0, document.body.scrollHeight);</script>", unsafe_allow_html=True)

def main():
    st.title("Near Instant Question Answering")
    st.write("Limitations: May generate unhelpful or outdated content, no chat history.")
    st.write("Warning: Do not share any personal/public information.")

    question = st.text_input("Ask your question:")

    if st.button("Ask"):
        if question:
            answer_placeholder = st.empty()
            answer_placeholder.write("Thinking...")
            answer = ask_question(question)
            answer_placeholder.empty()
            typewriter("Answer: " + answer, speed=300)
        else:
            st.write("Please enter a question.")

if __name__ == "__main__":
    main()