import streamlit as st
import os
from Components.para_utility import load_pdfs_from_file, load_pdfs_from_folder, save_uploaded_file, save_to_user_storage,create_user_storage, get_embedding_path
from Components.para_agent import initialize_model, ConversationalAgent,process_file,demo_file_load
# from whisper import load_model  # Importing Whisper AI
from transformers import pipeline  # Ensure transformers is updated
from Components.video_utility import save_uploaded_video, process_video_voice
import shutil
import atexit
import tempfile
import hashlib
from pathlib import Path

# __import__('pysqlite3')
# import sys
# sys.modules['sqlite3'] = sys.modules.pop('pysqlite3')

import sqlite3

# Page title
st.set_page_config(page_title='Ema Chatbot', page_icon='🤖')
st.title('🤖 Ema chatBot')

uploaded_file=None
uploaded_video_file=None
agent=None

with st.expander('About this app'):
  st.markdown('**What can this app do?**')
  st.info('This app allows users to upload a PDF or Video file about a topic and get a Query response from LLM.')

  st.markdown('**How to use the app?**')
  st.warning('To engage with the app, go to the sidebar and upload a PDF or use the demo PDF. Send a Query and get your answer.'
  'Note: Bigger size of PDF take more time to process')

st.write("It may take a few minutes to generate query response.")


if 'file_processed' not in st.session_state:
    st.session_state.file_processed = False
    st.session_state.agent = None
    st.session_state.file_path = None

# Initialize session state for demo mode
if 'use_demo_pdf' not in st.session_state:
    st.session_state['use_demo_pdf'] = False
    st.session_state['agent']=False

# Sidebar for accepting input parameters
with st.sidebar:
    st.header('1.1. Input data')
    st.markdown('**1. Choose data source**')
    
    # Add demo PDF option
    use_demo = st.checkbox("Use demo LLM PDF ", value=st.session_state['use_demo_pdf'])
    
    if not use_demo:
        uploaded_file = st.file_uploader("Upload a PDF file", type=["pdf"])
        if uploaded_file and not st.session_state.file_processed:
            # Save to user's local storage
            file_path = save_uploaded_file(uploaded_file, "dataset")
            if file_path:
                st.success(f"File saved locally at: {file_path}")
                st.session_state.file_path = file_path
                st.session_state.agent = process_file(file_path)  # Process the file once
                st.session_state.file_processed = True
            else:
                st.error("Failed to save file locally")
    else:
        if not st.session_state.file_processed:
            st.success("Using demo PDF")
            # file_path="../dataset/LLM.pdf"
            st.session_state.agent = demo_file_load()  # Process demo file once
            st.session_state.file_processed = True
    # Use session state to control the checkbox state
    if 'generate_questions' not in st.session_state:
        st.session_state['generate_questions'] = False

    generate_questions_checkbox = st.checkbox("Generate 5 questions from the content", value=st.session_state['generate_questions'])
    st.header('1.2. Upload Video')
    uploaded_video_file = st.file_uploader("Upload a video file", type=["mp4", "avi", "mov"])

if 'query_responses' not in st.session_state:
    st.session_state['query_responses'] = []

def add_query_response(query, response):
    st.session_state.query_responses.append({'query': query, 'response': response})

def summarize_text(text):
    summarizer = pipeline("summarization")
    summary = summarizer(text, max_length=200, min_length=30, do_sample=False)
    return summary[0]['summary_text']

# Custom CSS fpr From
st.markdown("""
<style>
    .flex-container {
        display: flex;
        align-items: center;
        gap: 10px;
    }
    .flex-container .stTextInput {
        flex: 4;
    }
    .flex-container .stButton {
        flex: 1;
    }
    .stTextInput input {
        border: 2px solid #4CAF50;
        padding: 10px;
        border-radius: 5px;
    }
    .stButton button {
        background-color: #4CAF50;
        color: white;
        padding: 10px 20px;
        border: none;
        border-radius: 5px;
        cursor: pointer;
        justify-content:center;
        align-item:basline;
    }
    .stButton button:hover {
        background-color: #45a049;
    }
</style>
""", unsafe_allow_html=True)

with st.form(key="my_form"):
    col1, col2 = st.columns([4,1])  # Create two columns with ratio 4:1
    with col1:
        query = st.text_input("Enter your query", key="query_input")
    with col2:
        submit_button = st.form_submit_button("Enter", type="primary")  # Add a primary colored button
    st.markdown('</div>', unsafe_allow_html=True)



# Reset the checkbox after the operation
if submit_button  and generate_questions_checkbox and (uploaded_file or use_demo):
    query = f"Generate 5 flashcard questions based Context: {query}"
    st.write(query)
    if st.session_state.agent:
        response, Source = st.session_state.agent.ask(query)
        add_query_response(query, response)

        # Uncheck the checkbox after processing
        st.session_state['generate_questions'] = False

        # # Displaying the sources
        # for doc in Source:
        #     page = doc.metadata['page']
        #     snippet = doc.page_content[:200]
        #     Source = {doc.metadata['source']}
        #     source=Source.split('/')[-1]
        #     Content = {doc.page_content[:50]}
        #     st.write(doc.page_content)
        
        # if page:
        #     st.write(response)
        #     st.write("Data taken from source:", Source, " and page No: ", page)
        # if Content:
        #     st.write("Taken content from:", Content)
        query = ""
    else:
        st.write("No documents found.")

# Modify the query processing section
if submit_button and query and not uploaded_video_file and not generate_questions_checkbox:  # Check if button is pressed
    if st.session_state.agent:
        response, Source = st.session_state.agent.ask(query)
        add_query_response(query, response)

        # Displaying the sources
        for doc in Source:
            page = doc.metadata['page']
            snippet = doc.page_content[:200]
            Source = {doc.metadata['source']}
            source=str(Source).split("/")[-1]
            Content = {doc.page_content}
            # print(Source)
        
        if Source and page:
                st.write(response)
                st.write("Data taken from source:", source, " and page No: ", page)
        if Source and Content:
                st.write("Taken content from:", Content)
        # Clear the query input after processing
        # st.session_state.query_input = ""
    else:
        st.write("No documents found.")
elif not query:
    st.write("Enter query.")

if  uploaded_video_file:
    # Save the uploaded video file to a temporary location
    try:
        video_file_path = save_uploaded_video(uploaded_video_file)
        
        # Process the video to extract voice and summarize
        voice_text = process_video_voice(video_file_path)
        # st.write("Voice Data:", voice_text)
        st.markdown(f"**Voice Data:** <span style='font-size: 20px;'>{voice_text}</span>", unsafe_allow_html=True)
        summary = summarize_text(voice_text)
        # st.write("Voice Summary:", summary)
        st.markdown(f"**Voice Summary:** <span style='font-size: 20px;'>{summary}</span>", unsafe_allow_html=True)
    except Exception as e:
        st.error(f"An error occurred: {e}")

st.header('Previous Queries and Responses')

if st.session_state.query_responses:
    for i, qr in enumerate(st.session_state.query_responses, 1):
        st.write(f"{i}. Query: {qr['query']}")
        st.write(f"   Response: {qr['response']}")
else:
    st.write("No queries yet.")

# Add this after session state initialization
if 'needs_cleanup' not in st.session_state:
    st.session_state.needs_cleanup = False