memory.py

"""
Cogni-Engine v1 — TiDB Persistence Layer
Handles all database communication: schema, CRUD, buffering, sync.
The "long-term memory" — survives reboots, crashes, rebuilds.
"""

import time
import json
import threading
import traceback
from typing import List, Dict, Optional, Any, Tuple
from contextlib import contextmanager

import pymysql
from pymysql.cursors import DictCursor

import config
import utils


# ═══════════════════════════════════════════════════════════
# CONNECTION MANAGER
# ═══════════════════════════════════════════════════════════

class ConnectionPool:
    """
    Simple connection pool for TiDB.
    Manages multiple reusable connections with auto-reconnect.
    """

    def __init__(self):
        self._connections: List[pymysql.Connection] = []
        self._lock = threading.Lock()
        self._available: List[pymysql.Connection] = []
        self._max_size = config.TIDB_POOL_SIZE
        self._initialized = False

    def _create_connection(self) -> pymysql.Connection:
        """Create a new TiDB connection."""
        connect_kwargs = {
            "host": config.TIDB_HOST,
            "port": config.TIDB_PORT,
            "user": config.TIDB_USER,
            "password": config.TIDB_PASSWORD,
            "database": config.TIDB_DATABASE,
            "connect_timeout": config.TIDB_CONNECT_TIMEOUT,
            "read_timeout": config.TIDB_READ_TIMEOUT,
            "write_timeout": config.TIDB_WRITE_TIMEOUT,
            "charset": "utf8mb4",
            "cursorclass": DictCursor,
            "autocommit": True
        }

        if config.TIDB_SSL:
            connect_kwargs["ssl"] = {"ssl_mode": "VERIFY_IDENTITY"}

        conn = pymysql.connect(**connect_kwargs)
        return conn

    def _test_connection(self, conn: pymysql.Connection) -> bool:
        """Test if connection is still alive."""
        try:
            conn.ping(reconnect=False)
            return True
        except Exception:
            return False

    def acquire(self) -> pymysql.Connection:
        """Get a connection from the pool."""
        with self._lock:
            # Try to reuse existing connection
            while self._available:
                conn = self._available.pop()
                if self._test_connection(conn):
                    return conn
                else:
                    # Dead connection, discard
                    try:
                        conn.close()
                    except Exception:
                        pass

            # Create new if under limit
            if len(self._connections) < self._max_size:
                conn = self._create_connection()
                self._connections.append(conn)
                return conn

        # All connections busy, create temporary one
        return self._create_connection()

    def release(self, conn: pymysql.Connection):
        """Return a connection to the pool."""
        with self._lock:
            if self._test_connection(conn):
                self._available.append(conn)
            else:
                try:
                    conn.close()
                except Exception:
                    pass
                # Remove from tracked connections
                if conn in self._connections:
                    self._connections.remove(conn)

    def close_all(self):
        """Close all connections."""
        with self._lock:
            for conn in self._connections:
                try:
                    conn.close()
                except Exception:
                    pass
            self._connections.clear()
            self._available.clear()

    @contextmanager
    def connection(self):
        """Context manager for auto acquire/release."""
        conn = None
        try:
            conn = self.acquire()
            yield conn
        finally:
            if conn:
                self.release(conn)


# ═══════════════════════════════════════════════════════════
# WRITE BUFFER
# ═══════════════════════════════════════════════════════════

class WriteBuffer:
    """
    Buffers write operations and flushes in batches.
    Prevents excessive DB writes during rapid thinking cycles.
    """

    def __init__(self):
        self._lock = threading.Lock()
        self._node_upserts: Dict[str, dict] = {}    # id → node_data
        self._edge_upserts: Dict[str, dict] = {}    # id → edge_data
        self._chain_upserts: Dict[str, dict] = {}   # id → chain_data
        self._node_deletes: set = set()              # ids to delete
        self._edge_deletes: set = set()              # ids to delete
        self._state_update: Optional[dict] = None    # thinking state
        self._checksum_updates: Dict[str, str] = {}  # filename → checksum
        self._operation_count = 0
        self._last_flush_time = time.time()

    def buffer_node(self, node_data: dict):
        """Buffer a node upsert."""
        with self._lock:
            self._node_upserts[node_data["id"]] = node_data
            self._operation_count += 1

    def buffer_edge(self, edge_data: dict):
        """Buffer an edge upsert."""
        with self._lock:
            self._edge_upserts[edge_data["id"]] = edge_data
            self._operation_count += 1

    def buffer_chain(self, chain_data: dict):
        """Buffer a chain upsert."""
        with self._lock:
            self._chain_upserts[chain_data["id"]] = chain_data
            self._operation_count += 1

    def buffer_node_delete(self, node_id: str):
        """Buffer a node deletion."""
        with self._lock:
            self._node_deletes.add(node_id)
            self._node_upserts.pop(node_id, None)
            self._operation_count += 1

    def buffer_edge_delete(self, edge_id: str):
        """Buffer an edge deletion."""
        with self._lock:
            self._edge_deletes.add(edge_id)
            self._edge_upserts.pop(edge_id, None)
            self._operation_count += 1

    def buffer_state(self, state: dict):
        """Buffer thinking state update."""
        with self._lock:
            self._state_update = state

    def buffer_checksum(self, filename: str, checksum: str):
        """Buffer file checksum update."""
        with self._lock:
            self._checksum_updates[filename] = checksum

    @property
    def pending_count(self) -> int:
        """Number of pending operations."""
        return self._operation_count

    @property
    def seconds_since_flush(self) -> float:
        """Seconds elapsed since last flush."""
        return time.time() - self._last_flush_time

    def should_flush(self) -> bool:
        """Check if buffer should be flushed based on config thresholds."""
        if self._operation_count == 0 and self._state_update is None:
            return False
        if self._operation_count >= config.SYNC_INTERVAL_CYCLES:
            return True
        if self.seconds_since_flush >= config.SYNC_INTERVAL_SECONDS:
            return True
        return False

    def drain(self) -> dict:
        """
        Extract all buffered operations and reset buffer.
        Returns dict with all pending operations.
        """
        with self._lock:
            data = {
                "node_upserts": dict(self._node_upserts),
                "edge_upserts": dict(self._edge_upserts),
                "chain_upserts": dict(self._chain_upserts),
                "node_deletes": set(self._node_deletes),
                "edge_deletes": set(self._edge_deletes),
                "state_update": self._state_update,
                "checksum_updates": dict(self._checksum_updates)
            }
            self._node_upserts.clear()
            self._edge_upserts.clear()
            self._chain_upserts.clear()
            self._node_deletes.clear()
            self._edge_deletes.clear()
            self._state_update = None
            self._checksum_updates.clear()
            self._operation_count = 0
            self._last_flush_time = time.time()
            return data


# ═══════════════════════════════════════════════════════════
# MAIN MEMORY CLASS
# ═══════════════════════════════════════════════════════════

class Memory:
    """
    TiDB persistence layer.
    Handles schema creation, CRUD, buffered writes, and full state load/save.
    """

    def __init__(self):
        self.pool = ConnectionPool()
        self.buffer = WriteBuffer()
        self._connected = False
        self._schema_ready = False

    # ───────────────────────────────────────────────────
    # INITIALIZATION
    # ───────────────────────────────────────────────────

    def initialize(self) -> bool:
        """
        Initialize database: test connection and create schema.
        Returns True if successful.
        """
        if not config.TIDB_HOST:
            print("[MEMORY] No TiDB host configured. Running without persistence.")
            return False

        for attempt in range(config.TIDB_RETRY_ATTEMPTS):
            try:
                with self.pool.connection() as conn:
                    with conn.cursor() as cur:
                        cur.execute("SELECT 1")
                        result = cur.fetchone()
                        if result:
                            self._connected = True
                            print(f"[MEMORY] Connected to TiDB at {config.TIDB_HOST}")
                            self._create_schema()
                            return True
            except Exception as e:
                print(f"[MEMORY] Connection attempt {attempt + 1} failed: {e}")
                if attempt < config.TIDB_RETRY_ATTEMPTS - 1:
                    time.sleep(config.TIDB_RETRY_DELAY)

        print("[MEMORY] Failed to connect to TiDB after all retries.")
        return False

    def _create_schema(self):
        """Create all required tables if they don't exist."""
        schema_sql = [
            """
            CREATE TABLE IF NOT EXISTS nodes (
                id VARCHAR(32) PRIMARY KEY,
                type VARCHAR(32) NOT NULL,
                content TEXT NOT NULL,
                vector JSON,
                weight FLOAT DEFAULT 1.0,
                connections INT DEFAULT 0,
                source VARCHAR(16) DEFAULT 'data',
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
                INDEX idx_type (type),
                INDEX idx_weight (weight),
                INDEX idx_source (source)
            )
            """,
            """
            CREATE TABLE IF NOT EXISTS edges (
                id VARCHAR(32) PRIMARY KEY,
                from_node VARCHAR(32) NOT NULL,
                to_node VARCHAR(32) NOT NULL,
                relation VARCHAR(32) NOT NULL,
                weight FLOAT DEFAULT 1.0,
                confidence FLOAT DEFAULT 1.0,
                source VARCHAR(16) DEFAULT 'data',
                used_count INT DEFAULT 0,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
                INDEX idx_from (from_node),
                INDEX idx_to (to_node),
                INDEX idx_relation (relation),
                INDEX idx_weight (weight),
                INDEX idx_confidence (confidence)
            )
            """,
            """
            CREATE TABLE IF NOT EXISTS reasoning_chains (
                id VARCHAR(32) PRIMARY KEY,
                path JSON NOT NULL,
                conclusion TEXT,
                confidence FLOAT DEFAULT 0.5,
                used_count INT DEFAULT 0,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                INDEX idx_confidence (confidence),
                INDEX idx_used (used_count)
            )
            """,
            """
            CREATE TABLE IF NOT EXISTS thinking_state (
                id INT PRIMARY KEY DEFAULT 1,
                current_cycle BIGINT DEFAULT 0,
                total_cycles BIGINT DEFAULT 0,
                cursor_position VARCHAR(64) DEFAULT '',
                phase VARCHAR(32) DEFAULT 'init',
                metrics JSON,
                started_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP
            )
            """,
            """
            CREATE TABLE IF NOT EXISTS file_checksums (
                filename VARCHAR(255) PRIMARY KEY,
                checksum VARCHAR(64) NOT NULL,
                processed_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                line_count INT DEFAULT 0
            )
            """,
            """
            CREATE TABLE IF NOT EXISTS config_store (
                k VARCHAR(64) PRIMARY KEY,
                v TEXT NOT NULL,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP
            )
            """
        ]

        try:
            with self.pool.connection() as conn:
                with conn.cursor() as cur:
                    for sql in schema_sql:
                        cur.execute(sql)

                    # Initialize thinking_state if empty
                    cur.execute("SELECT COUNT(*) as cnt FROM thinking_state")
                    row = cur.fetchone()
                    if row["cnt"] == 0:
                        cur.execute("""
                            INSERT INTO thinking_state 
                            (id, current_cycle, total_cycles, cursor_position, phase, metrics)
                            VALUES (1, 0, 0, '', 'init', '{}')
                        """)

            self._schema_ready = True
            print("[MEMORY] Schema ready.")
        except Exception as e:
            print(f"[MEMORY] Schema creation failed: {e}")
            traceback.print_exc()

    @property
    def is_connected(self) -> bool:
        return self._connected and self._schema_ready

    # ───────────────────────────────────────────────────
    # EXECUTE HELPERS
    # ───────────────────────────────────────────────────

    def _execute(self, sql: str, params: tuple = None, fetch: str = "none") -> Any:
        """
        Execute SQL with auto-retry.
        fetch: "none", "one", "all"
        """
        if not self._connected:
            return None

        for attempt in range(config.TIDB_RETRY_ATTEMPTS):
            try:
                with self.pool.connection() as conn:
                    with conn.cursor() as cur:
                        cur.execute(sql, params)
                        if fetch == "one":
                            return cur.fetchone()
                        elif fetch == "all":
                            return cur.fetchall()
                        return True
            except pymysql.err.OperationalError as e:
                if attempt < config.TIDB_RETRY_ATTEMPTS - 1:
                    time.sleep(config.TIDB_RETRY_DELAY)
                else:
                    print(f"[MEMORY] Execute failed after retries: {e}")
                    return None
            except Exception as e:
                print(f"[MEMORY] Execute error: {e}")
                return None

    def _execute_many(self, sql: str, params_list: List[tuple]) -> bool:
        """Execute SQL for multiple parameter sets (batch insert/update)."""
        if not self._connected or not params_list:
            return False

        for attempt in range(config.TIDB_RETRY_ATTEMPTS):
            try:
                with self.pool.connection() as conn:
                    with conn.cursor() as cur:
                        cur.executemany(sql, params_list)
                return True
            except pymysql.err.OperationalError as e:
                if attempt < config.TIDB_RETRY_ATTEMPTS - 1:
                    time.sleep(config.TIDB_RETRY_DELAY)
                else:
                    print(f"[MEMORY] ExecuteMany failed: {e}")
                    return False
            except Exception as e:
                print(f"[MEMORY] ExecuteMany error: {e}")
                return False

    # ───────────────────────────────────────────────────
    # NODE OPERATIONS
    # ───────────────────────────────────────────────────

    def save_node(self, node: dict):
        """Buffer a node for batch writing."""
        self.buffer.buffer_node(node)

    def save_nodes_immediate(self, nodes: List[dict]) -> bool:
        """Write nodes directly to DB (bypass buffer). Used for bulk import."""
        if not nodes:
            return True

        sql = """
            INSERT INTO nodes (id, type, content, vector, weight, connections, source)
            VALUES (%s, %s, %s, %s, %s, %s, %s)
            ON DUPLICATE KEY UPDATE
                content = VALUES(content),
                vector = VALUES(vector),
                weight = VALUES(weight),
                connections = VALUES(connections),
                source = VALUES(source),
                updated_at = CURRENT_TIMESTAMP
        """
        params = [
            (
                n["id"],
                n.get("type", "fact"),
                n.get("content", ""),
                json.dumps(n.get("vector", [])),
                n.get("weight", 1.0),
                n.get("connections", 0),
                n.get("source", "data")
            )
            for n in nodes
        ]
        return self._execute_many(sql, params)

    def load_all_nodes(self) -> List[dict]:
        """Load all nodes from DB. Used at startup."""
        rows = self._execute(
            "SELECT id, type, content, vector, weight, connections, source, "
            "created_at, updated_at FROM nodes",
            fetch="all"
        )
        if not rows:
            return []

        nodes = []
        for row in rows:
            vector_data = row.get("vector")
            if isinstance(vector_data, str):
                vector_data = json.loads(vector_data)
            nodes.append({
                "id": row["id"],
                "type": row["type"],
                "content": row["content"],
                "vector": vector_data if vector_data else [],
                "weight": float(row["weight"]),
                "connections": int(row["connections"]),
                "source": row["source"],
                "created_at": str(row["created_at"]) if row.get("created_at") else "",
                "updated_at": str(row["updated_at"]) if row.get("updated_at") else ""
            })
        return nodes

    def delete_node(self, node_id: str):
        """Buffer a node deletion."""
        self.buffer.buffer_node_delete(node_id)

    def get_node_count(self) -> int:
        """Get total node count from DB."""
        row = self._execute("SELECT COUNT(*) as cnt FROM nodes", fetch="one")
        return row["cnt"] if row else 0

    # ───────────────────────────────────────────────────
    # EDGE OPERATIONS
    # ───────────────────────────────────────────────────

    def save_edge(self, edge: dict):
        """Buffer an edge for batch writing."""
        self.buffer.buffer_edge(edge)

    def save_edges_immediate(self, edges: List[dict]) -> bool:
        """Write edges directly to DB (bypass buffer)."""
        if not edges:
            return True

        sql = """
            INSERT INTO edges (id, from_node, to_node, relation, weight, confidence, source, used_count)
            VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
            ON DUPLICATE KEY UPDATE
                weight = VALUES(weight),
                confidence = VALUES(confidence),
                used_count = VALUES(used_count),
                updated_at = CURRENT_TIMESTAMP
        """
        params = [
            (
                e["id"],
                e["from_node"],
                e["to_node"],
                e.get("relation", "related_to"),
                e.get("weight", 1.0),
                e.get("confidence", 1.0),
                e.get("source", "data"),
                e.get("used_count", 0)
            )
            for e in edges
        ]
        return self._execute_many(sql, params)

    def load_all_edges(self) -> List[dict]:
        """Load all edges from DB. Used at startup."""
        rows = self._execute(
            "SELECT id, from_node, to_node, relation, weight, confidence, "
            "source, used_count, created_at FROM edges",
            fetch="all"
        )
        if not rows:
            return []

        return [
            {
                "id": row["id"],
                "from_node": row["from_node"],
                "to_node": row["to_node"],
                "relation": row["relation"],
                "weight": float(row["weight"]),
                "confidence": float(row["confidence"]),
                "source": row["source"],
                "used_count": int(row["used_count"]),
                "created_at": str(row["created_at"]) if row.get("created_at") else ""
            }
            for row in rows
        ]

    def delete_edge(self, edge_id: str):
        """Buffer an edge deletion."""
        self.buffer.buffer_edge_delete(edge_id)

    def get_edge_count(self) -> int:
        """Get total edge count from DB."""
        row = self._execute("SELECT COUNT(*) as cnt FROM edges", fetch="one")
        return row["cnt"] if row else 0

    # ───────────────────────────────────────────────────
    # CHAIN OPERATIONS
    # ───────────────────────────────────────────────────

    def save_chain(self, chain: dict):
        """Buffer a reasoning chain for batch writing."""
        self.buffer.buffer_chain(chain)

    def load_top_chains(self, limit: int = 10000) -> List[dict]:
        """Load top-scored reasoning chains from DB."""
        rows = self._execute(
            "SELECT id, path, conclusion, confidence, used_count, created_at "
            "FROM reasoning_chains ORDER BY confidence DESC, used_count DESC LIMIT %s",
            (limit,),
            fetch="all"
        )
        if not rows:
            return []

        chains = []
        for row in rows:
            path_data = row.get("path")
            if isinstance(path_data, str):
                path_data = json.loads(path_data)
            chains.append({
                "id": row["id"],
                "path": path_data if path_data else [],
                "conclusion": row.get("conclusion", ""),
                "confidence": float(row["confidence"]),
                "used_count": int(row["used_count"]),
                "created_at": str(row["created_at"]) if row.get("created_at") else ""
            })
        return chains

    def get_chain_count(self) -> int:
        """Get total chain count."""
        row = self._execute("SELECT COUNT(*) as cnt FROM reasoning_chains", fetch="one")
        return row["cnt"] if row else 0

    def prune_weak_chains(self, min_confidence: float = 0.2, max_age_days: int = 30) -> int:
        """Delete chains with low confidence and old age. Returns count deleted."""
        result = self._execute(
            "DELETE FROM reasoning_chains WHERE confidence < %s "
            "AND used_count = 0 AND created_at < DATE_SUB(NOW(), INTERVAL %s DAY)",
            (min_confidence, max_age_days)
        )
        return 0  # executemany doesn't return rowcount easily

    # ───────────────────────────────────────────────────
    # THINKING STATE
    # ───────────────────────────────────────────────────

    def save_thinking_state(self, state: dict):
        """Buffer thinking state update."""
        self.buffer.buffer_state(state)

    def load_thinking_state(self) -> dict:
        """Load thinking state from DB."""
        row = self._execute(
            "SELECT current_cycle, total_cycles, cursor_position, phase, "
            "metrics, started_at, updated_at FROM thinking_state WHERE id = 1",
            fetch="one"
        )
        if not row:
            return {
                "current_cycle": 0,
                "total_cycles": 0,
                "cursor_position": "",
                "phase": "init",
                "metrics": {},
                "started_at": utils.timestamp_now(),
                "updated_at": utils.timestamp_now()
            }

        metrics = row.get("metrics")
        if isinstance(metrics, str):
            try:
                metrics = json.loads(metrics)
            except (json.JSONDecodeError, TypeError):
                metrics = {}

        return {
            "current_cycle": int(row.get("current_cycle", 0)),
            "total_cycles": int(row.get("total_cycles", 0)),
            "cursor_position": row.get("cursor_position", ""),
            "phase": row.get("phase", "init"),
            "metrics": metrics if metrics else {},
            "started_at": str(row["started_at"]) if row.get("started_at") else "",
            "updated_at": str(row["updated_at"]) if row.get("updated_at") else ""
        }

    # ───────────────────────────────────────────────────
    # FILE CHECKSUMS
    # ───────────────────────────────────────────────────

    def save_file_checksum(self, filename: str, checksum: str, line_count: int = 0):
        """Buffer file checksum update."""
        self.buffer.buffer_checksum(filename, checksum)

    def load_file_checksums(self) -> Dict[str, str]:
        """Load all file checksums. Returns {filename: checksum}."""
        rows = self._execute(
            "SELECT filename, checksum FROM file_checksums",
            fetch="all"
        )
        if not rows:
            return {}
        return {row["filename"]: row["checksum"] for row in rows}

    # ───────────────────────────────────────────────────
    # CONFIG STORE
    # ───────────────────────────────────────────────────

    def save_config(self, key: str, value: str):
        """Save a config key-value pair."""
        self._execute(
            "INSERT INTO config_store (k, v) VALUES (%s, %s) "
            "ON DUPLICATE KEY UPDATE v = VALUES(v), updated_at = CURRENT_TIMESTAMP",
            (key, value)
        )

    def load_config(self, key: str, default: str = "") -> str:
        """Load a config value."""
        row = self._execute(
            "SELECT v FROM config_store WHERE k = %s",
            (key,),
            fetch="one"
        )
        return row["v"] if row else default

    # ───────────────────────────────────────────────────
    # FLUSH (Buffer → DB)
    # ───────────────────────────────────────────────────

    def flush(self) -> dict:
        """
        Flush all buffered operations to TiDB.
        Returns summary of what was flushed.
        """
        if not self._connected:
            return {"status": "not_connected", "flushed": 0}

        data = self.buffer.drain()
        summary = {
            "nodes_upserted": 0,
            "edges_upserted": 0,
            "chains_upserted": 0,
            "nodes_deleted": 0,
            "edges_deleted": 0,
            "state_updated": False,
            "checksums_updated": 0
        }

        try:
            # ── Upsert nodes ──
            if data["node_upserts"]:
                nodes = list(data["node_upserts"].values())
                if self.save_nodes_immediate(nodes):
                    summary["nodes_upserted"] = len(nodes)

            # ── Upsert edges ──
            if data["edge_upserts"]:
                edges = list(data["edge_upserts"].values())
                if self.save_edges_immediate(edges):
                    summary["edges_upserted"] = len(edges)

            # ── Upsert chains ──
            if data["chain_upserts"]:
                chains = list(data["chain_upserts"].values())
                sql = """
                    INSERT INTO reasoning_chains (id, path, conclusion, confidence, used_count)
                    VALUES (%s, %s, %s, %s, %s)
                    ON DUPLICATE KEY UPDATE
                        confidence = VALUES(confidence),
                        used_count = VALUES(used_count)
                """
                params = [
                    (
                        c["id"],
                        json.dumps(c.get("path", [])),
                        c.get("conclusion", ""),
                        c.get("confidence", 0.5),
                        c.get("used_count", 0)
                    )
                    for c in chains
                ]
                if self._execute_many(sql, params):
                    summary["chains_upserted"] = len(chains)

            # ── Delete nodes ──
            if data["node_deletes"]:
                for node_id in data["node_deletes"]:
                    self._execute("DELETE FROM nodes WHERE id = %s", (node_id,))
                    # Also delete connected edges
                    self._execute(
                        "DELETE FROM edges WHERE from_node = %s OR to_node = %s",
                        (node_id, node_id)
                    )
                summary["nodes_deleted"] = len(data["node_deletes"])

            # ── Delete edges ──
            if data["edge_deletes"]:
                for edge_id in data["edge_deletes"]:
                    self._execute("DELETE FROM edges WHERE id = %s", (edge_id,))
                summary["edges_deleted"] = len(data["edge_deletes"])

            # ── Update thinking state ──
            if data["state_update"]:
                state = data["state_update"]
                self._execute(
                    """
                    UPDATE thinking_state SET
                        current_cycle = %s,
                        total_cycles = %s,
                        cursor_position = %s,
                        phase = %s,
                        metrics = %s,
                        updated_at = CURRENT_TIMESTAMP
                    WHERE id = 1
                    """,
                    (
                        state.get("current_cycle", 0),
                        state.get("total_cycles", 0),
                        state.get("cursor_position", ""),
                        state.get("phase", ""),
                        json.dumps(state.get("metrics", {}))
                    )
                )
                summary["state_updated"] = True

            # ── Update file checksums ──
            if data["checksum_updates"]:
                sql = """
                    INSERT INTO file_checksums (filename, checksum, processed_at)
                    VALUES (%s, %s, CURRENT_TIMESTAMP)
                    ON DUPLICATE KEY UPDATE
                        checksum = VALUES(checksum),
                        processed_at = CURRENT_TIMESTAMP
                """
                params = [
                    (fname, chk)
                    for fname, chk in data["checksum_updates"].items()
                ]
                if self._execute_many(sql, params):
                    summary["checksums_updated"] = len(params)

        except Exception as e:
            print(f"[MEMORY] Flush error: {e}")
            traceback.print_exc()
            summary["error"] = str(e)

        total = (
            summary["nodes_upserted"] + summary["edges_upserted"] +
            summary["chains_upserted"] + summary["nodes_deleted"] +
            summary["edges_deleted"]
        )
        if total > 0:
            print(f"[MEMORY] Flushed: {summary}")

        return summary

    def flush_if_needed(self) -> Optional[dict]:
        """Flush only if buffer thresholds are met."""
        if self.buffer.should_flush():
            return self.flush()
        return None

    # ───────────────────────────────────────────────────
    # FULL STATE LOAD (Startup)
    # ───────────────────────────────────────────────────

    def load_full_state(self) -> dict:
        """
        Load complete brain state from DB.
        Called once at startup.
        Returns dict with all components.
        """
        if not self._connected:
            return {
                "nodes": [],
                "edges": [],
                "chains": [],
                "thinking_state": {
                    "current_cycle": 0,
                    "total_cycles": 0,
                    "cursor_position": "",
                    "phase": "init",
                    "metrics": {}
                },
                "file_checksums": {},
                "loaded": False
            }

        print("[MEMORY] Loading full state from TiDB...")
        start = time.time()

        nodes = self.load_all_nodes()
        print(f"[MEMORY]   Loaded {len(nodes)} nodes")

        edges = self.load_all_edges()
        print(f"[MEMORY]   Loaded {len(edges)} edges")

        chains = self.load_top_chains(limit=10000)
        print(f"[MEMORY]   Loaded {len(chains)} chains")

        state = self.load_thinking_state()
        print(f"[MEMORY]   Loaded thinking state (cycle {state['total_cycles']})")

        checksums = self.load_file_checksums()
        print(f"[MEMORY]   Loaded {len(checksums)} file checksums")

        elapsed = time.time() - start
        print(f"[MEMORY] Full state loaded in {elapsed:.1f}s")

        return {
            "nodes": nodes,
            "edges": edges,
            "chains": chains,
            "thinking_state": state,
            "file_checksums": checksums,
            "loaded": True
        }

    # ───────────────────────────────────────────────────
    # MAINTENANCE
    # ───────────────────────────────────────────────────

    def prune_weak_edges(self, threshold: float = None) -> int:
        """Delete edges below weight threshold directly from DB."""
        if threshold is None:
            threshold = config.PRUNE_WEIGHT_THRESHOLD
        self._execute(
            "DELETE FROM edges WHERE weight < %s AND source = 'inferred'",
            (threshold,)
        )
        return 0

    def prune_orphan_nodes(self) -> int:
        """Delete nodes with no edges and low weight."""
        self._execute(
            """
            DELETE FROM nodes WHERE connections = 0 
            AND weight < %s AND source = 'inferred'
            """,
            (config.WEIGHT_MIN * 2,)
        )
        return 0

    def get_db_stats(self) -> dict:
        """Get database-level statistics."""
        if not self._connected:
            return {"connected": False}

        node_count = self.get_node_count()
        edge_count = self.get_edge_count()
        chain_count = self.get_chain_count()

        # Count by source
        inferred_nodes = self._execute(
            "SELECT COUNT(*) as cnt FROM nodes WHERE source = 'inferred'",
            fetch="one"
        )
        inferred_edges = self._execute(
            "SELECT COUNT(*) as cnt FROM edges WHERE source = 'inferred'",
            fetch="one"
        )

        return {
            "connected": True,
            "total_nodes": node_count,
            "total_edges": edge_count,
            "total_chains": chain_count,
            "inferred_nodes": inferred_nodes["cnt"] if inferred_nodes else 0,
            "inferred_edges": inferred_edges["cnt"] if inferred_edges else 0,
            "buffer_pending": self.buffer.pending_count
        }

    # ───────────────────────────────────────────────────
    # CLEANUP
    # ───────────────────────────────────────────────────

    def shutdown(self):
        """Graceful shutdown: flush buffer and close connections."""
        print("[MEMORY] Shutting down...")
        # Final flush
        if self._connected:
            try:
                self.flush()
                print("[MEMORY] Final flush completed.")
            except Exception as e:
                print(f"[MEMORY] Final flush error: {e}")
        # Close pool
        self.pool.close_all()
        print("[MEMORY] Connections closed.")