# backend/memory/semantic_memory.py
# Knowledge graph stored in MongoDB as an adjacency list using two collections: nodes and edges
import uuid
from backend.db.mongodb import MongoDBClient

class SemanticMemory:
    def __init__(self, db_path: str = None):
        # db_path is ignored now since we use MongoDB
        self.nodes = MongoDBClient.get_db().nodes
        self.edges = MongoDBClient.get_db().edges

    async def _get_or_create_node(self, label: str) -> str:
        label = label.lower()
        node = await self.nodes.find_one({'label': label})
        if node:
            return node['_id']
        
        node_id = str(uuid.uuid4())
        await self.nodes.insert_one({
            '_id': node_id,
            'label': label,
            'properties': {}
        })
        return node_id

    async def add_fact(self, subject: str, predicate: str, obj: str, confidence: float = 1.0):
        from_id = await self._get_or_create_node(subject)
        to_id = await self._get_or_create_node(obj)
        predicate = predicate.lower()
        
        # Upsert the edge based on from_id, to_id, and relation
        await self.edges.update_one(
            {
                'from_id': from_id,
                'to_id': to_id,
                'relation': predicate
            },
            {
                '$set': {'weight': confidence}
            },
            upsert=True
        )

    async def query_facts(self, subject: str) -> list[dict]:
        subject_lower = subject.lower()
        node = await self.nodes.find_one({'label': subject_lower})
        if not node:
            return []
            
        subject_id = node['_id']
        results = []
        
        # Outgoing edges (where subject is the source)
        async for edge in self.edges.find({'from_id': subject_id}):
            obj_node = await self.nodes.find_one({'_id': edge['to_id']})
            if obj_node:
                results.append({
                    "subject": subject_lower,
                    "predicate": edge['relation'],
                    "object": obj_node['label'],
                    "confidence": edge.get('weight', 1.0)
                })
                
        # Incoming edges (where subject is the target)
        async for edge in self.edges.find({'to_id': subject_id}):
            subj_node = await self.nodes.find_one({'_id': edge['from_id']})
            if subj_node:
                results.append({
                    "subject": subj_node['label'],
                    "predicate": edge['relation'],
                    "object": subject_lower,
                    "confidence": edge.get('weight', 1.0)
                })
                
        return results