import asyncio
import json
import os
from datetime import datetime
from pathlib import Path
from typing import Any, Optional
from fnmatch import fnmatch

from fastmcp import FastMCP, Client
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 FilesystemBackend, LocalShellBackend, CompositeBackend

# ------------------- Memory Server -------------------
class MemoryServer:
    def __init__(self):
        self.mcp = FastMCP("Memory-Server")
        self.storage_path = Path("./memory_data.json")

    def _load_memory(self) -> dict:
        if not self.storage_path.exists():
            return {}
        with open(self.storage_path, 'r', encoding='utf-8') as f:
            return json.load(f)

    def _save_memory(self, data: dict):
        self.storage_path.parent.mkdir(parents=True, exist_ok=True)
        with open(self.storage_path, 'w', encoding='utf-8') as f:
            json.dump(data, f, indent=2, ensure_ascii=False)

    @self.mcp.tool()
    def save(self, key: str, value: Any) -> bool:
        data = self._load_memory()
        data[key] = {
            "value": value,
            "timestamp": datetime.utcnow().isoformat()
        }
        self._save_memory(data)
        return True

    @self.mcp.tool()
    def get(self, key: str) -> Optional[dict]:
        data = self._load_memory()
        if key in data:
            return {"key": key, "value": data[key]["value"], "timestamp": data[key]["timestamp"]}
        return None

    @self.mcp.tool()
    def delete(self, key: str) -> bool:
        data = self._load_memory()
        if key in data:
            del data[key]
            self._save_memory(data)
            return True
        return False

    @self.mcp.tool()
    def list_keys(self, pattern: str = "*") -> list[str]:
        data = self._load_memory()
        return [k for k in data.keys() if fnmatch(k, pattern)]

    @self.mcp.tool()
    def save_with_namespace(self, key: str, value: Any, namespace: str = "default") -> bool:
        full_key = f"{namespace}:{key}"
        return self.save(full_key, value)

    @self.mcp.tool()
    def get_by_namespace(self, namespace: str = "default") -> list[dict]:
        data = self._load_memory()
        result = []
        prefix = f"{namespace}:"
        for k, v in data.items():
            if k.startswith(prefix):
                result.append({"key": k, "value": v["value"], "timestamp": v["timestamp"]})
        return result

# ------------------- DeepAgent that uses Memory Server -------------------
# The agent will use a simple tool that calls the memory server via MCP client
@tool
def memory_tool(action: str, params: dict) -> str:
    """Calls the memory server tool specified by action.
    Parameters must include the tool name and its arguments.
    """
    client = Client("python memory_server.py")
    asyncio.run(client.connect())
    try:
        result = asyncio.run(client.call_tool(action, params))
        return str(result)
    finally:
        asyncio.run(client.close())

# LLM and backend setup
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(),
])

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

async def run_agent_demo():
    response = await agent.ainvoke(
        {"messages": [HumanMessage(content="Store key 'greeting' with value 'Hello' in default namespace.")]},
        {"configurable": {"thread_id": "demo-1"}},
    )
    print("Agent response:", response["messages"][-1].content)

# ------------------- Server run -------------------
if __name__ == "__main__":
    # Start server in background if called directly
    server = MemoryServer()
    # Run server in a separate thread so that the agent demo can also run
    import threading
    server_thread = threading.Thread(target=lambda: server.mcp.run(transport="stdio", show_banner=False, log_level="ERROR"), daemon=True)
    server_thread.start()
    # Give server a moment to start
    asyncio.run(asyncio.sleep(0.5))
    # Run agent demo
    asyncio.run(run_agent_demo())
    # Keep server alive for a short while to allow manual testing
    try:
        while True:
            pass
    except KeyboardInterrupt:
        print("Shutting down.")