import os
# Increase HF Hub timeout to 60 seconds before any model downloads
os.environ["HF_HUB_TIMEOUT"] = "60"

import streamlit as st
import streamlit.components.v1 as components
from backend import process_medical_query, docs_cache
from visualization import create_medical_graph

def main():
    st.title("AI-Powered Medical Knowledge Graph Assistant")
    st.markdown(
        "**Using BioGPT-Large-PubMedQA + PubMed + Chroma** for advanced retrieval-augmented generation."
    )

    user_query = st.text_input("Enter biomedical/medical query", "Malaria and cough treatment")
    if st.button("Submit"):
        with st.spinner("Generating answer..."):
            final_answer, sub_questions, initial_answer, critique = process_medical_query(user_query)

        st.subheader("Sub-Question Decomposition")
        st.write(sub_questions)

        st.subheader("Initial AI Answer")
        st.write(initial_answer)

        st.subheader("Self-Critique")
        st.write(critique)

        st.subheader("Refined AI Answer")
        st.write(final_answer)

        st.subheader("Knowledge Graph")
        docs = docs_cache.get(user_query, [])
        if docs:
            graph_html = create_medical_graph(user_query, docs)
            components.html(graph_html, height=600, scrolling=True)
        else:
            st.info("No documents to visualize.")

if __name__ == "__main__":
    main()