from flask import Flask, request, jsonify, abort

app = Flask(__name__)

# In-memory storage: {agent_id: {key: value}}
memory_store = {}

@app.route('/memory/<agent_id>', methods=['GET'])
def list_keys(agent_id):
    agent_memory = memory_store.get(agent_id, {})
    return jsonify(list(agent_memory.keys()))

@app.route('/memory/<agent_id>/<key>', methods=['GET'])
def get_value(agent_id, key):
    agent_memory = memory_store.get(agent_id, {})
    if key not in agent_memory:
        abort(404, description='Key not found')
    return jsonify({"value": agent_memory[key]})

@app.route('/memory/<agent_id>/<key>', methods=['POST'])
def set_value(agent_id, key):
    if not request.is_json:
        abort(400, description='Request body must be JSON')
    data = request.get_json()
    if "value" not in data:
        abort(400, description='JSON must contain "value" key')
    agent_memory = memory_store.setdefault(agent_id, {})
    agent_memory[key] = data["value"]
    return jsonify({"message": "Value set"}), 201

@app.route('/memory/<agent_id>/<key>', methods=['DELETE'])
def delete_value(agent_id, key):
    agent_memory = memory_store.get(agent_id, {})
    if key not in agent_memory:
        abort(404, description='Key not found')
    del agent_memory[key]
    return jsonify({"message": "Key deleted"})

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000)