"""
Example: LangChain Agent with Memstate AI Memory

This example demonstrates a full LangChain agent with persistent memory
powered by Memstate AI. The agent can remember decisions across sessions
and recall relevant context using semantic search.

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

import os
import sys

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

from memstate_langchain import create_memstate_tools

try:
    from langchain_openai import ChatOpenAI
    from langchain.agents import create_tool_calling_agent, AgentExecutor
    from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
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:
    print("Set MEMSTATE_API_KEY environment variable")
    sys.exit(1)

if not OPENAI_API_KEY:
    print("Set OPENAI_API_KEY environment variable")
    sys.exit(1)

# ---------------------------------------------------------------------------
# Set up Memstate tools
# ---------------------------------------------------------------------------

PROJECT_ID = "my-fastapi-project"

tools = create_memstate_tools(
    api_key=MEMSTATE_API_KEY,
    project_id=PROJECT_ID,
)

print(f"✓ Loaded {len(tools)} Memstate tools for project '{PROJECT_ID}'")

# ---------------------------------------------------------------------------
# Create the agent
# ---------------------------------------------------------------------------

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

prompt = ChatPromptTemplate.from_messages([
    ("system", f"""You are a helpful software engineering assistant with persistent memory.

You have access to Memstate AI memory tools that let you remember and recall information
across sessions. Use them proactively:

**Before answering questions:**
- Use `memstate_get` to check what you already know about the project
- Use `memstate_search` when you need to find specific information by meaning

**After completing tasks:**
- Use `memstate_remember` to save task summaries and decisions
- Use `memstate_set` for specific key-value facts (config, versions, status)

**Memory workflow:**
1. CHECK: memstate_get() → see full project knowledge
2. ACT: Complete the task using retrieved context
3. SAVE: memstate_remember(content="## Summary\\n- What was done") → persist learnings

Current project: {PROJECT_ID}
"""),
    MessagesPlaceholder(variable_name="chat_history"),
    ("human", "{input}"),
    MessagesPlaceholder(variable_name="agent_scratchpad"),
])

agent = create_tool_calling_agent(llm, tools, prompt)
executor = AgentExecutor(
    agent=agent,
    tools=tools,
    verbose=True,
    max_iterations=8,
    handle_parsing_errors=True,
)

# ---------------------------------------------------------------------------
# Run a realistic multi-turn conversation
# ---------------------------------------------------------------------------

print("\n" + "="*60)
print("Session 1: Setting up the project")
print("="*60)

chat_history = []

# Turn 1: Store initial decisions
response1 = executor.invoke({
    "input": "We've decided to use FastAPI with PostgreSQL and deploy on Railway. "
             "Store these decisions in memory.",
    "chat_history": chat_history,
})
print(f"\nAgent: {response1['output']}\n")

# Turn 2: Store auth decision
response2 = executor.invoke({
    "input": "We're using OAuth2 with Google for authentication. "
             "The JWT tokens use RS256 signing with 24-hour expiry. Remember this.",
    "chat_history": chat_history,
})
print(f"\nAgent: {response2['output']}\n")

print("\n" + "="*60)
print("Session 2: Retrieving context (simulates a new agent session)")
print("="*60)

# New session — agent should retrieve context from Memstate
response3 = executor.invoke({
    "input": "What do you know about this project? What stack are we using?",
    "chat_history": [],  # Fresh chat history
})
print(f"\nAgent: {response3['output']}\n")

response4 = executor.invoke({
    "input": "What authentication approach did we decide on?",
    "chat_history": [],
})
print(f"\nAgent: {response4['output']}\n")

print("\n" + "="*60)
print("Session 3: Updating a decision")
print("="*60)

response5 = executor.invoke({
    "input": "We've decided to add MFA with TOTP to our auth system. "
             "Update the authentication memory to reflect this.",
    "chat_history": [],
})
print(f"\nAgent: {response5['output']}\n")

print("\n✓ Example complete! Check https://memstate.ai/dashboard to see your memories.")
