graph/graph_store.py

from __future__ import annotations

from datetime import datetime, timezone
from typing import Any, Dict, List

from shield.graph.relationships import build_relationship_edges
from shield.memory.entity_store import EncryptedEntityStore
from shield.utils.hashing import hash_entity


def _now() -> str:
    return datetime.now(timezone.utc).isoformat()


class GraphStore:
    def __init__(self, store: EncryptedEntityStore | None = None) -> None:
        self.store = store or EncryptedEntityStore()

    def lookup_entity(self, entity_type: str, value: str) -> Dict[str, Any]:
        entity_hash = hash_entity(entity_type, value)
        return self.store.read("graph", "entities", entity_type, f"{entity_hash}.json") or {}

    def find_linked_accounts(self, entities: Dict[str, str]) -> Dict[str, List[Dict[str, Any]]]:
        matches: Dict[str, List[Dict[str, Any]]] = {}
        for entity_type, value in entities.items():
            record = self.lookup_entity(entity_type, value)
            accounts = record.get("accounts", []) if isinstance(record, dict) else []
            if accounts:
                matches[entity_type] = accounts
        return matches

    def add_account_entities(
        self,
        *,
        customer_id: str | None,
        account_hash: str,
        entities: Dict[str, str],
    ) -> None:
        customer_hash = hash_entity("customer", customer_id or "anonymous")
        account_record = self.store.read("graph", "accounts", f"{account_hash}.json") or {
            "account_hash": account_hash,
            "first_seen": _now(),
            "last_seen": None,
            "edges": [],
        }
        account_record["last_seen"] = _now()

        for entity_type, value in entities.items():
            entity_hash = hash_entity(entity_type, value)
            entity_record = self.store.read("graph", "entities", entity_type, f"{entity_hash}.json") or {
                "entity_hash": entity_hash,
                "entity_type": entity_type,
                "first_seen": _now(),
                "last_seen": None,
                "accounts": [],
                "customer_hashes": [],
            }
            entity_record["last_seen"] = _now()
            account_entry = {"account_hash": account_hash, "customer_hash": customer_hash}
            if account_entry not in entity_record["accounts"]:
                entity_record["accounts"].append(account_entry)
            if customer_hash not in entity_record["customer_hashes"]:
                entity_record["customer_hashes"].append(customer_hash)
            self.store.write(entity_record, "graph", "entities", entity_type, f"{entity_hash}.json")

        for edge in build_relationship_edges(account_hash, entities):
            if edge not in account_record["edges"]:
                account_record["edges"].append(edge)
        self.store.write(account_record, "graph", "accounts", f"{account_hash}.json")