import streamlit as st
import requests
import json
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/chat"

def ask_question(question):
    if question in data:
        return data[question]
    else:
        response = requests.post(client_endpoint, json={"context": question})
        if response.status_code == 200:
            answer = response.json()['response']
            data[question] = answer
            with open('data.json', 'w') as f:
                json.dump(data, f)
            return answer
        else:
            return "Error: Failed to retrieve an 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("Question Answering System")

    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=3)
        else:
            st.write("Please enter a question.")

if __name__ == "__main__":
    main()