"""
RAG Chain module for Study Ocean Student Mentor.
Optimized for Groq (Llama 3.3) and High-Speed Streaming.
"""

from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate
import json
import config

# ─── System Prompt for Student Mentor ────────────────────────
STUDENT_MENTOR_PROMPT = PromptTemplate(
    input_variables=["context", "question"],
    template="""You are **Study Ocean AI Mentor** — a friendly, patient, and knowledgeable tutor who helps students understand their study material deeply.

Your role:
- You solve student doubts with clear, step-by-step explanations
- You use simple language that students can easily understand
- You provide examples and analogies when helpful
- You reference the study material when answering
- If the answer is in the provided context, use it; if not, say so honestly
- You encourage students and make learning enjoyable

**Important Guidelines:**
1. Start with a brief, direct answer, then explain in detail
2. Use bullet points, numbered steps, or tables for clarity
3. If it's a math/science problem, show the solution step by step
4. If it's a theory question, break it into key points
5. Mention which topic/chapter the information relates to when possible
6. 🌟 Use relevant emojis generously in every response!

---
📚 **Relevant Study Material:**
{context}
---
🙋 **Student's Question:**
{question}
---
🎓 **Mentor's Response:**
"""
)

def load_vector_store() -> FAISS:
    """Load the pre-built FAISS vector store from disk. Fallback to ingestion if it fails."""
    embeddings = HuggingFaceEmbeddings(
        model_name=config.EMBEDDING_MODEL,
        model_kwargs={"device": "cpu"},
        encode_kwargs={"normalize_embeddings": True},
    )
    
    try:
        print("📚 Loading vector store...")
        vector_store = FAISS.load_local(
            config.VECTOR_STORE_PATH,
            embeddings,
            allow_dangerous_deserialization=True,
        )
        print("   ✅ Vector store loaded!")
    except Exception as e:
        print(f"   ⚠ Failed to load vector store: {e}")
        print("   🔄 Automatically rebuilding index from PDFs...")
        from ingest import run_ingestion
        run_ingestion()
        vector_store = FAISS.load_local(
            config.VECTOR_STORE_PATH,
            embeddings,
            allow_dangerous_deserialization=True,
        )
    return vector_store

def build_rag_chain(vector_store: FAISS) -> RetrievalQA:
    """Build the RAG chain with Groq LLM."""
    retriever = vector_store.as_retriever(
        search_type="similarity",
        search_kwargs={"k": config.TOP_K_RESULTS},
    )

    # Initialize Groq LLM (Super Fast)
    from langchain_groq import ChatGroq
    import re
    api_key = re.sub(r'[^a-zA-Z0-9_\-]', '', str(config.GROQ_API_KEY))
    
    print(f"   ⚡ Groq Engine active. Key length: {len(api_key)}")
    
    llm = ChatGroq(
        model_name=config.LLM_MODEL,
        groq_api_key=api_key,
        temperature=config.TEMPERATURE,
        max_tokens=config.MAX_OUTPUT_TOKENS,
    )

    chain = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=retriever,
        return_source_documents=True,
        chain_type_kwargs={"prompt": STUDENT_MENTOR_PROMPT},
    )
    return chain

def ask_question(chain: RetrievalQA, question: str) -> dict:
    result = chain.invoke({"query": question})
    sources = []
    seen = set()
    for doc in result.get("source_documents", []):
        source_name = doc.metadata.get("source", "Unknown")
        subject = doc.metadata.get("subject", "General")
        key = f"{subject}/{source_name}"
        if key not in seen:
            seen.add(key)
            sources.append({"file": source_name, "subject": subject})
    return {"answer": result["result"], "sources": sources}

def stream_question(chain: RetrievalQA, question: str):
    """Generator that streams responses from Groq."""
    try:
        print(f"📡 Request: '{question[:50]}...'")
        docs = chain.retriever.get_relevant_documents(question)
        context = "\n\n".join([doc.page_content for doc in docs])
        
        sources = []
        seen = set()
        for doc in docs:
            name = doc.metadata.get("source", "Unknown")
            sub = doc.metadata.get("subject", "General")
            if f"{sub}/{name}" not in seen:
                seen.add(f"{sub}/{name}")
                sources.append({"file": name, "subject": sub})

        llm = chain.combine_documents_chain.llm_chain.llm
        prompt = STUDENT_MENTOR_PROMPT.format(context=context, question=question)
        
        yield f"DATA:{json.dumps({'sources': sources})}\n\n"
        
        for chunk in llm.stream(prompt):
            yield chunk.content
        print("   ✅ Stream complete.")

    except Exception as e:
        print(f"   ❌ Error: {str(e)}")
        yield f"\n\n**Error:** {str(e)}"

# Singleton
_chain_instance = None
def get_chain():
    global _chain_instance
    if _chain_instance is None:
        vector_store = load_vector_store()
        _chain_instance = build_rag_chain(vector_store)
    return _chain_instance