from langgraph.store.mongodb.base import MongoDBStore, VectorIndexConfig
from langchain_voyageai import VoyageAIEmbeddings

client = MongoClient(MONGODB_URI)
db = client["memories"]
collection = db["memory_store"]

store = MongoDBStore(
  collection=collection,
  index_config=VectorIndexConfig(
    fields=None,
    filters=None,
    dims=1536,
    embed=VoyageAIEmbeddings(model="voyage-3.5")
  ),
  auto_index_timeout=70
)

from langgraph.checkpoint.mongodb import MongoDBSaver

checkpointer = MongoDBSaver(
  client,
  db_name="memories",
  collection_name="thread_checkpoints"
)

from langgraph.prebuilt import create_react_agent
from langmem import create_manage_memory_tool

def prompt(state, store):
  memories = store.search(
    ("memories",),
    query=state["messages"][-1].content,
  )
  system_msg = f"""
  You are a shopping assistant with persistent memory.
  ## Relevant Memories
  <memories>{memories}</memories>
  Use these memories to provide personalized responses.
  """
  return [{"role": "system", "content": system_msg}, *state["messages"]]

agent = create_react_agent(
  "openai:gpt-4o",
  prompt=lambda state: prompt(state, store),
  tools=[
    create_manage_memory_tool(namespace=("memories",)),
    search_products # Công cụ tìm kiếm sản phẩm riêng biệt
  ],
  store=store,
  checkpointer=checkpointer
)

@tool
def search_products(query: str) -> str:
  vectorstore = MongoDBAtlasVectorSearch(
    collection,
    OpenAIEmbeddings(),
    text_key="title",
    embedding_key="embedding",
    index_name="vector_index"
  )
  docs = vectorstore.similarity_search(query, k=5)
  return "\n".join([str(doc.metadata) for doc in docs])

def create_shopping_agent():
  store = MongoDBStore(
    collection=client.memories.user_preferences,
    index_config=VectorIndexConfig(
      dims=1536,
      embed=VoyageAIEmbeddings(model="voyage-3.5"),
      fields=["content"],
      filters=["active"]
    )
  )
  checkpointer = MongoDBSaver(client, "shopping_assistant", "conversations")

  def enhanced_prompt(state, store):
    user_query = state["messages"][-1].content
    memories = store.search(("preferences",), query=user_query, limit=3, filter={"active": True})
    purchase_history = store.search(("purchases",), query=user_query, limit=2)
    system_msg = f"""
    Bạn là trợ lý mua sắm chuyên nghiệp có khả năng truy cập:
    - Tìm kiếm sản phẩm
    - Bộ nhớ sở thích người dùng
    - Lịch sử mua hàng gần đây
    ## Sở thích người dùng
    {memories}
    ## Lịch sử mua sắm gần đây
    {purchase_history}
    Hãy cung cấp lời khuyên mua sắm cá nhân hóa, hữu ích.
    """
    return [{"role": "system", "content": system_msg}, *state["messages"]]

  return create_react_agent(
    "openai:gpt-4o",
    prompt=lambda state: enhanced_prompt(state, store),
    tools=[
      create_manage_memory_tool(namespace=("preferences",)),
      create_manage_memory_tool(namespace=("purchases",)),
      search_products
    ],
    store=store,
    checkpointer=checkpointer
  )

def create_user_scoped_agent(user_id: str):
  user_namespace = ("users", user_id, "memories")
  store = MongoDBStore(
    collection=collection,
    index_config=VectorIndexConfig(
      dims=1536,
      embed=OpenAIEmbeddings(),
      filters={"metadata.user_id": user_id}
    )
  )
  memory_tool = create_manage_memory_tool(
    namespace=user_namespace,
    metadata_filter={"user_id": user_id}
  )
  return create_react_agent(
    "openai:gpt-4o",
    tools=[memory_tool],
    store=store,
    checkpointer=MongoDBSaver(client, db_name="multi_tenant_memories", collection_name=f"user_{user_id}_checkpoints")
  )