# ---- Standard libs ----
import os, tempfile, shutil, uuid
from pathlib import Path
from datetime import datetime
import streamlit as st
from dotenv import load_dotenv

# ---- LangChain stack ----
from langchain_groq import ChatGroq
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import Chroma
from langchain.chains import RetrievalQA
from langchain_community.document_loaders import UnstructuredFileLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter

# ---- Load .env ----
load_dotenv()

# ---- Session Management ----
if "session_id" not in st.session_state:
    st.session_state.session_id = str(uuid.uuid4())

def cleanup_session_data():
    session_path = Path(f"./rag_db/{st.session_state.session_id}")
    if session_path.exists():
        try:
            shutil.rmtree(session_path)
        except Exception as e:
            print(f"⚠️ Could not delete {session_path}: {e}")

if "cleaned_up" not in st.session_state:
    cleanup_session_data()
    st.session_state["cleaned_up"] = True

# ---- Initialize RAG Chain ----
@st.cache_resource(show_spinner=False)
def init_rag_chain():
    embedding = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
    vectordb = Chroma(
        persist_directory=f"./rag_db/{st.session_state.session_id}",
        embedding_function=embedding
    )
    llm = ChatGroq(
        groq_api_key=os.getenv("GROQ_API_KEY"),
        model="llama3-70b-8192",
        temperature=0.2,
    )
    return RetrievalQA.from_chain_type(
        llm=llm,
        retriever=vectordb.as_retriever(search_kwargs={"k": 4}),
        return_source_documents=False
    )

qa_chain = init_rag_chain()

def ingest_pdf(file_bytes: bytes, file_name: str):
    embedding = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
    vectordb = Chroma(
        persist_directory=f"./rag_db/{st.session_state.session_id}",
        embedding_function=embedding
    )
    with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp:
        tmp.write(file_bytes)
        tmp_path = tmp.name

    docs = UnstructuredFileLoader(tmp_path).load()
    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
    chunks = splitter.split_documents(docs)

    if chunks:
        vectordb.add_documents(chunks)
        vectordb.persist()

# ---- Streamlit Config ----
st.set_page_config(
    page_title="📄 Multimodal Chat Assistant",
    page_icon="🤖",
    layout="wide",
    initial_sidebar_state="expanded"
)

# ---- Custom CSS ----
st.markdown("""
<style>
.block-container {
    padding-top: 2rem;
    padding-bottom: 2rem;
    max-width: 800px;
    margin: auto;
}
.title {
    font-size: 3rem;
    font-weight: bold;
    color: #4A90E2;
    text-align: center;
    margin-bottom: 1.5rem;
}
.main-container {
    max-width: 800px;
    margin: 0 auto;
    padding: 1rem;
}
.chat-history {
    max-height: 450px;
    overflow-y: auto;
    background-color: #F9F9F9;
    border-radius: 12px;
    padding: 1rem;
    border: 1px solid #DDD;
    margin-bottom: 1rem;
    min-height: 100px;
}
.message-user {
    background-color: #DDF3FF;
    color: #000;
    padding: 0.8rem;
    border-radius: 12px;
    margin-bottom: 0.5rem;
    text-align: left;
    word-wrap: break-word;
}
.message-bot {
    background-color: #E3FFE1;
    color: #000;
    padding: 0.8rem;
    border-radius: 12px;
    margin-bottom: 0.5rem;
    text-align: left;
    word-wrap: break-word;
}
.stTextInput > div > div > input {
    font-size: 1rem;
    padding: 0.75rem;
    border-radius: 8px;
    border: 1px solid #DDD;
    height: 48px;
    min-height: 48px;
}
.stButton > button {
    border-radius: 8px;
    background-color: #4A90E2;
    color: white;
    border: none;
    padding: 0.75rem 1rem;
    font-weight: bold;
    width: 100%;
    height: 48px;
    min-height: 48px;
}
.stButton > button:hover {
    background-color: #357ABD;
}
.stFileUploader > div {
    padding: 1rem;
    border: 2px dashed #4A90E2;
    border-radius: 8px;
    text-align: center;
    margin-bottom: 1rem;
}
.sidebar-logo {
    text-align: center;
    margin-bottom: 1rem;
}
.input-row {
    display: flex;
    gap: 0.5rem;
    align-items: end;
    margin-top: 1rem;
}
.empty-chat {
    text-align: center;
    color: #888;
    font-style: italic;
    padding: 1rem;
}
</style>
""", unsafe_allow_html=True)

# ---- Sidebar ----
with st.sidebar:
    st.markdown("<div class='sidebar-logo'><img src='https://cdn-icons-png.flaticon.com/512/4712/4712027.png' width='100'></div>", unsafe_allow_html=True)
    st.title("🤖 RAG Assistant")
    st.write("• Upload a PDF 📄")
    st.write("• Ask context-specific questions 💬")
    st.write("• Powered by **LLaMA 3 + Groq** ⚡")
    st.markdown("---")
    if st.button("🗑️ Clear Chat", use_container_width=True):
        cleanup_session_data()
        st.session_state.chat_history = []
        st.rerun()

# ---- Title ----
st.markdown("<div class='title'>📄 Multimodal Chat Assistant</div>", unsafe_allow_html=True)

# ---- Main container ----
st.markdown("<div class='main-container'>", unsafe_allow_html=True)

# Upload UI
uploaded_file = st.file_uploader("📤 Upload a PDF file", type=["pdf"])
if uploaded_file:
    with st.spinner("📚 Ingesting document..."):
        ingest_pdf(uploaded_file.read(), uploaded_file.name)
    st.success("✅ Document ingested successfully! You can now ask questions.")

# Initialize chat history
if "chat_history" not in st.session_state:
    st.session_state.chat_history = []

# Chat history display
if st.session_state.chat_history:
    st.markdown("<div class='chat-history'>", unsafe_allow_html=True)
    for msg in st.session_state.chat_history:
        if msg["role"] == "user":
            st.markdown(f"<div class='message-user'><b>🧑 You:</b> {msg['content']}</div>", unsafe_allow_html=True)
        else:
            st.markdown(f"<div class='message-bot'><b>🤖 Assistant:</b> {msg['content']}</div>", unsafe_allow_html=True)
    st.markdown("</div>", unsafe_allow_html=True)

# Input & button
col1, col2 = st.columns([5, 1])
with col1:
    prompt = st.text_input("💬 Type your question", placeholder="Ask about your PDF...", label_visibility="collapsed")

with col2:
    send_clicked = st.button("Send ➡️", use_container_width=True)

# Response handling
if (prompt and send_clicked) or (prompt and st.session_state.get('last_prompt', '') != prompt):
    if prompt.strip():
        st.session_state.chat_history.append({"role": "user", "content": prompt})
        with st.spinner("🤔 Thinking..."):
            try:
                response = qa_chain.run(prompt)
                st.session_state.chat_history.append({"role": "assistant", "content": response})
            except Exception as e:
                error_msg = f"Sorry, I encountered an error: {str(e)}"
                st.session_state.chat_history.append({"role": "assistant", "content": error_msg})
        st.session_state.last_prompt = prompt
        st.rerun()

st.markdown("</div>", unsafe_allow_html=True)