import streamlit as st
import json
from gradio_client import Client
import time

# Load data.json
with open('data.json') as f:
    data = json.load(f)

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

def ask_question(question):
    if question in data:
        return data[question]
    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]
        data[question] = answer
        with open('data.json', 'w') as f:
            json.dump(data, f)
        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)

def main():
    st.title("Near Instant Question Answering")
    st.write("Limitations: May generate unhelpfull or outdated content, no chat history.")

    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()