import streamlit as st
import streamlit.components.v1 as components
from generate_knowledge_graph import generate_knowledge_graph, answer_question_with_graph

st.set_page_config(
    page_icon="None",
    layout="wide",
    initial_sidebar_state="auto",
    menu_items=None
)

st.title("Knowledge Graph From Text")

# Initialize session state variables
if 'graph_version' not in st.session_state:
    st.session_state['graph_version'] = 0
if 'qa_version' not in st.session_state:
    st.session_state['qa_version'] = 0

st.sidebar.title("Input document")
input_method = st.sidebar.radio(
    "Choose an input method:",
    ("Upload .txt", "Input text")
)

# Text extraction based on user choice
text = ""
if input_method == "Upload .txt":
    uploaded_file = st.sidebar.file_uploader(label="Upload file", type="txt")
    if uploaded_file is not None:
        text = uploaded_file.read().decode("utf-8")
else:
    text = st.sidebar.text_area("Input text", height=300)

if st.sidebar.button("1. Generate Knowledge Graph"):
    if text:
        with st.spinner("Generating knowledge graph..."):
            net, graph_docs = generate_knowledge_graph(text)
            st.session_state['graph_docs'] = graph_docs
            
            output_file = "knowledge_graph.html"
            net.save_graph(output_file)
            with open(output_file, 'r', encoding='utf-8') as f:
                st.session_state['graph_html'] = f.read()
            
            # Increment version to force iframe refresh
            st.session_state['graph_version'] += 1
            
            # Reset QA state for new graph
            st.session_state.pop('qa_answer', None)
            st.session_state.pop('qa_html', None)
            
            st.success("Knowledge graph generated successfully!")
    else:
        st.sidebar.error("Please provide some text to generate the graph.")

# Display the main graph if it exists in session state
if 'graph_html' in st.session_state:
    st.subheader("Initial Knowledge Graph")
    # Append version comment to force Streamlit to refresh the iframe when version changes
    components.html(
        st.session_state['graph_html'] + f"<!-- version {st.session_state['graph_version']} -->", 
        height=600
    )

# QA Section
if 'graph_docs' in st.session_state:
    st.markdown("---")
    st.subheader("Ask a question about the document")

    col1, col2 = st.columns([3, 1])
    with col1:
        question = st.text_input("Your question :")
    with col2:
        k_value = st.slider("Relationships to be analyzed (Top K)", min_value=1, max_value=30, value=15)

    if st.button("2. Analyze") and question:
        with st.spinner("Semantic search in the current graph..."):
            answer, filtered_net = answer_question_with_graph(                question, 
                st.session_state['graph_docs'],
                k_relations=k_value
            )
            
            st.session_state['qa_answer'] = answer
            
            # Read and save the filtered graph HTML content
            with open("filtered_graph.html", 'r', encoding='utf-8') as f:
                st.session_state['qa_html'] = f.read()
                
            st.session_state['qa_version'] += 1

    # Persist the QA results display
    if 'qa_answer' in st.session_state and 'qa_html' in st.session_state:
        st.info(f"**Answer :** {st.session_state['qa_answer']}")
        st.markdown("**Subgraph of the relationships used to answer the question :**")
        # Append version comment to force Streamlit to refresh the iframe when version changes
        components.html(
            st.session_state['qa_html'] + f"<!-- version {st.session_state['qa_version']} -->", 
            height=450
        )