"""
Example: LangChain LCEL Chain with Memstate Chat History

This example shows how to use MemstateChatMessageHistory with
RunnableWithMessageHistory for the modern LangChain Expression Language pattern.

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/chat_history.py
"""

import os
import sys

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

from memstate_langchain import MemstateChatMessageHistory

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

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)

PROJECT_ID = "chat-history-demo"

# ---------------------------------------------------------------------------
# Build the LCEL chain
# ---------------------------------------------------------------------------

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

prompt = ChatPromptTemplate.from_messages([
    ("system", "You are a helpful assistant. You remember the full conversation history."),
    MessagesPlaceholder(variable_name="history"),
    ("human", "{input}"),
])

chain = prompt | llm

# Wrap with Memstate chat history
chain_with_history = RunnableWithMessageHistory(
    chain,
    lambda session_id: MemstateChatMessageHistory(
        api_key=MEMSTATE_API_KEY,
        project_id=PROJECT_ID,
        session_id=session_id,
    ),
    input_messages_key="input",
    history_messages_key="history",
)

# ---------------------------------------------------------------------------
# Session 1: Multi-turn conversation
# ---------------------------------------------------------------------------

print("\n" + "="*60)
print("Session 1: Multi-turn conversation with Alice")
print("="*60)

SESSION_ID = "alice-chat-001"
config = {"configurable": {"session_id": SESSION_ID}}

r1 = chain_with_history.invoke({"input": "Hi! My name is Alice."}, config=config)
print(f"\nAI: {r1.content}\n")

r2 = chain_with_history.invoke(
    {"input": "I'm building a FastAPI application with PostgreSQL."},
    config=config,
)
print(f"\nAI: {r2.content}\n")

r3 = chain_with_history.invoke(
    {"input": "What's my name and what am I building?"},
    config=config,
)
print(f"\nAI: {r3.content}\n")

# ---------------------------------------------------------------------------
# Session 2: Resume the same conversation
# ---------------------------------------------------------------------------

print("\n" + "="*60)
print("Session 2: Resuming Alice's conversation (new chain instance)")
print("="*60)

# New chain instance, same session ID — history is loaded from Memstate
chain_with_history2 = RunnableWithMessageHistory(
    chain,
    lambda session_id: MemstateChatMessageHistory(
        api_key=MEMSTATE_API_KEY,
        project_id=PROJECT_ID,
        session_id=session_id,
    ),
    input_messages_key="input",
    history_messages_key="history",
)

r4 = chain_with_history2.invoke(
    {"input": "Do you remember what I told you earlier?"},
    config=config,
)
print(f"\nAI: {r4.content}\n")

# ---------------------------------------------------------------------------
# Different session: Bob's conversation (isolated)
# ---------------------------------------------------------------------------

print("\n" + "="*60)
print("Session 3: Bob's separate conversation (different session ID)")
print("="*60)

bob_config = {"configurable": {"session_id": "bob-chat-001"}}

r5 = chain_with_history2.invoke(
    {"input": "Hi! I'm Bob. Do you know anything about Alice?"},
    config=bob_config,
)
print(f"\nAI: {r5.content}\n")

print("\n✓ Chat history example complete!")
print(f"  Check https://memstate.ai/dashboard to see the stored messages.")
