import streamlit as st
import os
import base64
from langchain.memory import ConversationBufferWindowMemory
from langchain_community.chat_message_histories import StreamlitChatMessageHistory
from utils.ingestion import DocumentProcessor
from utils.llm import LLMProcessor
from utils.qa import QAEngine

# Configure Streamlit page
st.set_page_config(page_title="AI-Powered Document QA", layout="wide")

# Function to encode image in Base64 for avatars
def encode_image(image_path):
    with open(image_path, "rb") as file:
        return base64.b64encode(file.read()).decode()

# Load avatar images
user_avatar = encode_image("./icons/user.jpg")  # Change path if needed
ai_avatar = encode_image("./icons/ai.jpg")

# Initialize document processing & AI components
document_processor = DocumentProcessor()
llm_processor = LLMProcessor()
qa_engine = QAEngine()

# Ensure temp directory exists
os.makedirs("temp", exist_ok=True)

# Sidebar for file upload
st.sidebar.header("Upload a PDF")
uploaded_file = st.sidebar.file_uploader("Choose a PDF file", type=["pdf", "docx", "html", "pptx"])

# Initialize chat memory
memory_storage = StreamlitChatMessageHistory(key="chat_messages")
memory = ConversationBufferWindowMemory(
    memory_key="chat_history", human_prefix="User", chat_memory=memory_storage, k=5
)

# Document upload & processing
if uploaded_file and "document_uploaded" not in st.session_state:
    pdf_path = os.path.join("temp", uploaded_file.name)

    with open(pdf_path, "wb") as f:
        f.write(uploaded_file.read())

    st.sidebar.success("File uploaded successfully!")

    with st.spinner("Processing document..."):
        document_processor.process_document(pdf_path)

    st.sidebar.success("Document processed successfully!")
    st.session_state["document_uploaded"] = True

# Chat UI Header
st.markdown("<h2 style='text-align: center;'>AI Chat Assistant</h2>", unsafe_allow_html=True)
st.markdown("---")

# Display chat history with avatars
for message in memory_storage.messages:
    role = "user" if message.type == "human" else "assistant"
    avatar = user_avatar if role == "user" else ai_avatar  # Assign appropriate avatar

    with st.chat_message(role):
        st.markdown(
            f"""
            <div style="display: flex; align-items: center; margin-bottom: 10px;">
                <img src="data:image/jpeg;base64,{avatar}" width="40" height="40" style="border-radius: 50%; margin-right: 10px;">
                <div style="background-color: #f1f1f1; padding: 10px; border-radius: 10px; max-width: 80%;">{message.content}</div>
            </div>
            """,
            unsafe_allow_html=True
        )

# User input at the bottom
user_input = st.chat_input("Ask me anything...")

if user_input:
    memory_storage.add_user_message(user_input)

    with st.chat_message("user"):
        st.markdown(
            f"""
            <div style="display: flex; align-items: center; margin-bottom: 10px;">
                <img src="data:image/jpeg;base64,{user_avatar}" width="40" height="40" style="border-radius: 50%; margin-right: 10px;">
                <div style="background-color: #d9f7be; padding: 10px; border-radius: 10px; max-width: 80%;">{user_input}</div>
            </div>
            """,
            unsafe_allow_html=True
        )

    with st.spinner("Generating response..."):
        if st.session_state.get("document_uploaded", False):
            answer = qa_engine.query(user_input)
        else:
            answer = llm_processor.generate_answer("", user_input)
            st.warning("No document uploaded. This response is generated from general AI knowledge and may not be document-specific.")

    memory_storage.add_ai_message(answer)

    with st.chat_message("assistant"):
        st.markdown(
            f"""
            <div style="display: flex; align-items: center; margin-bottom: 10px;">
                <img src="data:image/jpeg;base64,{ai_avatar}" width="40" height="40" style="border-radius: 50%; margin-right: 10px;">
                <div style="background-color: #e6f7ff; padding: 10px; border-radius: 10px; max-width: 80%;">{answer.content}</div>
            </div>
            """,
            unsafe_allow_html=True
        )