from src.memory.manager import MemoryManager
from src.memory.base import MemoryType
from src.storage.mongodb_client import MongoDBClient, MongoDBConfig

# 1. Connect to MongoDB
config = MongoDBConfig(uri="mongodb+srv://...", database="my_app")
db_client = MongoDBClient()
await db_client.initialize(config)

# 2. Create Memory Manager
memory = MemoryManager(db_client.db)

# 3. Store memories (from your agent's conversations)
await memory.store_memory(
    content="User said they love Python and work at Google",
    memory_type=MemoryType.EPISODIC,
    agent_id="my_agent",
    user_id="user_123"
)

# 4. Retrieve relevant memories (before responding)
memories = await memory.episodic.retrieve(
    query="What programming language does the user like?",
    agent_id="my_agent",
    user_id="user_123",
    limit=5
)

# 5. Use memories in your agent's context
for mem in memories:
    print(f"Remembered: {mem.content}")

# Agent A's memories are separate from Agent B's
await memory.store_memory(..., agent_id="agent_a", user_id="user_1")
await memory.store_memory(..., agent_id="agent_b", user_id="user_1")

entities = await memory.extract_entities(
    text="I'm John, a software engineer at Google",
    agent_id="my_agent",
    llm=your_llm
)
# Returns: [{"name": "John", "type": "PERSON"}, {"name": "Google", "type": "ORGANIZATION"}]

from src.context.engineer import ContextEngineer

engineer = ContextEngineer()
if engineer.should_compress(context, model="gpt-4"):
    compressed = await engineer.compress(context, llm)

# Hybrid search is the DEFAULT - combines semantic + keyword matching
memories = await memory.episodic.retrieve(
    query="What does the user like?",
    agent_id="my_agent",
    user_id="user_123",
    limit=5,
    threshold=0.5,
    search_mode="hybrid"  # Default - can also use "semantic" or "text"
)

from langchain_openai import ChatOpenAI
from src.memory.manager import MemoryManager

llm = ChatOpenAI(model="gpt-4")
memory = MemoryManager(db)

# Before each LLM call, retrieve relevant memories
memories = await memory.episodic.retrieve(query=user_input, agent_id="my_agent")
context = "\n".join([m.content for m in memories])

response = llm.invoke(f"Context:\n{context}\n\nUser: {user_input}")

# After LLM response, store the interaction
await memory.store_memory(
    content=f"User: {user_input}\nAssistant: {response}",
    memory_type=MemoryType.EPISODIC,
    agent_id="my_agent"
)

from langgraph.graph import StateGraph
from src.memory.manager import MemoryManager

memory = MemoryManager(db)

def memory_node(state):
    # Retrieve memories before processing
    memories = await memory.episodic.retrieve(
        query=state["input"],
        agent_id=state["agent_id"]
    )
    state["context"] = memories
    return state

# Add to your graph
workflow.add_node("memory", memory_node)