# my_memory_logic.py
import os

# We'll import the session-based classes from langchain_core if you have them installed:
# If not, you'll need to install the correct package versions or adapt to your environment.
from langchain_core.chat_history import BaseChatMessageHistory
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.runnables.history import RunnableWithMessageHistory

# We'll assume you have a `rag_chain` from your pipeline code or can import it.
# For example:
# from pipeline import rag_chain

# For demonstration, let's just define a dummy "rag_chain" that returns "answer".
# In your real code, import your actual chain. 
class DummyRagChain:
    def invoke(self, inputs):
        # returns a dictionary with "answer"
        return {"answer": f"Dummy answer to '{inputs['input']}'."}

rag_chain = DummyRagChain()

###############################################################################
# 1) We'll keep an in-memory store of session_id -> ChatMessageHistory
###############################################################################
store = {}  # { "abc123": ChatMessageHistory(...) }

def get_session_history(session_id: str) -> BaseChatMessageHistory:
    """
    Retrieve or create a ChatMessageHistory object for the given session_id.
    """
    if session_id not in store:
        store[session_id] = ChatMessageHistory()
    return store[session_id]

###############################################################################
# 2) Create the RunnableWithMessageHistory (conversational chain)
###############################################################################
# If your snippet references `rag_chain`, combine it with get_session_history.
conversational_rag_chain = RunnableWithMessageHistory(
    rag_chain,               # your main chain (RAG or pipeline)
    get_session_history,     # function to fetch chat history for a session
    input_messages_key="input",
    history_messages_key="chat_history",
    output_messages_key="answer"
)

###############################################################################
# 3) A convenience function to run a query with session-based memory
###############################################################################
def run_with_session_memory(user_query: str, session_id: str) -> str:
    """
    A convenience wrapper that calls our `conversational_rag_chain` 
    with a specific session_id. This returns the final 'answer'.
    """
    response = conversational_rag_chain.invoke(
        {"input": user_query},
        config={
            "configurable": {
                "session_id": session_id
            }
        }
    )
    return response["answer"]