import json, os

class ContextGraph:
    def __init__(self, file_path="context_graph.json"):
        self.file_path = file_path
        self._init_graph()

    def _init_graph(self):
        if not os.path.exists(self.file_path):
            with open(self.file_path, "w") as f:
                json.dump({"context_links": {}}, f)

    def link_context(self, agent_id, category, key, value):
        with open(self.file_path, "r") as f:
            graph = json.load(f)

        if agent_id not in graph["context_links"]:
            graph["context_links"][agent_id] = {}
        if category not in graph["context_links"][agent_id]:
            graph["context_links"][agent_id][category] = {}

        graph["context_links"][agent_id][category][key] = value

        with open(self.file_path, "w") as f:
            json.dump(graph, f, indent=2)

    def query_context(self, agent_id, category=None, keyword=None):
        with open(self.file_path, "r") as f:
            graph = json.load(f)

        agent_data = graph["context_links"].get(agent_id, {})
        if category:
            data = agent_data.get(category, {})
        else:
            # Merge all categories
            data = {k: v for c in agent_data.values() for k, v in c.items()}

        results = [
            f"{k}: {v}" for k, v in data.items()
            if not keyword or keyword.lower() in k.lower() or keyword.lower() in v.lower()
        ]
        return results or ["No context found."]