import subprocess
import sys

# Force reinstall pydantic v1 before any other imports to ensure compatibility with Chromadb.
subprocess.run([sys.executable, "-m", "pip", "install", "--force-reinstall", "pydantic==1.10.7"])

import os
# Set HF Hub timeout to 60 seconds.
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."
    )

    # Clinical query input – designed for clarity and ease-of-use.
    user_query = st.text_input("Enter biomedical/medical query", "Malaria and cough treatment")
    if st.button("Submit"):
        with st.spinner("Generating answer..."):
            # Process the query using the streamlined backend.
            final_answer, sub_questions, initial_answer, critique = process_medical_query(user_query)

        st.subheader("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()