ReMe/reme/components/file_graph/local_file_graph.py

"""Pure-Python file-graph backend (no external deps)."""

from pathlib import Path

from .base_file_graph import BaseFileGraph
from ..component_registry import R
from ...enumeration import LinkScopeEnum
from ...schema import FileLink, FileNode
from ...utils.jsonl_zst import read_jsonl_zst, write_jsonl_zst


@R.register("local")
class LocalFileGraph(BaseFileGraph):
    """Dict-backed file graph; uses FileLink.target_path for adjacency."""

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._nodes: dict[str, FileNode] = {}
        self._inverse: dict[str, set[str]] = {}  # real target → sources
        self._pending: dict[str, set[str]] = {}  # virtual target → sources
        self._graph_file: Path = self.component_metadata_path / f"{self.name}.jsonl.zst"

    # -- Lifecycle ---------------------------------------------------------

    async def _start(self) -> None:
        self.component_metadata_path.mkdir(parents=True, exist_ok=True)
        await super()._start()  # base calls load()
        await self.rebuild_links()

    async def load(self) -> None:
        if not self._graph_file.exists():
            return
        try:
            for line in read_jsonl_zst(self._graph_file):
                if line.strip():
                    node = FileNode.model_validate_json(line)
                    self._nodes[node.path] = node
            self.logger.debug(f"Loaded {len(self._nodes)} nodes from {self._graph_file}")
        except Exception as e:
            self.logger.exception(f"Failed to load {self._graph_file}: {e}")

    async def dump(self) -> None:
        try:
            write_jsonl_zst(self._graph_file, (n.model_dump_json() for n in self._nodes.values()))
            self.logger.info(f"Saved {len(self._nodes)} nodes to {self._graph_file}")
        except Exception as e:
            self.logger.exception(f"Failed to write {self._graph_file}: {e}")

    # -- Internals ---------------------------------------------------------

    @staticmethod
    def _targets(node: FileNode) -> list[str]:
        return [lnk.target_path for lnk in node.links if lnk.target_path]

    def _add_edge(self, src: str, target: str) -> None:
        bucket = self._inverse if target in self._nodes else self._pending
        bucket.setdefault(target, set()).add(src)

    def _remove_edge(self, src: str, target: str) -> None:
        for bucket in (self._inverse, self._pending):
            srcs = bucket.get(target)
            if srcs and src in srcs:
                srcs.discard(src)
                if not srcs:
                    del bucket[target]

    @staticmethod
    def _normalize_scope(scope: LinkScopeEnum | str) -> LinkScopeEnum:
        return scope if isinstance(scope, LinkScopeEnum) else LinkScopeEnum(scope)

    def _scope_match(self, target: str, scope: LinkScopeEnum | str) -> bool:
        scope = self._normalize_scope(scope)
        if scope is LinkScopeEnum.ALL:
            return True
        is_real = target in self._nodes
        return is_real if scope is LinkScopeEnum.REAL else not is_real

    # -- Node CRUD ---------------------------------------------------------

    async def upsert_nodes(self, nodes: list[FileNode]) -> None:
        for node in nodes:
            path = node.path
            old = self._nodes.get(path)
            if old is not None:
                for target in self._targets(old):
                    self._remove_edge(path, target)
            self._nodes[path] = node
            for target in self._targets(node):
                self._add_edge(path, target)
            promoted = self._pending.pop(path, None)
            if promoted:
                self._inverse.setdefault(path, set()).update(promoted)

    async def delete_nodes(self, paths: list[str]) -> None:
        for path in paths:
            node = self._nodes.pop(path, None)
            if node is None:
                continue
            for target in self._targets(node):
                self._remove_edge(path, target)
            demoted = self._inverse.pop(path, None)
            if demoted:
                self._pending.setdefault(path, set()).update(demoted)

    async def get_nodes(self, paths: list[str] | None = None) -> list[FileNode]:
        if paths is None:
            return list(self._nodes.values())
        return [self._nodes[p] for p in paths if p in self._nodes]

    async def rebuild_links(self) -> None:
        self._inverse.clear()
        self._pending.clear()
        for src, node in self._nodes.items():
            for target in self._targets(node):
                self._add_edge(src, target)

    async def clear(self):
        self._nodes.clear()
        self._inverse.clear()
        self._pending.clear()
        self._graph_file.unlink(missing_ok=True)

    # -- Link access -------------------------------------------------------

    async def get_outlinks(self, path: str, scope: LinkScopeEnum | str = LinkScopeEnum.REAL) -> list[FileLink]:
        scope = self._normalize_scope(scope)
        node = self._nodes.get(path)
        if node is None:
            return []
        return [lnk for lnk in node.links if lnk.target_path and self._scope_match(lnk.target_path, scope)]

    async def get_inlinks(self, path: str, scope: LinkScopeEnum | str = LinkScopeEnum.REAL) -> list[FileLink]:
        scope = self._normalize_scope(scope)
        sources: set[str] = set()
        if scope in (LinkScopeEnum.REAL, LinkScopeEnum.ALL):
            sources |= self._inverse.get(path, set())
        if scope in (LinkScopeEnum.VIRTUAL, LinkScopeEnum.ALL):
            sources |= self._pending.get(path, set())
        return [
            link
            for src in sorted(sources)
            if src in self._nodes
            for link in self._nodes[src].links
            if link.target_path == path
        ]