import json
from brynq import decentralized_memory

def mcp_broker_dmem_store(args: str) -> str:
    """Store memory. Args: JSON containing content, type, tags (list), source."""
    try:
        data = json.loads(args)
        res = decentralized_memory.store_memory(data.get('content'), data.get('type', 'text'), data.get('tags'), data.get('source'))
        return json.dumps(res)
    except Exception as e:
        return json.dumps({"error": str(e)})

def mcp_broker_dmem_retrieve(args: str) -> str:
    """Retrieve memory. Args: JSON containing query, limit."""
    try:
        data = json.loads(args)
        res = decentralized_memory.retrieve_memory(data.get('query'), data.get('limit', 5))
        return json.dumps({"results": res})
    except Exception as e:
        return json.dumps({"error": str(e)})

TOOLS = [
    {
        "name": "broker_dmem_store",
        "description": "Store a decentralized memory. Args: JSON with content, type, tags, source",
        "inputSchema": {"type": "object", "properties": {"args": {"type": "string"}}},
        "handler": mcp_broker_dmem_store
    },
    {
        "name": "broker_dmem_retrieve",
        "description": "Retrieve memories by query. Args: JSON with query, limit",
        "inputSchema": {"type": "object", "properties": {"args": {"type": "string"}}},
        "handler": mcp_broker_dmem_retrieve
    }
]
