import os
import asyncio
from typing import Any, Dict

from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage
from langchain.tools import tool
from deepagents import create_deep_agent
from deepagents.backends import CompositeBackend, LocalShellBackend, FilesystemBackend

# LLM via OpenRouter
llm = ChatOpenAI(
    model="openai/gpt-oss-20b:free",
    base_url="https://openrouter.ai/api/v1",
    api_key=os.getenv("OPENAI_API_KEY"),
    temperature=0.0,
)

backend = CompositeBackend(
    [
        LocalShellBackend(workspace_dir="./workspace"),
        FilesystemBackend(),
    ]
)

# Tool that forwards calls to the MCP memory server
@tool
def memory_action(action: str, params: Dict[str, Any]) -> str:
    """
    Perform a memory operation via the MCP server.

    action: one of save_with_namespace, get_by_namespace, list_keys, delete, get, save
    params: dictionary of parameters required by the chosen action.
    Returns a JSON string with the server response.
    """
    import json
    import asyncio
    from fastmcp import Client

    async def _call():
        client = Client("python memory_server.py")
        await client.connect()
        try:
            result = await client.call_tool(action, params)
            return json.dumps(result, ensure_ascii=False)
        finally:
            await client.close()

    loop = asyncio.get_event_loop()
    return loop.run_until_complete(_call())


agent = create_deep_agent(
    model=llm,
    tools=[memory_action],
    backend=backend,
    system_prompt="You are an assistant that can store and retrieve data using a remote memory server.",
)

async def demo():
    # Example: store a name and then retrieve it
    query = (
        "Save the user name 'Алексей' in the default namespace using the memory_action tool. "
        "Then read back all entries from the default namespace and return them."
    )
    result = await agent.ainvoke(
        {"messages": [HumanMessage(content=query)]},
        {"configurable": {"thread_id": "demo-1"}},
    )
    print(result["messages"][-1].content)


if __name__ == "__main__":
    asyncio.run(demo())