import hashlib
import numpy as np
import os
import re
import streamlit as st
import streamlit.components.v1 as components
import torch
import glob
from datetime import datetime
from langchain_community.vectorstores import FAISS  # Add this line

#from huggingface_hub import HfFolder
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import CharacterTextSplitter

#from langchain.text_splitter import CharacterTextSplitter
from langchain_community.vectorstores import Chroma
from langchain_huggingface import HuggingFaceEmbeddings
#from sentence_transformers import util  
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

# ====== CONFIGURATION SECTION ======
APP_TITLE = "Educational PDF Chatbot"
APP_LAYOUT = "wide"
MODEL_NAME = "Qwen/Qwen2.5-14B-Instruct"

EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"

CHUNK_SIZE = 1200  # Reduced from 2000
CHUNK_OVERLAP = 100  # Reduced from 300
SEARCH_K = 4  # Reduced from 7
MIN_SIMILARITY_THRESHOLD = 0.1
MAX_CONVERSATION_HISTORY = 6

PDF_SEARCH_PATHS = [
    "*.pdf",
    "Data/*.pdf",
    "documents/*.pdf",
    "pdfs/*.pdf"
]

# Response style configurations
RESPONSE_STYLES = {
    'Balanced': {'temperature': 0.1, 'max_tokens': 500},
    'Concise': {'temperature': 0.05, 'max_tokens': 300},
    'Detailed': {'temperature': 0.2, 'max_tokens': 700}
}

SAFETY_CONFIG = {
    'enable_strict_mode': False,
    'educational_alternatives': True,
    'allow_general_knowledge': True
}

CHATBOT_PASSWORD = st.secrets.get("CHATBOT_PASSWORD", os.getenv("CHATBOT_PASSWORD", "edu123"))

# ====== END CONFIGURATION SECTION ======

os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
os.environ["HF_HUB_DISABLE_EXPERIMENTAL_WARNING"] = "1"
os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "0"

np.float_ = np.float64
st.set_page_config(page_title=APP_TITLE, layout=APP_LAYOUT)

# ====== COPY BUTTON FUNCTION ======
def create_copy_button(text, button_id):
    """Create a working copy button using Streamlit components."""
    button_html = f"""
    <div style="margin-top: 0.5rem;">
        <button 
            id="{button_id}" 
            onclick="copyText_{button_id}()"
            style="
                background-color: #5E35B1;
                color: white;
                border: none;
                padding: 0.4rem 0.8rem;
                border-radius: 0.3rem;
                cursor: pointer;
                font-size: 0.85rem;
                transition: all 0.2s ease;
                font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif;
            "
            onmouseover="this.style.backgroundColor='#7E57C2'; this.style.transform='translateY(-2px)'"
            onmouseout="this.style.backgroundColor='#5E35B1'; this.style.transform='translateY(0)'"
        >
            📋 Copy Response
        </button>
        <span id="feedback_{button_id}" style="color: #4CAF50; font-size: 0.85rem; margin-left: 0.5rem; font-weight: 500;"></span>
    </div>
    <script>
        function copyText_{button_id}() {{
            const text = {repr(text)};
            navigator.clipboard.writeText(text).then(function() {{
                const button = document.getElementById("{button_id}");
                const feedback = document.getElementById("feedback_{button_id}");
                const originalText = button.innerHTML;
                button.innerHTML = '✓ Copied!';
                button.style.backgroundColor = '#4CAF50';
                feedback.innerHTML = '';
                setTimeout(function() {{
                    button.innerHTML = originalText;
                    button.style.backgroundColor = '#5E35B1';
                }}, 2000);
            }}).catch(function(err) {{
                const feedback = document.getElementById("feedback_{button_id}");
                feedback.innerHTML = '❌ Copy failed';
                feedback.style.color = '#f44336';
                console.error('Copy failed:', err);
                setTimeout(function() {{
                    feedback.innerHTML = '';
                }}, 3000);
            }});
        }}
    </script>
    """
    components.html(button_html, height=60)

# ====== AUTHENTICATION SYSTEM ======
def check_password():
    """Check password and manage authentication."""
    if "authenticated" not in st.session_state:
        st.session_state.authenticated = False
    if "login_attempts" not in st.session_state:
        st.session_state.login_attempts = 0
    
    if st.session_state.authenticated:
        return True
    
    st.markdown("""
    <div style='text-align: center; padding: 2rem; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); border-radius: 10px; margin-bottom: 2rem;'>
        <h1 style='color: white; margin-bottom: 1rem;'>🔒 Educational PDF Chatbot</h1>
        <p style='color: white; font-size: 1.1em;'>Authentication required</p>
    </div>
    """, unsafe_allow_html=True)
    
    if st.session_state.login_attempts >= 5:
        st.error("Too many failed attempts. Please wait a few minutes.")
        st.stop()
    
    with st.form("login_form"):
        st.subheader("Enter Password")
        password = st.text_input("Password", type="password", placeholder="Access password")
        submit_button = st.form_submit_button("Access", use_container_width=True)
        
        if submit_button:
            if password == CHATBOT_PASSWORD:
                st.session_state.authenticated = True
                st.session_state.login_attempts = 0
                st.rerun()
            else:
                st.session_state.login_attempts += 1
                remaining = 5 - st.session_state.login_attempts
                if remaining > 0:
                    st.error(f"Incorrect password. Attempts remaining: {remaining}")
                else:
                    st.error("Access temporarily blocked.")
    return False

if not check_password():
    st.stop()

# ====== SESSION STATE INITIALIZATION ======
if "messages" not in st.session_state:
    st.session_state.messages = []
if "conversation_id" not in st.session_state:
    st.session_state.conversation_id = 0
if "model_loaded" not in st.session_state:
    st.session_state.model_loaded = False
if "response_style" not in st.session_state:
    st.session_state.response_style = "Balanced"
if "retriever" not in st.session_state:  # ← ADDED
    st.session_state.retriever = None       # ← ADDED
if "model" not in st.session_state:  # ← ADD
    st.session_state.model = None     # ← ADD
if "tokenizer" not in st.session_state:  # ← ADD
    st.session_state.tokenizer = None     # ← ADD
    
# ====== API CONFIGURATION ======
HF_API_KEY = st.secrets.get("HF_TOKEN", os.getenv("HF_TOKEN"))

#if HF_API_KEY:
    #HfFolder.save_token(HF_API_KEY)

if not HF_API_KEY:
    st.error("Hugging Face API key is missing.")
    st.stop()

@st.cache_resource
def load_quantized_model():    
    """Load model with 4-bit quantization."""
    try:
        quantization_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_compute_dtype=torch.float16,
            bnb_4bit_use_double_quant=True,
            bnb_4bit_quant_type="nf4"
        )
        
        tokenizer = AutoTokenizer.from_pretrained(
            MODEL_NAME, 
            token=HF_API_KEY,
            trust_remote_code=True,
            use_fast=True,
            padding_side="left",
        )
        
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
            tokenizer.pad_token_id = tokenizer.eos_token_id
        
        model = AutoModelForCausalLM.from_pretrained(
            MODEL_NAME,
            device_map="auto",
            torch_dtype=torch.float16,
            quantization_config=quantization_config,
            token=HF_API_KEY,
            low_cpu_mem_usage=True,
            trust_remote_code=True,
        )
        
        return model, tokenizer
    except Exception as e:
        st.error(f"Error loading model: {str(e)}")
        return None, None



def check_question_safety(question):
    """Enhanced safety check with improved precision."""
    question_lower = question.lower().strip()
    
    profanity_patterns = [
        r'\bfuck\b', r'\bfucking\b', r'\bfucked\b', r'\bfucker\b',
        r'\bshit\b', r'\bshitty\b', r'\bshitting\b',
        r'\bbitch\b', r'\bbitching\b', r'\bbitchy\b',
        r'\bass\b(?!\w)', r'\basshole\b',
        r'\bdamn\b(?!\w)',
        r'\bcrap\b', r'\bcrappy\b'
    ]
    
    for pattern in profanity_patterns:
        if re.search(pattern, question_lower):
            return False, "I'd prefer to keep our conversation respectful and educational."
    
    unsafe_patterns = [
        r'\b(kill|murder|hurt|harm|attack|violence|weapon|bomb|suicide)\b',
        r'\bself[\s\-]harm\b',
        r'\b(illegal\s+drugs|hack\s+into|steal\s+from|fraud|piracy|money\s+laundering)\b',
        r'\bhow\s+to\s+(hack|steal|forge|counterfeit)\b',
        r'\b(personal\s+address|phone\s+number|social\s+security|password|credit\s+card)\b',
        r'\bprivate\s+information\b',
        r'\b(hate\s+speech|racist\s+jokes|discrimination\s+against|offensive\s+slur)\b',
        r'\bextremist\s+(content|views|ideology)\b',
        r'\b(sexual\s+content|pornographic|explicit\s+content|adult\s+material)\b',
        r'\b(sexy|erotic|intimate|adult\s+humor)\b',
        r'\b(sexual\s+joke|dirty\s+joke|adult\s+joke)\b'
    ]
    
    for pattern in unsafe_patterns:
        if re.search(pattern, question_lower):
            return False, "I keep conversations appropriate and educational."
    
    educational_inappropriate = [
        'how to cheat on', 'academic dishonesty', 'plagiarism methods', 'fake certificates',
        'exam answers for', 'homework answers for', 'cheat codes for', 'bypass security'
    ]
    
    for phrase in educational_inappropriate:
        if phrase in question_lower:
            return False, "I'm designed to support ethical learning."
    
    return True, ""

def generate_educational_alternative(declined_topic):
    """Provide educational alternatives."""
    if not SAFETY_CONFIG['educational_alternatives']:
        return ""
    
    alternatives = {
        'violence': "I can help with conflict resolution, peace studies, or historical context.",
        'illegal': "I can provide information about legal systems, ethics, or policy studies.",
        'harm': "I can help with safety education, health information, or wellness topics.",
        'academic_dishonesty': "I can help you understand the topic better or provide study strategies."
    }
    
    for key, alternative in alternatives.items():
        if key in declined_topic.lower():
            return f"\n\n{alternative}"
    
    return "\n\nI'm here to help with educational topics and research questions."

def clean_document_text(text):
    """Clean document text."""
    if not text:
        return text
    import unicodedata
    try:
        text = unicodedata.normalize('NFKD', text)
    except:
        pass
    text = re.sub(r'[\u0600-\u06FF\u0750-\u077F\u08A0-\u08FF\uFB50-\uFDFF\uFE70-\uFEFF]', '', text)
    text = re.sub(r'[^\x00-\x7F\u00C0-\u00FF]', '', text)
    text = re.sub(r'\s+', ' ', text)
    return text.strip()

def get_pdf_files():
    """Discover all PDF files."""
    pdf_files = []
    for search_path in PDF_SEARCH_PATHS:
        found_files = glob.glob(search_path)
        pdf_files.extend(found_files)
    pdf_files = list(set(pdf_files))
    pdf_files.sort()
    return pdf_files

PDF_FILES = get_pdf_files()

if not PDF_FILES:
    st.error("No PDF files found.")
    st.stop()


@st.cache_resource
def get_embeddings():
    """Load embeddings on CPU to avoid GPU contention."""
    return HuggingFaceEmbeddings(
        model_name=EMBEDDING_MODEL,
        model_kwargs={"token": HF_API_KEY}  # CPU only - no device="cuda"
    )

@st.cache_resource
def load_and_index_pdfs():
    """Fast in-memory FAISS index - no disk persistence issues."""
    from langchain_community.vectorstores import FAISS
    
    status = st.empty()
    progress = st.progress(0)

    try:
        # 1) Load embeddings first
        status.write("🔢 Loading embedding model...")
        embeddings = get_embeddings()
        status.write("✅ Embedding model ready")
        progress.progress(10)

        # 2) Load PDFs
        status.write("📄 Loading PDFs...")
        documents = []
        total_pdfs = len(PDF_FILES)
        
        for idx, pdf in enumerate(PDF_FILES, start=1):
            status.write(f"📄 Processing {os.path.basename(pdf)} ({idx}/{total_pdfs})...")
            try:
                loader = PyPDFLoader(pdf)
                docs = loader.load()

                for doc in docs:
                    doc.metadata["source"] = f"{os.path.basename(pdf)} (Page {doc.metadata.get('page', 0)+1})"
                    doc.page_content = clean_document_text(doc.page_content)

                documents.extend(docs)
                status.write(f"  ✅ {len(docs)} pages loaded")
            except Exception as e:
                status.write(f"⚠️ Skipping {os.path.basename(pdf)}: {str(e)}")

            progress.progress(10 + int(30 * idx / total_pdfs))

        if not documents:
            status.write("❌ No documents loaded")
            progress.empty()
            status.empty()
            return None

        status.write(f"✅ Loaded {len(documents)} pages total")
        progress.progress(40)

        # 3) Split documents
        status.write("✂️ Splitting into chunks...")
        text_splitter = CharacterTextSplitter(chunk_size=CHUNK_SIZE, chunk_overlap=CHUNK_OVERLAP)
        splits = text_splitter.split_documents(documents)
        status.write(f"✅ Created {len(splits)} chunks")
        progress.progress(50)

        # 4) Embed in batches with FAISS
        status.write(f"🔢 Creating vector index for {len(splits)} chunks...")
        
        BATCH_SIZE = 25  # Small batches for progress visibility
        total_batches = (len(splits) - 1) // BATCH_SIZE + 1
        vectorstore = None
        successful_batches = 0
        
        for i in range(0, len(splits), BATCH_SIZE):
            batch = splits[i:i+BATCH_SIZE]
            batch_num = i // BATCH_SIZE + 1
            
            status.write(f"🔢 Batch {batch_num}/{total_batches} ({len(batch)} chunks)...")
            
            try:
                if vectorstore is None:
                    # Create FAISS index with first batch
                    vectorstore = FAISS.from_documents(batch, embeddings)
                else:
                    # Add subsequent batches
                    vectorstore.add_documents(batch)
                
                successful_batches += 1
                
                # Update progress (50% to 100%)
                done = min(i + len(batch), len(splits))
                batch_progress = 50 + int(50 * done / len(splits))
                progress.progress(batch_progress)
                
            except Exception as e:
                status.write(f"⚠️ Error in batch {batch_num}: {str(e)}")
                # Don't continue - we need to know if embedding is broken
                if successful_batches == 0:
                    raise Exception(f"Failed to embed first batch: {str(e)}")

        if successful_batches == 0:
            status.write("❌ No batches were successfully embedded")
            progress.empty()
            status.empty()
            return None

        progress.progress(100)
        status.write(f"✅ Knowledge base ready! ({successful_batches}/{total_batches} batches)")
        
        # Clear UI after brief pause (but don't block)
        progress.empty()
        status.empty()

        return vectorstore.as_retriever(search_kwargs={"k": SEARCH_K})
        
    except Exception as e:
        status.write(f"❌ Failed: {str(e)}")
        st.error(f"Error building knowledge base: {str(e)}")
        progress.empty()
        status.empty()
        return None    
    
# ← ADDED: Lazy loading helper function
def get_retriever():
    """Lazy load retriever only when needed."""
    if st.session_state.retriever is None:
        with st.spinner("📚 Loading knowledge base... (first time only)"):
            st.session_state.retriever = load_and_index_pdfs()
    return st.session_state.retriever

def get_model():
    """Lazy load model only when generating response."""
    if st.session_state.model is None:
        with st.spinner("🤖 Loading AI model..."):
            st.session_state.model, st.session_state.tokenizer = load_quantized_model()
    return st.session_state.model, st.session_state.tokenizer
    
def clean_message_content(content):
    """Clean message content."""
    if not content:
        return ""
    content = re.sub(r'Source:.*?(?=\n|$)', '', content, flags=re.DOTALL)
    content = re.sub(r'Follow-up.*?(?=\n|$)', '', content, flags=re.DOTALL)
    content = re.sub(r'\n{3,}', '\n\n', content)
    return content.strip()

def needs_pronoun_resolution(query):
    """Check if query contains pronouns."""
    query_lower = query.lower()
    pronouns_to_check = ['they', 'them', 'their', 'it', 'its', 'this', 'that', 'these', 'those']
    return any(f' {pronoun} ' in f' {query_lower} ' or 
              query_lower.startswith(f'{pronoun} ') or 
              query_lower.endswith(f' {pronoun}') 
              for pronoun in pronouns_to_check)

def detect_pronouns_and_resolve(query, conversation_history):
    """Detect and resolve pronouns using context."""
    query_lower = query.lower()
    pronouns = {
        'they': [], 'them': [], 'their': [], 'theirs': [],
        'it': [], 'its': [], 'this': [], 'that': [], 'these': [], 'those': [],
        'he': [], 'him': [], 'his': [], 'she': [], 'her': [], 'hers': []
    }
    
    found_pronouns = []
    for pronoun in pronouns.keys():
        if f' {pronoun} ' in f' {query_lower} ' or query_lower.startswith(f'{pronoun} ') or query_lower.endswith(f' {pronoun}'):
            found_pronouns.append(pronoun)
    
    if not found_pronouns:
        return query, False
    
    if len(conversation_history) < 2:
        return query, False
    
    last_user_msg = ""
    last_assistant_msg = ""
    
    for msg in reversed(conversation_history):
        if msg["role"] == "user" and not last_user_msg:
            last_user_msg = msg["content"]
        elif msg["role"] == "assistant" and not last_assistant_msg:
            last_assistant_msg = clean_message_content(msg["content"])
        if last_user_msg and last_assistant_msg:
            break
    
    potential_referents = []
    entity_patterns = [
        r'\b([A-Z][a-z]+ [A-Z][a-z]+)\b',
        r'\b([a-z]+ [a-z]+(?:ies|tion|ment|ness|ity))\b',
        r'\b(organizations?|institutions?|companies?|governments?|agencies?|groups?)\b',
        r'\b(students?|teachers?|researchers?|scientists?|experts?|professionals?)\b',
        r'\b(countries?|nations?|regions?|communities?|populations?)\b'
    ]
    
    combined_text = f"{last_user_msg} {last_assistant_msg}"
    for pattern in entity_patterns:
        matches = re.findall(pattern, combined_text, re.IGNORECASE)
        potential_referents.extend(matches)
    
    best_referent = None
    for ref in potential_referents:
        if len(ref.split()) > 1:
            best_referent = ref
            break
    
    if not best_referent and potential_referents:
        best_referent = potential_referents[0]
    
    if best_referent:
        expanded_query = query
        for pronoun in found_pronouns:
            if pronoun in ['they', 'them', 'their', 'theirs']:
                if pronoun == 'they':
                    expanded_query = re.sub(rf'\bthey\b', best_referent, expanded_query, flags=re.IGNORECASE)
                elif pronoun == 'them':
                    expanded_query = re.sub(rf'\bthem\b', best_referent, expanded_query, flags=re.IGNORECASE)
                elif pronoun == 'their':
                    expanded_query = re.sub(rf'\btheir\b', f"{best_referent}'s", expanded_query, flags=re.IGNORECASE)
        return expanded_query, True
    
    return query, False

def get_document_topics():
    """Extract clean topics from documents."""
    if not PDF_FILES:
        return []
    topics = []
    for pdf in PDF_FILES:
        filename = os.path.basename(pdf).lower()
        clean_name = filename
        if clean_name.endswith('.pdf'):
            clean_name = clean_name[:-4]
        clean_name = re.sub(r'[_-]+', ' ', clean_name)
        clean_name = re.sub(r'^\d+\s*', '', clean_name)
        stop_words = ['document', 'file', 'report', 'briefing', 'overview', 'web', 'pdf']
        words = [word for word in clean_name.split() if word not in stop_words and len(word) > 2]
        if words:
            clean_topic = ' '.join(words[:4])
            clean_topic = ' '.join(word.capitalize() for word in clean_topic.split())
            topics.append(clean_topic)
    unique_topics = list(dict.fromkeys(topics))[:5]
    return unique_topics

def handle_topic_questions(prompt):
    """Handle questions about available topics."""
    prompt_lower = prompt.lower()
    topic_question_patterns = [
        'what are the other', 'what are the 3 other', 'what are all the topics',
        'what topics', 'what information do you have', 'what can you help with',
        'what documents', 'what subjects', 'what areas', 'list topics',
        'show me the topics', 'what else do you know', 'what other topics'
    ]
    is_topic_question = any(pattern in prompt_lower for pattern in topic_question_patterns)
    if is_topic_question:
        topics = get_document_topics()
        if topics:
            response = f"I have information on these topics:\n\n"
            for i, topic in enumerate(topics, 1):
                response += f"{i}. {topic}\n"
            response += "\nWhich topic would you like to explore?"
            return response, True
    return None, False

def classify_query_type(prompt):
    """Classify query type."""
    prompt_lower = prompt.lower()
    meta_patterns = [
        'what topics', 'what are the other', 'what information',
        'what can you help', 'what documents', 'list topics'
    ]
    if any(pattern in prompt_lower for pattern in meta_patterns):
        return "meta_question"
    if any(phrase in prompt_lower for phrase in ['summarize', 'summarise', 'summary', 'overview']):
        return "summarization"
    return "factual_question"

def validate_response_uses_documents(response, document_content):
    """Check if response uses documents."""
    if not document_content or not response:
        return False
    not_in_docs_phrases = [
        "not in my uploaded documents", "not available in the provided documents",
        "not covered in my documents", "this specific information isn't in"
    ]
    if any(phrase in response.lower() for phrase in not_in_docs_phrases):
        return True
    decline_phrases = ["cannot find", "not in the documents", "not mentioned"]
    if any(phrase in response.lower() for phrase in decline_phrases):
        return False
    if not SAFETY_CONFIG['allow_general_knowledge']:
        general_knowledge_flags = [
            "generally", "typically", "usually", "commonly", "in general"
        ]
        if any(flag in response.lower() for flag in general_knowledge_flags):
            return False
    response_words = set(response.lower().split())
    doc_words = set(document_content.lower().split())
    common_words = {
        'the', 'and', 'or', 'but', 'in', 'on', 'at', 'to', 'for', 'of', 'with', 'by', 
        'is', 'are', 'was', 'were', 'a', 'an', 'this', 'that', 'these', 'those',
        'can', 'will', 'would', 'should', 'could', 'may', 'might', 'must'
    }
    response_words -= common_words
    doc_words -= common_words
    if len(response_words) > 0:
        overlap = len(response_words.intersection(doc_words))
        overlap_ratio = overlap / len(response_words)
        return overlap_ratio >= 0.15
    return False

def build_conversation_context():
    """Build conversation context."""
    if len(st.session_state.messages) <= 1:
        return ""
    start_idx = 1 if st.session_state.messages[0]["role"] == "assistant" else 0
    recent_messages = st.session_state.messages[start_idx:-MAX_CONVERSATION_HISTORY-1:-1]
    recent_messages.reverse()
    context_parts = []
    for msg in recent_messages:
        role = msg["role"]
        content = clean_message_content(msg["content"])
        if content:
            if role == "user":
                context_parts.append(f"User: {content}")
            elif role == "assistant":
                context_parts.append(f"Assistant: {content}")
    return "\n".join(context_parts)

def format_text(text):
    """Basic text formatting."""
    replacements = {
        'alpha': 'α', 'beta': 'β', 'pi': 'π', 'sum': '∑',
        'leq': '≤', 'geq': '≥', 'neq': '≠', 'approx': '≈'
    }
    for latex, unicode_char in replacements.items():
        text = text.replace('\\' + latex, unicode_char)
    return text

def is_self_reference_request(query):
    """Check if asking about previous response."""
    query_lower = query.lower().strip()
    self_reference_patterns = [
        r'\b(your|that)\s+(answer|response|explanation)\b',
        r'\bsummariz(e|ing)\s+(that|your|the)\s+(answer|response)\b',
        r'\b(sum up|recap)\s+(that|your|the)\s+(answer|response)\b',
        r'\bmake\s+(that|your|the)\s+(answer|response)\s+(shorter|brief|concise)\b',
        r'\b(that|your)\s+(previous|last)\s+(answer|response)\b',
        r'\bwhat\s+you\s+just\s+(said|explained|told)\b'
    ]
    simple_self_ref = [
        "can you summarize", "can you summarise", "can you sum up",
        "summarize that", "summarise that", "sum that up",
        "make it shorter", "shorten it", "brief version",
        "recap that", "condense that", "in summary"
    ]
    if any(re.search(pattern, query_lower) for pattern in self_reference_patterns):
        return True
    if any(phrase in query_lower for phrase in simple_self_ref):
        return True
    if query_lower in ["summarize", "summarise", "summary", "sum up", "recap", "brief"]:
        return True
    return False

def is_follow_up_request(query):
    """Check if asking for more information."""
    if is_self_reference_request(query):
        return False
    query_lower = query.lower()
    follow_up_words = [
        "more", "elaborate", "explain", "clarify", "expand", "further", 
        "continue", "what else", "tell me more", "go on", "details",
        "can you", "could you", "please", "also", "additionally"
    ]
    return any(word in query_lower for word in follow_up_words)

def clean_model_output(raw_response):
    """Clean model output."""
    artifacts = [
        "You are an educational assistant", "GUIDELINES:", "DOCUMENT CONTENT:",
        "RECENT CONVERSATION:", "Current question:", "Based on the provided",
        "According to the document", "STRICT RULES:", "Use ONLY", "Do NOT use",
        "SAFETY GUIDELINES", "INSTRUCTIONS:"
    ]
    for artifact in artifacts:
        raw_response = raw_response.replace(artifact, "").strip()
    unwanted_patterns = [
        r'I apologize if.*?[.!]?\s*',
        r'I\'m sorry if.*?[.!]?\s*',
        r'I\'m here to help with.*?[.!]?\s*',
    ]
    for pattern in unwanted_patterns:
        raw_response = re.sub(pattern, '', raw_response, flags=re.IGNORECASE)
    lines = raw_response.split("\n")
    skip_patterns = [
        "answer this question", "question:", "you are an", "be concise",
        "i apologize", "i'm sorry"
    ]
    cleaned_lines = [
        line for line in lines 
        if not any(line.lower().strip().startswith(pattern) for pattern in skip_patterns)
    ]
    cleaned_text = "\n".join(cleaned_lines)
    cleaned_text = re.sub(r'\n{3,}', '\n\n', cleaned_text)
    return cleaned_text.strip()

def create_system_message(has_docs, is_self_ref, document_content="", conversation_context="", last_response=""):
    """Create system message with style adjustment."""
    base_safety_rules = """
SAFETY GUIDELINES:
- Only provide helpful, educational, legal, and appropriate information
- Never provide instructions for illegal activities, violence, or harm
- Do not generate content that could be used to discriminate or harass  
- Refuse inappropriate requests politely and suggest educational alternatives
"""
    
    # Adjust instructions based on response style
    style_instructions = ""
    if st.session_state.response_style == "Concise":
        style_instructions = "\n- Be brief and direct. Provide concise answers without unnecessary details."
    elif st.session_state.response_style == "Detailed":
        style_instructions = "\n- Provide comprehensive, detailed explanations with examples and context."
    else:  # Balanced
        style_instructions = "\n- Provide clear, balanced responses with appropriate detail."
    
    if is_self_ref and last_response:
        return f"""You are an educational assistant. The user is asking you to modify your previous response.

{base_safety_rules}

YOUR PREVIOUS RESPONSE:
{last_response}

CONTEXT:
{conversation_context}

INSTRUCTIONS:
- Provide the requested modification of your previous response
- Be concise and direct{style_instructions}"""
    elif has_docs and SAFETY_CONFIG['allow_general_knowledge']:
        return f"""You are an educational assistant that provides accurate, safe information.

{base_safety_rules}

CONTEXT:
{conversation_context}

DOCUMENT CONTENT:
{document_content}

STRATEGY:
1. Check if the question can be answered using the documents
2. If YES: Provide answer based on document content
3. If NO but question is educational: State "This information isn't in my documents, but I can provide context:" then give general information
4. If inappropriate: Politely decline
5. Be direct, educational, and helpful{style_instructions}"""
    elif has_docs:
        return f"""You are an educational assistant focused on documents.

{base_safety_rules}

CONTEXT:
{conversation_context}

CONTENT:
{document_content}

INSTRUCTIONS:
- Use ONLY the provided content
- If not in documents, state this clearly
- Be direct and educational{style_instructions}"""
    else:
        return f"""You are an educational assistant.

{base_safety_rules}

CONTEXT:
{conversation_context}

The question doesn't match documents. If educational and appropriate, provide general information while being transparent.{style_instructions}"""

def generate_response_from_model(prompt, relevant_docs=None):
    """Generate response with safety checks and style control."""
    try:
        model, tokenizer = get_model()
        
        if model is None or tokenizer is None:
            return "Error: Model could not be loaded."
        
        is_safe, safety_message = check_question_safety(prompt)
        if not is_safe:
            return safety_message + generate_educational_alternative(prompt)
        
        is_self_ref = is_self_reference_request(prompt)
        conversation_context = build_conversation_context()
        
        last_assistant_response = ""
        if is_self_ref and len(st.session_state.messages) >= 2:
            for msg in reversed(st.session_state.messages[:-1]):
                if msg["role"] == "assistant":
                    last_assistant_response = clean_message_content(msg["content"])
                    break
        
        document_content = ""
        has_relevant_docs = False
        if relevant_docs:
            doc_texts = []
            for doc in relevant_docs[:3]:
                doc_texts.append(doc.page_content[:800])
            document_content = "\n\n".join(doc_texts)
            has_relevant_docs = len(doc_texts) > 0
        
        system_message = create_system_message(
            has_docs=has_relevant_docs,
            is_self_ref=is_self_ref,
            document_content=document_content,
            conversation_context=conversation_context,
            last_response=last_assistant_response
        )
        
        user_message = f"Question: {prompt}"
        
        # Get style parameters
        style_config = RESPONSE_STYLES[st.session_state.response_style]
        temperature = style_config['temperature']
        max_tokens = style_config['max_tokens']
        
        # Determine device safely
        try:
            model_device = next(model.parameters()).device
        except Exception as e:
            print(f"⚠️ Could not get model device: {e}")
            model_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        
        print(f"🔍 Using device: {model_device}")
        
        # Try chat template first, fallback to manual formatting
        try:
            if hasattr(tokenizer, "apply_chat_template") and tokenizer.chat_template is not None:
                messages = [
                    {"role": "system", "content": system_message},
                    {"role": "user", "content": user_message}
                ]
                inputs = tokenizer.apply_chat_template(
                    messages, 
                    return_tensors="pt",
                    add_generation_prompt=True,
                    tokenize=True,
                    padding=False
                )
                inputs = inputs.to(model_device)
                input_length = inputs.shape[1]
            else:
                raise AttributeError("No chat template available")
                
        except (AttributeError, Exception) as e:
            print(f"⚠️ Chat template failed, using fallback: {e}")
            # Fallback to manual prompt formatting
            formatted_prompt = f"<|im_start|>system\n{system_message}<|im_end|>\n<|im_start|>user\n{user_message}<|im_end|>\n<|im_start|>assistant\n"
            tokenized = tokenizer(formatted_prompt, return_tensors="pt", padding=False)
            
            # Extract input_ids tensor from tokenizer output
            inputs = tokenized["input_ids"].to(model_device)
            input_length = inputs.shape[1]

        # Generate
        print(f"🔍 Generating with {input_length} input tokens...")
        
        outputs = model.generate(
            inputs, 
            max_new_tokens=max_tokens,
            temperature=temperature,
            top_p=0.8, 
            do_sample=True,
            eos_token_id=tokenizer.eos_token_id,
            pad_token_id=tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id,
            repetition_penalty=1.1,
            use_cache=True
        )

        # Decode output
        if hasattr(outputs, 'shape'):
            # outputs is a tensor
            raw_response = tokenizer.decode(outputs[0][input_length:], skip_special_tokens=True)
        else:
            # outputs might be in a different format
            raw_response = tokenizer.decode(outputs[input_length:], skip_special_tokens=True)
        
        print(f"✅ Generated {len(raw_response)} characters")
        return raw_response.strip()
        
    except Exception as e:
        import traceback
        error_details = traceback.format_exc()
        print(f"❌ Full error traceback:")
        print(error_details)
        return f"Error generating response: {type(e).__name__}: {str(e)}"
def is_conversational_input(prompt):
    """Check if input is conversational."""
    prompt_lower = prompt.lower().strip()
    
    # Remove trailing punctuation for matching
    prompt_clean = re.sub(r'[!.?]+$', '', prompt_lower).strip()
    
    # Exact match patterns
    conversational_patterns = [
        r'^(hi|hello|hey|greetings|howdy)[\s!.?]*$',
        r'^(how\s+are\s+you|how\'s\s+it\s+going|what\'s\s+up|wassup)[\s!.?]*$',
        r'^(good\s+morning|good\s+afternoon|good\s+evening|good\s+night)[\s!.?]*$',
        r'^(thanks|thank\s+you|thx|ty|thank\s+u|thanx)[\s!.?]*$',
        r'^(bye|goodbye|see\s+you|farewell|see\s+ya|later|cya)[\s!.?]*$',
        r'^(clear|reset|start\s+over|new\s+conversation)[\s!.?]*$',
        r'^(ok|okay|alright|sure|yes|yep|yeah|no|nope|got\s+it|understood|i\s+see)[\s!.?]*$',
        r'^(cool|nice|great|awesome|perfect|fine|good)[\s!.?]*$',
        r'^(hmm|hm|mhm|uh\s+huh|aha|oh|ooh|wow)[\s!.?]*$'
    ]
    
    # Check exact patterns first
    if any(re.match(pattern, prompt_clean) for pattern in conversational_patterns):
        return True
    
    # Catch common farewell/thank you phrases that might have extra words
    casual_phrases = [
        'thank you', 'thanks', 'thank u', 'thanx', 'amazing', 'fantastic',
        'have a good day', 'have a great day', 'have a nice day',
        'good day', 'great day', 'nice day',
        'goodbye', 'good bye', 'bye', 'see you', 'see ya',
        'take care', 'cheers'
    ]
    
    # Check if the entire prompt is just a casual phrase (even with punctuation)
    for phrase in casual_phrases:
        if prompt_clean == phrase or prompt_clean.startswith(phrase + ' '):
            return True
    
    return False

def generate_conversational_response(prompt):
    """Generate conversational responses."""
    prompt_lower = prompt.lower().strip()
    document_topics = get_document_topics()
    topic_hint = ""
    if document_topics:
        if len(document_topics) == 1:
            topic_hint = f" I can help with {document_topics[0]}."
        elif len(document_topics) == 2:
            topic_hint = f" I can help with {document_topics[0]} and {document_topics[1]}."
        else:
            topic_hint = f" I can help with {document_topics[0]}, {document_topics[1]}, and more."
    
    conversational_patterns = {
        r'^(hi|hello|hey|greetings|howdy)[\s!.?]*$': 
            (f"Hello!{topic_hint} What would you like to learn?", True),
        r'^(how\s+are\s+you|how\'s\s+it\s+going|what\'s\s+up)[\s!.?]*$': 
            (f"Ready to help!{topic_hint} What interests you?", True),
        r'^(good\s+morning|good\s+afternoon|good\s+evening)[\s!.?]*$': 
            (f"{prompt.capitalize()}!{topic_hint} What would you like to explore?", True),
        r'^(thanks|thank\s+you|thx|ty)[\s!.?]*$': 
            ("You're welcome! Anything else?", True),
        r'^(bye|goodbye|see\s+you|farewell)[\s!.?]*$': 
            ("Goodbye!", False),
        r'^(clear|reset|start\s+over|new\s+conversation)[\s!.?]*$': 
            ("Conversation cleared.", True),
        r'^(ok|okay|alright|sure|got\s+it|understood|i\s+see)[\s!.?]*$':
            ("What else can I help with?", True),
        r'^(yes|yep|yeah)[\s!.?]*$':
            ("What would you like to explore?", True),
        r'^(no|nope)[\s!.?]*$':
            ("Feel free to ask anytime.", True),
        r'^(cool|nice|great|awesome|perfect)[\s!.?]*$':
            ("What else?", True),
        r'^(fine|good)[\s!.?]*$':
            ("What's next?", True),
        r'^(hmm|hm|mhm|uh\s+huh|aha|oh|ooh|wow)[\s!.?]*$':
            ("Something specific you'd like to explore?", True)
    }
    
    for pattern, (response, continue_flag) in conversational_patterns.items():
        if re.match(pattern, prompt_lower):
            return response, continue_flag
    
    return f"I'm here to help.{topic_hint} What interests you?", True

def generate_follow_up_question(context, conversation_length, prompt=None):
    """Generate follow-up question."""
    if prompt and is_self_reference_request(prompt):
        return None
    context_lower = context.lower()
    if "process" in context_lower or "step" in context_lower:
        return "What are the key steps?"
    elif "method" in context_lower or "approach" in context_lower:
        return "How is this applied in practice?"
    elif "benefit" in context_lower or "advantage" in context_lower:
        return "What challenges might arise?"
    simple_questions = [
        "What interests you most?",
        "Would you like to explore related concepts?",
        "Need more details?"
    ]
    return simple_questions[conversation_length % len(simple_questions)]

def process_query(prompt, context_docs):
    """Process query with enhanced features."""
    is_safe, safety_message = check_question_safety(prompt)
    if not is_safe:
        return safety_message + generate_educational_alternative(prompt), None, False, None
    
    if is_conversational_input(prompt):
        response, should_continue = generate_conversational_response(prompt)
        reset_pattern = r'^(clear|reset|start\s+over|new\s+conversation)[\s!.?]*$'
        if re.match(reset_pattern, prompt.lower().strip()):
            return response, None, True, None
        return response, None, False, None
    
    query_type = classify_query_type(prompt)
    if query_type == "meta_question":
        response, handled = handle_topic_questions(prompt)
        if handled:
            return response, None, False, None
    
    is_self_ref = is_self_reference_request(prompt)
    if is_self_ref:
        raw_response = generate_response_from_model(prompt, relevant_docs=None)
        clean_response = clean_model_output(raw_response)
        clean_response = format_text(clean_response)
        return clean_response, None, False, None
    
    if needs_pronoun_resolution(prompt):
        expanded_prompt, was_expanded = detect_pronouns_and_resolve(prompt, st.session_state.messages)
        if was_expanded:
            prompt = expanded_prompt
            st.info(f"Understood: '{prompt}'")
    
    is_followup = is_follow_up_request(prompt)
    #relevant_docs, similarity_scores = check_document_relevance(prompt, context_docs, min_similarity=MIN_SIMILARITY_THRESHOLD)
    relevant_docs = context_docs  # Chroma already ranked these by relevance
    
    raw_response = generate_response_from_model(prompt, relevant_docs if relevant_docs else None)
    clean_response = clean_model_output(raw_response)
    clean_response = format_text(clean_response)
    
    sources = set()
    used_documents = False
    safety_decline_phrases = [
        "keep conversations appropriate", "focused on educational topics",
        "respectful and focused", "designed to support ethical learning"
    ]
    is_safety_decline = any(phrase in clean_response.lower() for phrase in safety_decline_phrases)
    
    if not is_safety_decline and relevant_docs:
        if any(phrase in clean_response.lower() for phrase in [
            "according to the document", "the document shows", "based on the provided",
            "from the document", "the text states", "as mentioned in"
        ]):
            used_documents = True
            for doc in relevant_docs:
                if hasattr(doc, "metadata") and "source" in doc.metadata:
                    sources.add(doc.metadata["source"])
        elif any(phrase in clean_response.lower() for phrase in [
            "not in my uploaded documents", "not available in the provided documents", 
            "not covered in my documents", "this specific information isn't in"
        ]):
            used_documents = False
        else:
            document_content = "\n\n".join([doc.page_content for doc in relevant_docs[:3]])
            if validate_response_uses_documents(clean_response, document_content):
                used_documents = True
                for doc in relevant_docs:
                    if hasattr(doc, "metadata") and "source" in doc.metadata:
                        sources.add(doc.metadata["source"])
    
    if not is_followup and not is_safety_decline and len(st.session_state.messages) % 3 == 0:
        follow_up = generate_follow_up_question(clean_response, len(st.session_state.messages), prompt)
        if follow_up:
            clean_response += f"\n\n{follow_up}"
    
    if used_documents and sources and not is_safety_decline:
        clean_response += f"\n\nSource: {', '.join(sorted(sources))}"
    
    return clean_response, ", ".join(sorted(sources)) if sources else None, False, None

# ====== STREAMLIT INTERFACE ======

# Custom CSS for better visual appearance
st.markdown("""
<style>
/* Message bubbles - different colors for user vs assistant */
[data-testid="stChatMessageContent"] {
    padding: 1rem;
    border-radius: 0.5rem;
}

/* User messages - light blue background */
[data-testid="stChatMessage"][data-testid*="user"] {
    background-color: #E3F2FD;
}

/* Assistant messages - light gray background */
[data-testid="stChatMessage"]:not([data-testid*="user"]) {
    background-color: #F5F5F5;
}

/* Sidebar sections styling */
.stSidebar [data-testid="stMarkdownContainer"] {
    background-color: #FAFAFA;
    padding: 0.5rem;
    border-radius: 0.3rem;
    margin-bottom: 0.5rem;
}

/* Settings section - light purple */
.stSidebar .element-container:has(> [data-testid="stSelectbox"]) {
    background-color: #F3E5F5;
    padding: 1rem;
    border-radius: 0.5rem;
    margin-bottom: 1rem;
}

/* Documents section header styling */
.stSidebar h3 {
    color: #5E35B1;
    font-size: 1rem;
    font-weight: 600;
}

/* Input area focus effect */
[data-testid="stChatInput"] textarea:focus {
    background-color: #F1F8F4;
    border-color: #4CAF50;
}

/* Info messages - light cyan for pronoun resolution */
.stAlert[data-baseweb="notification"] {
    background-color: #E0F7FA;
}

/* Success messages */
[data-testid="stAlert"][kind="success"] {
    background-color: #E8F5E9;
}

/* Error messages keep default red but soften */
[data-testid="stAlert"][kind="error"] {
    background-color: #FFEBEE;
}

/* Buttons styling */
.stButton button {
    border-radius: 0.5rem;
    font-weight: 500;
    transition: all 0.3s ease;
}

.stButton button:hover {
    transform: translateY(-2px);
    box-shadow: 0 4px 8px rgba(0,0,0,0.1);
}

/* Spinner text color */
.stSpinner > div {
    border-top-color: #5E35B1 !important;
}
</style>
""", unsafe_allow_html=True)

st.title(APP_TITLE)

# Sidebar
with st.sidebar:
    st.title("System")
    
    if st.button("Logout", use_container_width=True):
        st.session_state.authenticated = False
        st.session_state.messages = []
        st.session_state.conversation_id = 0
        st.rerun()
    
    st.markdown("---")
    st.subheader("Settings")
    
    # Response style selector
    response_style = st.selectbox(
        "Response Style",
        ["Balanced", "Concise", "Detailed"],
        index=["Balanced", "Concise", "Detailed"].index(st.session_state.response_style),
        help="Control the length and detail of responses"
    )
    
    # Update session state if changed
    if response_style != st.session_state.response_style:
        st.session_state.response_style = response_style
        st.rerun()
    
    st.markdown("---")
    st.write("**Documents:**")
    for pdf in PDF_FILES:
        st.write(f"• {os.path.basename(pdf)}")

# Initialize welcome message
if not st.session_state.messages:
    document_topics = get_document_topics()
    if document_topics:
        if len(document_topics) == 1:
            topic_preview = f"Hello, we can talk about **{document_topics[0]}**, exploring together interesting ideas and reflecting upon our shared challenges, using our curated knowledge base."
        elif len(document_topics) == 2:
            topic_preview = f"Hello, we can talk about **{document_topics[0]}** and **{document_topics[1]}**, exploring together interesting ideas and reflecting upon our shared challenges, using our curated knowledge base."
        else:
            topic_list = ", ".join([f"**{topic}**" for topic in document_topics[:-1]])
            last_topic = f"**{document_topics[-1]}**"
            topic_preview = f"Hello, we can talk about {topic_list}, and {last_topic}, exploring together interesting ideas and reflecting upon our shared challenges, using our curated knowledge base."
    else:
        topic_preview = "Let's explore interesting ideas together using our curated knowledge base."
    
    welcome_msg = topic_preview
    st.session_state.messages.append({"role": "assistant", "content": welcome_msg})

# Clear conversation button
col1, col2 = st.columns([4, 1])
with col2:
    if st.button("New Conversation", use_container_width=True):
        st.session_state.conversation_id += 1
        st.session_state.messages = []
        document_topics = get_document_topics()
        if document_topics:
            if len(document_topics) == 1:
                topic_preview = f"Hello, we can talk about **{document_topics[0]}**, exploring together interesting ideas and reflecting upon our shared challenges, using our curated knowledge base."
            elif len(document_topics) <= 3:
                topic_list = " and ".join([f"**{topic}**" for topic in document_topics])
                topic_preview = f"Hello, we can talk about {topic_list}, exploring together interesting ideas and reflecting upon our shared challenges, using our curated knowledge base."
            else:
                topic_list = ", ".join([f"**{topic}**" for topic in document_topics[:2]])
                topic_preview = f"Hello, we can talk about {topic_list}, and more, exploring together interesting ideas and reflecting upon our shared challenges, using our curated knowledge base."
        else:
            topic_preview = "Let's explore interesting ideas together using our curated knowledge base."
        welcome_msg = topic_preview
        st.session_state.messages.append({"role": "assistant", "content": welcome_msg})
        st.rerun()

# ← CHANGED: Use get_retriever() instead of checking if retriever exists
retriever = get_retriever()

if retriever:
    # Display messages with copy button for assistant responses
    for idx, message in enumerate(st.session_state.messages):
        with st.chat_message(message["role"]):
            st.markdown(message["content"])
            
            # Add copy button for assistant messages
            if message["role"] == "assistant":
                button_id = f"copy_btn_{idx}_{st.session_state.conversation_id}"
                create_copy_button(message["content"], button_id)

    # User input
    if prompt := st.chat_input("What would you like to learn?"):
        st.session_state.messages.append({"role": "user", "content": prompt})
        with st.chat_message("user"):
            st.markdown(prompt)
        
        with st.chat_message("assistant"):
            try:
                # Only retrieve documents if NOT conversational
                if is_conversational_input(prompt):
                    retrieved_docs = None
                else:
                    retrieved_docs = retriever.invoke(prompt)
                
                answer, sources, should_reset, new_follow_up = process_query(prompt, retrieved_docs)                
                if should_reset:
                    st.session_state.conversation_id += 1
                    st.session_state.messages = []
                    st.session_state.messages.append({"role": "assistant", "content": answer})
                    st.rerun()
                
                st.session_state.messages.append({"role": "assistant", "content": answer})
                st.markdown(answer)
                
                # Add copy button for the new response
                button_id = f"copy_btn_new_{st.session_state.conversation_id}_{len(st.session_state.messages)}"
                create_copy_button(answer, button_id)
            except Exception as e:
                error_msg = f"Error: {str(e)}"
                st.error(error_msg)
                st.session_state.messages.append({"role": "assistant", "content": error_msg})
else:
    st.error("Failed to load document retrieval system.")