import streamlit as st
import ollama
import faiss
import pickle
import numpy as np
from sentence_transformers import SentenceTransformer
import base64
import os
import re
import pandas as pd
import torch
import sys
from typing import List, Dict, Any
import json

# --- Page Configuration ---
st.set_page_config(
    page_title="IISER-K RAG System (Stable)",
    page_icon="✅",
    layout="wide"
)

# --- Constants & Global State ---
VECTOR_STORE_DIR = 'data'
VECTOR_STORE_PATH = os.path.join(VECTOR_STORE_DIR, 'vector_store.faiss')
METADATA_PATH = os.path.join(VECTOR_STORE_DIR, 'chunks_metadata.pkl')
CONFIG_PATH = os.path.join(VECTOR_STORE_DIR, 'config.json')
FOLDER_STRUCTURE_PATH = 'data/folder_structure.txt'
PDF_BASE_PATH = '../'

# --- Caching and Resource Loading ---
@st.cache_data
def load_folder_structure(path: str) -> List[str]:
    if not os.path.exists(path): return []
    with open(path, 'r', encoding='utf-8') as f:
        filenames = [line.strip().split(' -> ')[-1] for line in f if line.strip()]
        return [fn for fn in filenames if '.' in fn]

@st.cache_resource
def load_resources():
    if not all(os.path.exists(p) for p in [METADATA_PATH, VECTOR_STORE_PATH, CONFIG_PATH]):
        st.error(f"Data files not found! Run your ingestion script first.", icon="🚨")
        return [None] * 5
    
    print("Loading resources...")
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    
    with open(CONFIG_PATH, 'r') as f: config = json.load(f)
    with open(METADATA_PATH, 'rb') as f: metadata = pickle.load(f)
    if not metadata:
        st.error("Metadata is empty.", icon="🚨")
        return [None] * 5
    index = faiss.read_index(VECTOR_STORE_PATH)
    model = SentenceTransformer(config['embedding_model'], trust_remote_code=True, device=device)
    folder_structure_list = load_folder_structure(FOLDER_STRUCTURE_PATH)
    
    print("All resources loaded successfully.")
    return index, metadata, model, device, folder_structure_list

faiss_index, metadata, embed_model, device, folder_structure_list = load_resources()

if not faiss_index:
    st.warning("Application halted due to missing data.")
    sys.exit()

# --- ADVANCED RAG LOGIC ---
def generate_grounded_queries(query: str, chat_history: List[Dict], file_list: List[str]) -> Dict:
    history_str = "\n".join([f"{msg['role']}: {msg['content']}" for msg in chat_history[-4:]])
    file_sample = "\n".join(file_list[:100])
    prompt = f"""You are an advanced query analysis system. Your goal is to deconstruct a user's request and generate multiple, high-quality search queries grounded in an available document list. You also create a hypothetical document answer.

**Conversation History:**
---
{history_str}
---

**User's Latest Query:** "{query}"

**Available Documents (Sample):**
---
{file_sample}
---

**Your Task:**
Generate a JSON object with two keys: "query_variations" and "hypothetical_answer".

1.  **`query_variations` (List of 3 strings):** Rewrite the query from three perspectives: grounded, conceptual, and keyword-focused.
2.  **`hypothetical_answer` (String):** Write a detailed, hypothetical answer rich with keywords.

**CRITICAL INSTRUCTION: You MUST output a single, valid JSON object and NOTHING else.**
"""
    try:
        response = ollama.chat(
            model=st.session_state.get('ollama_model', 'gemma3:27b'),
            messages=[{'role': 'user', 'content': prompt}],
            stream=False,
            format='json'
        )
        return json.loads(response['message']['content'])
    except Exception as e:
        st.error(f"Error during query refinement: {e}")
        return {"query_variations": [query], "hypothetical_answer": query}

def reciprocal_rank_fusion(search_results_lists: List[List[int]], k: int = 60) -> Dict[int, float]:
    fused_scores = {}
    for doc_list in search_results_lists:
        for rank, doc_id in enumerate(doc_list):
            if doc_id not in fused_scores: fused_scores[doc_id] = 0
            fused_scores[doc_id] += 1 / (rank + k)
    return fused_scores

def advanced_retrieval(user_query: str, chat_history: List[Dict], file_list: List[str], top_k: int = 10) -> tuple[List[Dict], Dict]:
    generation_result = generate_grounded_queries(user_query, chat_history, file_list)
    query_variations = generation_result.get("query_variations", [user_query])
    hypothetical_answer = generation_result.get("hypothetical_answer", "")
    
    all_queries = [user_query] + query_variations + ([hypothetical_answer] if hypothetical_answer else [])
    query_embeddings = embed_model.encode(all_queries, normalize_embeddings=True)
    
    k_per_query = 20
    distances, indices = faiss_index.search(query_embeddings, k_per_query)
    fused_scores = reciprocal_rank_fusion(indices.tolist())
    sorted_docs = sorted(fused_scores.items(), key=lambda item: item[1], reverse=True)
    
    top_results = []
    for doc_id, score in sorted_docs[:top_k]:
        chunk_meta = metadata[doc_id].copy()
        chunk_meta['retrieval_score'] = score
        top_results.append(chunk_meta)
    
    return top_results, generation_result

def format_generation_prompt(query: str, context_chunks: List[Dict]) -> str:
    prompt = f"""You are a highly knowledgeable academic research assistant specializing in detailed, comprehensive responses. Your expertise lies in synthesizing complex information from multiple sources into thorough, well-structured answers.

**RESPONSE REQUIREMENTS:**
1. **Length & Detail**: Provide a comprehensive, detailed response (minimum 300-500 words when sufficient context is available)
2. **Structure**: Use clear paragraphs, explanations, and logical flow
3. **Citations**: After EVERY sentence that uses information from the chunks, cite using `[Source: chunk_id]`
4. **Synthesis**: Don't just summarize each chunk - integrate information across sources to create a cohesive narrative
5. **Context**: Include relevant background information, explanations of concepts, and detailed examples when available
6. **Completeness**: Address all aspects of the user's question thoroughly

**ANSWER STYLE:**
- Write in an academic but accessible tone
- Provide detailed explanations of concepts and processes
- Include specific examples, data, and evidence from the sources
- Explain the significance and implications of the information
- Connect related ideas across different sources
- Use transitional phrases to create smooth flow between ideas

**User's Question:** {query}

---
**CONTEXT SOURCES:**
---
"""
    for chunk in context_chunks:
        prompt += f"**Chunk ID:** {chunk['chunk_id']}\n**Source File:** {chunk['filename']}\n**Content:** {chunk['original_text']}\n\n---\n"
    
    prompt += """
**INSTRUCTIONS FOR YOUR RESPONSE:**
Write a detailed, comprehensive academic response that thoroughly addresses the user's question. Ensure you:
- Synthesize information from multiple sources where applicable
- Provide detailed explanations and context
- Include specific examples and evidence
- Cite every factual statement with [Source: chunk_id]
- Maintain academic rigor while being accessible
- Create a well-structured, flowing narrative

**DETAILED RESPONSE:**"""
    return prompt

# --- FIXED DISPLAY FUNCTION ---
def display_response_with_sources(response_text: str, context_chunks: List[Dict], message_index: int):
    cited_chunk_ids = sorted(list(set(re.findall(r'\[Source: ([\w-]+)\]', response_text))))
    display_text = re.sub(r'\[Source: [\w-]+\]', '', response_text).strip()
    st.markdown(display_text)
    
    if cited_chunk_ids:
        st.markdown("---")
        st.subheader("📚 Referenced Sources")
        context_map = {chunk['chunk_id']: chunk for chunk in context_chunks}
        
        for chunk_id in cited_chunk_ids:
            chunk_data = context_map.get(chunk_id)
            if chunk_data:
                expander_title = (
                    f"**Source: `{chunk_data['filename']}` "
                    f"(Chunk: {chunk_data['chunk_index'] + 1}/{chunk_data['total_chunks']})"
                )
                with st.expander(expander_title):
                    st.markdown(f"> {chunk_data['original_text']}")
                    col1, col2 = st.columns(2)
                    with col1:
                        # FIXED: Remove the 'key' parameter from link_button
                        st.link_button("Open Web Link 🔗", chunk_data['link'], use_container_width=True)
                    with col2:
                        # Keep the key for regular button as it's supported
                        if st.button("Open Local PDF 📄", key=f"pdf_{message_index}_{chunk_id}", use_container_width=True):
                            st.session_state.pdf_to_display = chunk_data['local_filepath']
                            st.rerun()

# --- UI ---
st.title("⚙️ IISER-K RAG System (Stable)")
st.caption("A professional RAG system powered by a configurable ingestion pipeline.")

with st.sidebar:
    st.header("Settings")
    st.selectbox("Ollama Model", ["gemma3:27b", "llama3:8b", "mistral"], key="ollama_model")
    
    # Add generation parameters controls
    st.subheader("Generation Parameters")
    temperature = st.slider("Temperature", 0.0, 1.0, 0.3, 0.1, help="Controls creativity (0=focused, 1=creative)")
    max_tokens = st.slider("Max Response Length", 512, 4096, 2048, 256, help="Maximum tokens in response")
    
    # Store these in session state for use in generation
    st.session_state.temperature = temperature
    st.session_state.max_tokens = max_tokens
    
    st.markdown("---")
    
    st.header("PDF Viewer")
    if 'pdf_to_display' in st.session_state and st.session_state.pdf_to_display:
        full_path = os.path.join(PDF_BASE_PATH, st.session_state.pdf_to_display)
        if os.path.exists(full_path):
            with open(full_path, "rb") as f:
                base64_pdf = base64.b64encode(f.read()).decode('utf-8')
            st.markdown(f'<iframe src="data:application/pdf;base64,{base64_pdf}" width="100%" height="800"></iframe>', unsafe_allow_html=True)
        else: 
            st.error("Local PDF not found.")

if "messages" not in st.session_state: 
    st.session_state.messages = []

# Display chat history with fixed function
for i, message in enumerate(st.session_state.messages):
    with st.chat_message(message["role"]):
        if message["role"] == "user":
            st.markdown(message["content"])
        else:
            display_response_with_sources(message["content"], message.get("context_chunks", []), i)

if prompt := st.chat_input("Ask a question about IISER-K documents..."):
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"): 
        st.markdown(prompt)
    
    with st.chat_message("assistant"):
        with st.spinner("Analyzing query, performing search, and generating response..."):
            retrieved_chunks, refinement_data = advanced_retrieval(prompt, st.session_state.messages, folder_structure_list)
            
            with st.expander("🔍 Search & Retrieval Strategy"):
                st.write("**Query Generation Details:**")
                st.json(refinement_data)
                st.write("**Top Retrieved Chunks:**")
                df_results = pd.DataFrame(retrieved_chunks)[['filename', 'chunk_id', 'chunk_index', 'retrieval_score']]
                st.dataframe(df_results, use_container_width=True)
            
            generation_prompt = format_generation_prompt(prompt, retrieved_chunks)
            
            # Enhanced generation parameters for more detailed responses
            stream = ollama.chat(
                model=st.session_state.get('ollama_model', 'gemma3:27b'),
                messages=[{'role': 'user', 'content': generation_prompt}],
                stream=True,
                options={
                    'temperature': st.session_state.get('temperature', 0.3),
                    'top_p': 0.9,       # Good balance of coherence and diversity
                    'top_k': 50,        # Reasonable vocabulary diversity
                    'repeat_penalty': 1.1,  # Prevent repetition
                    'num_predict': st.session_state.get('max_tokens', 2048),
                    'num_ctx': 8192,        # Larger context window
                }
            )
            
            # Collect the full response
            response_placeholder = st.empty()
            full_response = ""
            for chunk in stream:
                full_response += chunk['message']['content']
                response_placeholder.markdown(full_response + "▌")
            response_placeholder.empty()
    
    st.session_state.messages.append({
        "role": "assistant",
        "content": full_response,
        "context_chunks": retrieved_chunks
    })
    # Rerun to display the new message
    st.rerun()