import os

# Import the PipelineRunnable from pipeline.py
from pipeline import pipeline_runnable

from langchain_core.chat_history import BaseChatMessageHistory
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.runnables.history import RunnableWithMessageHistory

###############################################################################
# 1) In-memory store: session_id -> ChatMessageHistory
###############################################################################
store = {}  # e.g. { "abc123": ChatMessageHistory() }

def get_session_history(session_id: str) -> BaseChatMessageHistory:
    if session_id not in store:
        store[session_id] = ChatMessageHistory()
    return store[session_id]

###############################################################################
# 2) Create RunnableWithMessageHistory
###############################################################################
conversational_rag_chain = RunnableWithMessageHistory(
    pipeline_runnable,
    get_session_history,
    input_messages_key="input",
    history_messages_key="chat_history",
    output_messages_key="answer"
)

###############################################################################
# 3) Convenience function for session-based memory
###############################################################################
def run_with_session_memory(user_query: str, session_id: str) -> str:
    if not user_query:
        raise ValueError("User query is missing. Ensure the 'input' key is provided.")
    
    # Prepare input dictionary with the correct structure
    input_data = {"input": user_query}
    
    # Ensure the session ID is passed correctly in the configuration
    response = conversational_rag_chain.invoke(
        input_data,  # Pass the dictionary containing the 'input' key
        config={"configurable": {"session_id": session_id}}  # Pass session ID in the config
    )
    
    return response.get("answer", "No answer returned from the chain.")