"""
Example: Conversational Memory with Memstate AI (Modern LCEL Pattern)

This example demonstrates how to use ``MemstateMemory`` (alias for
``MemstateSemanticMemory``) to give an LCEL chain persistent, cross-session
semantic memory backed by Memstate AI.

Before each chain call, relevant memories are retrieved via semantic search
and injected into the prompt under the ``relevant_memories`` key.  After each
call, the conversation turn is ingested by the Memstate server, which
automatically extracts keypaths and creates structured memories.

This replaces the deprecated ``ConversationChain`` + ``BaseMemory`` pattern
from LangChain < 0.3.  For LangChain 0.3+, use ``memory.wrap(chain)``.

Prerequisites:
    pip install "memstate-langchain[langchain]"
    export MEMSTATE_API_KEY=mst_your_key_here
    export OPENAI_API_KEY=sk_your_key_here

Run:
    python examples/conversational_memory.py
"""

import os
import sys

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

try:
    from langchain_core.prompts import ChatPromptTemplate
    from langchain_openai import ChatOpenAI
except ImportError:
    print("Install langchain extras: pip install 'memstate-langchain[langchain]'")
    sys.exit(1)

from memstate_langchain import MemstateMemory

MEMSTATE_API_KEY = os.environ.get("MEMSTATE_API_KEY", "")
OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", "")

if not MEMSTATE_API_KEY or not OPENAI_API_KEY:
    print("Set MEMSTATE_API_KEY and OPENAI_API_KEY environment variables")
    sys.exit(1)

# ---------------------------------------------------------------------------
# Build the LCEL chain with Memstate semantic memory
# ---------------------------------------------------------------------------

llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)

# The {relevant_memories} placeholder is populated by MemstateMemory.retrieve()
# before each chain call.
prompt = ChatPromptTemplate.from_messages([
    (
        "system",
        "You are a helpful software engineering assistant.\n\n"
        "{relevant_memories}",
    ),
    ("human", "{input}"),
])

memory = MemstateMemory(
    api_key=MEMSTATE_API_KEY,
    project_id="conversation-demo",
    session_id="alice-session",
    search_limit=5,
    memory_key="relevant_memories",
)

# memory.wrap() returns a RunnableLambda that:
#   1. Calls memory.retrieve(inputs) to inject relevant memories
#   2. Invokes the wrapped chain
#   3. Calls memory.save(inputs, output) to persist the turn
chain = memory.wrap(prompt | llm)

# ---------------------------------------------------------------------------
# Session 1: Establish user preferences
# ---------------------------------------------------------------------------

print("\n" + "=" * 60)
print("Session 1: Learning about the user")
print("=" * 60)

r1 = chain.invoke({"input": "My name is Alice and I'm a Python developer working on a FastAPI project."})
print(f"\nAI: {r1.content}\n")

r2 = chain.invoke({"input": "I prefer PostgreSQL for databases and I'm deploying on Railway."})
print(f"\nAI: {r2.content}\n")

r3 = chain.invoke({"input": "I've decided to use OAuth2 with Google for authentication."})
print(f"\nAI: {r3.content}\n")

# ---------------------------------------------------------------------------
# Session 2: New chain instance — memory persists from Memstate
# ---------------------------------------------------------------------------

print("\n" + "=" * 60)
print("Session 2: New chain instance — testing memory persistence")
print("=" * 60)

memory2 = MemstateMemory(
    api_key=MEMSTATE_API_KEY,
    project_id="conversation-demo",
    session_id="alice-session-2",
    search_limit=5,
    memory_key="relevant_memories",
)

chain2 = memory2.wrap(prompt | llm)

r4 = chain2.invoke({"input": "What do you remember about my project?"})
print(f"\nAI: {r4.content}\n")

r5 = chain2.invoke({"input": "What database am I using?"})
print(f"\nAI: {r5.content}\n")

print("\n✓ Conversational memory example complete!")
print("  Check https://memstate.ai/dashboard to see the stored memories.")
