import streamlit as st
import os
from model_processor import LlamaProcessor

st.set_page_config(page_title="Llama PDF Expert", layout="wide")

# --- UI Header ---
st.title("📄 PDF QA with Llama 3.2")
st.markdown("Upload a document and ask questions using a local Llama model.")

# --- Sidebar Configuration ---
with st.sidebar:
    st.header("1. Authentication")
    token = st.text_input("Hugging Face Token", type="password", help="Enter your HF 'Read' token.")
    model_choice = st.selectbox("Select Model", ["meta-llama/Llama-3.2-1B-Instruct", "meta-llama/Llama-3.2-3B-Instruct"])
    
    st.divider()
    st.header("2. Document Upload")
    uploaded_file = st.file_uploader("Upload PDF", type="pdf")

# --- Session State Initialization ---
if "processor" not in st.session_state:
    st.session_state.processor = None
if "vector_db" not in st.session_state:
    st.session_state.vector_db = None
if "messages" not in st.session_state:
    st.session_state.messages = []

# --- Logic: Model Loading & Processing ---
if uploaded_file and token:
    if st.session_state.processor is None:
        try:
            with st.spinner("Initializing Llama... (this may take a minute)"):
                st.session_state.processor = LlamaProcessor(model_choice, token)
            
            # Save and process PDF
            with open("temp_upload.pdf", "wb") as f:
                f.write(uploaded_file.getbuffer())
            
            with st.spinner("Indexing document..."):
                st.session_state.vector_db = st.session_state.processor.process_pdf("temp_upload.pdf")
            st.success("Document processed! Ready to chat.")
        except Exception as e:
            st.error(f"Error initializing: {str(e)}")

# --- Logic: Chat Interface ---
# Display chat history
for msg in st.session_state.messages:
    with st.chat_message(msg["role"]):
        st.markdown(msg["content"])

# User Input
if prompt := st.chat_input("Ask a question about the PDF..."):
    if not st.session_state.vector_db:
        st.warning("Please upload a PDF and provide a token first.")
    else:
        # User message
        st.session_state.messages.append({"role": "user", "content": prompt})
        with st.chat_message("user"):
            st.markdown(prompt)

        # Assistant response
        with st.chat_message("assistant"):
            with st.spinner("Thinking..."):
                answer = st.session_state.processor.get_answer(prompt, st.session_state.vector_db)
                st.markdown(answer)
                st.session_state.messages.append({"role": "assistant", "content": answer})