helix_state.py

# -*- coding: utf-8 -*-
"""
helix_state.py — In-Process Graph State Manager
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
Implements the same graph API that HelixDB would expose via Python bindings,
but runs entirely in-process (zero network, zero daemon, ~0 MB overhead).

Why not Redis?
  Redis requires a separate daemon process (~200 MB RAM) and TCP round-trips.
  This module stores the same data as a Python dict-of-dicts with O(1) node
  lookup and O(k) edge traversal where k = number of edges per node.

Why not flat dicts in backend.py?
  A graph model lets us express relationships that flat dicts cannot:
    Project → HAS_TASK → Task
    Task    → USES_FILE → File
    File    → HAD_BUG  → Bug
    Bug     → FIXED_BY → Snippet (in Second Brain)
  This powers the Bell Curve apex prompt: we can query "What file is this
  task working on?" and load exactly that file — nothing more.

Drop-in swap: When HelixDB ships stable Python bindings, replace this
file with:  from helixdb import HelixDB as HelixStateDB
The public API (add_node, add_edge, get_node, get_neighbors, update_node,
remove_node, query_path) is kept identical to the planned HelixDB SDK.
"""

import time
import logging
import threading
from typing import Any, Dict, List, Optional, Tuple

logger = logging.getLogger("helix_state")


class HelixStateDB:
    """
    In-process graph database.

    Graph model:
      Nodes: { node_type: { node_id: { **properties } } }
      Edges: { (src_type, src_id, edge_label, dst_type, dst_id): { **properties } }

    Thread-safe for concurrent FastAPI request handlers.
    """

    def __init__(self):
        self._nodes: Dict[str, Dict[str, Dict[str, Any]]] = {}
        self._edges: Dict[Tuple, Dict[str, Any]] = {}
        self._lock = threading.RLock()
        logger.info("[HelixState] In-process graph database initialised.")

    # ── Node Operations ───────────────────────────────────────────────────────

    def add_node(self, node_type: str, node_id: str, **props) -> bool:
        """Insert or replace a node. Returns True on success."""
        with self._lock:
            if node_type not in self._nodes:
                self._nodes[node_type] = {}
            self._nodes[node_type][node_id] = {
                **props,
                "_created_at": time.time(),
                "_updated_at": time.time(),
            }
            logger.debug("[HelixState] add_node(%s, %s)", node_type, node_id)
            return True

    def update_node(self, node_type: str, node_id: str, **props) -> bool:
        """Merge props into an existing node. Returns False if node not found."""
        with self._lock:
            node = self._nodes.get(node_type, {}).get(node_id)
            if node is None:
                return False
            node.update(props)
            node["_updated_at"] = time.time()
            logger.debug("[HelixState] update_node(%s, %s)", node_type, node_id)
            return True

    def get_node(self, node_type: str, node_id: str) -> Optional[Dict[str, Any]]:
        """Return a node's property dict, or None."""
        with self._lock:
            return self._nodes.get(node_type, {}).get(node_id)

    def remove_node(self, node_type: str, node_id: str) -> bool:
        """Remove a node and all its edges."""
        with self._lock:
            if node_id not in self._nodes.get(node_type, {}):
                return False
            del self._nodes[node_type][node_id]
            # Prune orphaned edges
            dead = [k for k in self._edges
                    if (k[0] == node_type and k[1] == node_id) or
                       (k[3] == node_type and k[4] == node_id)]
            for k in dead:
                del self._edges[k]
            logger.debug("[HelixState] remove_node(%s, %s) + %d edges", node_type, node_id, len(dead))
            return True

    def list_nodes(self, node_type: str) -> List[str]:
        """Return all node IDs of a given type."""
        with self._lock:
            return list(self._nodes.get(node_type, {}).keys())

    # ── Edge Operations ───────────────────────────────────────────────────────

    def add_edge(
        self,
        src_type: str, src_id: str,
        dst_type: str, dst_id: str,
        label: str,
        **props,
    ) -> bool:
        """Add a directed edge (src)-[label]->(dst). Overwrites if exists."""
        with self._lock:
            key = (src_type, src_id, label, dst_type, dst_id)
            self._edges[key] = {**props, "_created_at": time.time()}
            logger.debug("[HelixState] add_edge %s:%s -[%s]-> %s:%s", src_type, src_id, label, dst_type, dst_id)
            return True

    def get_neighbors(
        self,
        src_type: str,
        src_id: str,
        label: str,
        dst_type: Optional[str] = None,
    ) -> List[Dict[str, Any]]:
        """
        Return list of destination node property dicts reachable from
        (src_type, src_id) via edges with the given label.
        Optionally filter by dst_type.
        """
        with self._lock:
            results = []
            for key, edge_props in self._edges.items():
                s_type, s_id, e_label, d_type, d_id = key
                if s_type != src_type or s_id != src_id or e_label != label:
                    continue
                if dst_type and d_type != dst_type:
                    continue
                node = self._nodes.get(d_type, {}).get(d_id)
                if node:
                    results.append({"_type": d_type, "_id": d_id, **node})
            return results

    def remove_edge(
        self,
        src_type: str, src_id: str,
        dst_type: str, dst_id: str,
        label: str,
    ) -> bool:
        """Remove a specific directed edge."""
        with self._lock:
            key = (src_type, src_id, label, dst_type, dst_id)
            if key in self._edges:
                del self._edges[key]
                return True
            return False

    # ── Query Helpers ─────────────────────────────────────────────────────────

    def query_path(
        self,
        start_type: str, start_id: str,
        *edge_labels: str,
    ) -> List[Dict[str, Any]]:
        """
        Traverse a chain of edges and return the terminal nodes.
        Example: query_path("project", "calc_v1", "HAS_TASK", "USES_FILE")
        Returns all File nodes reachable via the two-hop path.
        """
        current: List[Dict] = [{"_type": start_type, "_id": start_id}]
        for label in edge_labels:
            next_level = []
            for node in current:
                ntype, nid = node["_type"], node["_id"]
                neighbors = self.get_neighbors(ntype, nid, label)
                next_level.extend(neighbors)
            current = next_level
        return current

    def dump(self) -> Dict[str, Any]:
        """Serialise the full graph to a JSON-compatible dict (for /api/metrics)."""
        with self._lock:
            return {
                "node_counts": {t: len(ids) for t, ids in self._nodes.items()},
                "edge_count":  len(self._edges),
                "nodes":       {t: dict(ids) for t, ids in self._nodes.items()},
            }

    # ── Project State Helpers (Eternity Loop convenience) ────────────────────

    def upsert_project(self, name: str, goal: str, mode: str, priority: str):
        """Convenience: add or update a project node."""
        if not self.get_node("project", name):
            self.add_node("project", name, goal=goal, mode=mode, priority=priority, cycle=0)
        else:
            self.update_node("project", name, mode=mode, priority=priority)

    def record_cycle(self, project_name: str, summary: str, status: str):
        """Increment cycle counter and store last summary on the project node."""
        node = self.get_node("project", project_name)
        if node:
            cycle = node.get("cycle", 0) + 1
            self.update_node("project", project_name,
                             cycle=cycle,
                             last_summary=summary,
                             last_status=status,
                             last_cycle_at=time.time())

    def get_active_project_names(self) -> List[str]:
        """Return IDs of all project nodes where is_active == True."""
        with self._lock:
            return [
                pid for pid, props in self._nodes.get("project", {}).items()
                if props.get("is_active", True)
            ]

    def link_task_to_file(self, project_name: str, task_id: str, file_path: str):
        """Record which file a task is working on, for targeted brain loading."""
        self.add_node("task", task_id, project=project_name, file=file_path)
        self.add_edge("project", project_name, "task", task_id, "HAS_TASK")
        if file_path:
            self.add_node("file", file_path)
            self.add_edge("task", task_id, "file", file_path, "USES_FILE")

    def record_bug_fix(self, file_path: str, bug_summary: str, fix_summary: str, brain_path: str):
        """Record that a bug in a file was fixed and persisted to the Second Brain."""
        bug_id = f"bug_{int(time.time())}"
        self.add_node("bug", bug_id, file=file_path, summary=bug_summary)
        self.add_node("fix", bug_id, summary=fix_summary, brain_path=brain_path)
        self.add_edge("file", file_path, "bug",  bug_id, "HAD_BUG")
        self.add_edge("bug",  bug_id,    "fix",  bug_id, "FIXED_BY")
        self.add_edge("fix",  bug_id,    "brain", brain_path, "PERSISTED_TO")


# Singleton — import and use directly
helix_db = HelixStateDB()