import os
import asyncio
from typing import Any

from fastmcp import Client
from deepagents import create_deep_agent
from deepagents.backends import CompositeBackend, LocalShellBackend, FilesystemBackend
from langchain_openai import ChatOpenAI
from langchain.tools import tool


# LLM configuration (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
def memory_save(key: str, value: Any, namespace: str = "default") -> bool:
    """Save a value in the remote memory server."""
    async def _call():
        client = Client("python memory_server.py")
        await client.connect()
        try:
            result = await client.call_tool(
                "save_with_namespace",
                {"key": key, "value": value, "namespace": namespace},
            )
            return result
        finally:
            await client.close()
    return asyncio.run(_call())


@tool
def memory_get(namespace: str = "default") -> str:
    """Retrieve all key-value pairs from a namespace as a formatted string."""
    async def _call():
        client = Client("python memory_server.py")
        await client.connect()
        try:
            data = await client.call_tool("get_by_namespace", {"namespace": namespace})
            if not data:
                return "No data."
            lines = [f"{item['key']}: {item['value']}" for item in data]
            return "\\n".join(lines)
        finally:
            await client.close()
    return asyncio.run(_call())


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


async def demo():
    # Store a fact
    await agent.ainvoke(
        {"messages": [{"role": "user", "content": "Запомни, что мой любимый цвет - синий."}]},
        {"configurable": {"thread_id": "demo-1"}},
    )
    # Retrieve stored facts
    result = await agent.ainvoke(
        {"messages": [{"role": "user", "content": "Что я просил запомнить?"}]},
        {"configurable": {"thread_id": "demo-1"}},
    )
    print(result["messages"][-1].content)


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