"""
RAG chatbot module using latest LangChain with LCEL
Handles question-answering with conversation memory using modern patterns
Uses MultiQueryRetriever for improved document retrieval
"""

import os
import time
import logging
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_qdrant import QdrantVectorStore
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.chat_history import BaseChatMessageHistory
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_core.output_parsers import StrOutputParser
from langchain_core.documents import Document
from langchain_classic.retrievers.multi_query import MultiQueryRetriever
from qdrant_client import QdrantClient
from typing import Tuple, List
from operator import itemgetter

# Configure logging for MultiQueryRetriever to see generated query variations
logging.basicConfig()
logging.getLogger("langchain_classic.retrievers.multi_query").setLevel(logging.INFO)

load_dotenv()

# Store for chat sessions: {session_id: {"history": ChatMessageHistory, "last_access": timestamp}}
session_store = {}

# Sessions expire after 1 hour of inactivity
SESSION_TTL_SECONDS = 3600


def cleanup_expired_sessions():
    """
    Remove sessions that haven't been accessed within SESSION_TTL_SECONDS.
    Called on every request to prevent memory buildup (OOM).
    """
    now = time.time()
    expired = [
        sid for sid, data in session_store.items()
        if now - data["last_access"] > SESSION_TTL_SECONDS
    ]
    for sid in expired:
        del session_store[sid]
    if expired:
        logging.info(f"Cleaned up {len(expired)} expired sessions. Active: {len(session_store)}")


def get_session_history(session_id: str) -> BaseChatMessageHistory:
    """
    Get or create chat history for a session.
    Also cleans up expired sessions on each call.
    
    Args:
        session_id: Unique identifier for the session
        
    Returns:
        Chat message history object
    """
    # Clean up expired sessions first
    cleanup_expired_sessions()
    
    if session_id not in session_store:
        session_store[session_id] = {
            "history": ChatMessageHistory(),
            "last_access": time.time()
        }
    else:
        session_store[session_id]["last_access"] = time.time()
    
    return session_store[session_id]["history"]


def format_docs(docs: List[Document]) -> str:
    """
    Format retrieved documents into a single string
    
    Args:
        docs: List of retrieved documents
        
    Returns:
        Formatted string with document contents
    """
    return "\n\n".join(doc.page_content for doc in docs)


def create_rag_chain():
    """
    Create RAG question-answering chain using LCEL (LangChain Expression Language)
    Modern approach with pipe operator for better composability.
    
    Uses MultiQueryRetriever for improved document retrieval.
    
    Returns:
        Tuple of (conversational_rag_chain, retriever, llm)
    """
    
    # 1. Connect to Qdrant
    client = QdrantClient(
        url=os.getenv("QDRANT_URL"),
        api_key=os.getenv("QDRANT_API_KEY")
    )
    
    embeddings = OpenAIEmbeddings(
        model=os.getenv("OPEN_AI_EMBEDDING_MODEL", "text-embedding-3-small")
    )
    
    vectorstore = QdrantVectorStore(
        client=client,
        collection_name=os.getenv("QDRANT_COLLECTION", "hr-intervals"),
        embedding=embeddings
    )
    
    # 2. Base retriever
    base_retriever = vectorstore.as_retriever(
        search_type="similarity",
        search_kwargs={"k": 8}
    )
    
    # 3. Create LLM
    llm = ChatOpenAI(
        model=os.getenv("OPEN_AI_CHAT_MODEL", "gpt-4o"),
        temperature=0.3
    )
    
    # 4. Multi-Query retriever
    retriever = MultiQueryRetriever.from_llm(
        retriever=base_retriever,
        llm=llm,
    )
    
    # 5. System prompt
    system_prompt = """You are an HR assistant for nonprofit organizations in Canada. 
Use the following context to answer questions accurately and helpfully.

IMPORTANT DISCLAIMERS:
- This tool provides general HR information only
- Not a substitute for professional legal or HR advice
- Consult qualified professionals before implementing policies
- Do NOT share personal information about specific individuals

Context:
{context}

INSTRUCTIONS:
- If the context above is empty or contains no relevant information for the question, you MUST explicitly state that you cannot answer based on the knowledge base and recommend the user consult an HR professional or legal advisor for their specific situation. Do NOT make up or guess answers when context is lacking.
- Otherwise, provide a clear, helpful answer grounded in the context. If you're not certain, say so. Always remind users to consult HR/legal professionals for important decisions."""

    prompt = ChatPromptTemplate.from_messages([
        ("system", system_prompt),
        MessagesPlaceholder(variable_name="chat_history"),
        ("human", "{input}")
    ])
    
    # 6. Build RAG chain using LCEL (pipe operator)
    rag_chain = (
        {
            "context": itemgetter("context"),
            "input": itemgetter("input"),
            "chat_history": itemgetter("chat_history")
        }
        | prompt
        | llm
        | StrOutputParser()
    )
    
    # 7. Add chat history with message management
    conversational_rag_chain = RunnableWithMessageHistory(
        rag_chain,
        get_session_history,
        input_messages_key="input",
        history_messages_key="chat_history",
    )
    
    return conversational_rag_chain, retriever, llm


def is_hr_related_question(question: str, llm: ChatOpenAI) -> Tuple[bool, str]:
    """
    Check if the question is HR-related using an LLM classifier.
    
    Args:
        question: User's question
        llm: Reusable ChatOpenAI instance
        
    Returns:
        Tuple of (is_hr_related: bool, rejection_message: str)
    """
    classification_prompt = ChatPromptTemplate.from_messages([
        ("system", """You are a question classifier for an HR chatbot system for nonprofit organizations in Canada.
        
Your task: Determine if the question is related to HR (Human Resources) topics.

HR-related topics include:
- Recruitment, hiring, onboarding
- Employee policies, handbooks, procedures
- Compensation, benefits, payroll
- Performance management, reviews
- Training and development
- Employee relations, workplace issues
- Labor laws, employment standards
- Termination, layoffs, severance
- Health and safety, workplace accommodation
- Diversity, equity, inclusion
- Nonprofit-specific HR matters
- Volunteer management (for nonprofits)

NON-HR topics that should be REJECTED:
- General questions unrelated to workplace/employment
- Technical questions (programming, IT, etc.) unless about HR systems/tools
- Personal advice not related to employment
- Math problems, trivia, general knowledge
- Questions about other business areas (finance, marketing, operations) unless HR-related

Respond with ONLY "YES" if the question is HR-related, or "NO" if it is not."""),
        ("human", "{question}")
    ])
    
    chain = classification_prompt | llm | StrOutputParser()
    result = chain.invoke({"question": question}).strip().upper()
    
    is_hr = result.startswith("YES")
    
    rejection_message = """I apologize, but I am specialized in answering HR (Human Resources) questions only. Your question appears to be outside my area of expertise.

I can help you with HR-related questions such as:
- Recruitment, hiring, and onboarding
- Employee policies and handbooks
- Compensation and benefits management
- Performance management
- Training and development
- Employee relations
- Labor laws and employment standards
- Nonprofit organization HR matters
- Volunteer management

Do you have an HR-related question I can help you with?"""
    
    return is_hr, rejection_message


def ask_question(
    rag_chain, 
    retriever,
    llm: ChatOpenAI,
    question: str, 
    session_id: str = "default",
) -> Tuple[str, List[Document]]:
    """
    Ask a question and get answer with sources.
    
    Args:
        rag_chain: The conversational RAG chain
        retriever: MultiQueryRetriever for document retrieval
        llm: Reusable ChatOpenAI instance
        question: User's question
        session_id: Session identifier for conversation history
        
    Returns:
        Tuple of (answer, source_documents)
    """
    
    is_hr, rejection_message = is_hr_related_question(question, llm)
    if not is_hr:
        return rejection_message, []
    
    sources = retriever.invoke(question)
    context = format_docs(sources)
    
    answer = rag_chain.invoke(
        {"input": question, "context": context},
        config={"configurable": {"session_id": session_id}}
    )
    
    return answer, sources


if __name__ == "__main__":
    print("Initializing chatbot with latest LangChain (LCEL)...")
    rag_chain, retriever, llm = create_rag_chain()
    
    print("\nReady! Enter your question (type 'quit' to exit):\n")
    
    session_id = "test_session"
    
    while True:
        question = input("You: ")
        if question.lower() in ['quit', 'exit', 'q']:
            break
        
        try:
            answer, sources = ask_question(
                rag_chain, retriever, llm, question, session_id
            )
            
            print(f"\nBot: {answer}\n")
            
            if sources:
                print("Sources:")
                for i, doc in enumerate(sources[:3], 1):
                    source = doc.metadata.get("source", "Unknown")
                    print(f"  {i}. {source}")
                print()
        except Exception as e:
            print(f"\nError: {str(e)}")
            print("Make sure you have uploaded some documents first.\n")