import streamlit as st
import os
from src.utils.ingest_text import create_vector_database
from src.utils.ingest_image import extract_and_store_images
from src.utils.text_qa import qa_bot
from src.utils.image_qa import query_and_print_results
import nest_asyncio
nest_asyncio.apply()

from dotenv import load_dotenv
load_dotenv()

def get_answer(query, chain):
    try:
        response = chain.invoke(query)
        return response['result']
    except Exception as e:
        st.error(f"Error in get_answer: {e}")
        return None

st.title("MULTIMODAL DOC QA")

uploaded_file = st.file_uploader("File upload", type="pdf")
if uploaded_file is not None:
    # Save the uploaded file to a temporary location
    temp_file_path = os.path.join("temp", uploaded_file.name)
    os.makedirs("temp", exist_ok=True)  # Ensure the temp directory exists
    with open(temp_file_path, "wb") as f:
        f.write(uploaded_file.getbuffer())

    # Get the absolute path of the saved file
    path = os.path.abspath(temp_file_path)
    st.write(f"File saved to: {path}")
    print(path)

    st.write("Document uploaded successfully!")

if st.button("Start Processing"):
    if uploaded_file is not None:
        with st.spinner("Processing"):
            try:
                client = create_vector_database(path)
                image_vdb = extract_and_store_images(path)
                chain = qa_bot(client)
                st.session_state['chain'] = chain  # Store chain in session state
                st.session_state['image_vdb'] = image_vdb  # Store image_vdb in session state
                st.success("Processing complete.")
            except Exception as e:
                st.error(f"Error during processing: {e}")
    else:
        st.error("Please upload a file before starting processing.")

if user_input := st.chat_input("User Input"):
    if 'chain' in st.session_state and 'image_vdb' in st.session_state:
        chain = st.session_state['chain']
        image_vdb = st.session_state['image_vdb']

        with st.chat_message("user"):
            st.markdown(user_input)

        with st.spinner("Generating Response..."):
            response = get_answer(user_input, chain)
            if response:
                st.markdown(response)
                try:
                    query_and_print_results(image_vdb, user_input)
                except Exception as e:
                    st.error(f"Error querying image database: {e}")
            else:
                st.error("Failed to generate response.")
    else:
        st.error("Please start processing before entering user input.")