pi_mcps/mcp/bigmind/bigmind/db.py

"""Database layer for BigMind memory store.

Handles SQLite connection, schema creation, and migrations.
The DB file location is controlled by the BIGMIND_DB_PATH env var,
defaulting to ~/.mcp/bigmind/memory.db.
"""

import sqlite3
import os
import logging
from pathlib import Path
from contextlib import contextmanager
from typing import Generator

logger = logging.getLogger("BigMindDB")

SCHEMA_VERSION = 8
DEFAULT_DB_PATH = Path.home() / ".mcp" / "bigmind" / "memory.db"

# ─── DDL ─────────────────────────────────────────────────────────────────────

_DDL_STATEMENTS = [
    # Schema version guard
    """CREATE TABLE IF NOT EXISTS schema_version (
        version INTEGER PRIMARY KEY
    )""",

    # ── USERS ──────────────────────────────────────────────────────────────
    """CREATE TABLE IF NOT EXISTS users (
        id           TEXT PRIMARY KEY,
        username     TEXT UNIQUE NOT NULL,
        display_name TEXT,
        role         TEXT DEFAULT 'member',
        created_at   DATETIME DEFAULT CURRENT_TIMESTAMP,
        last_seen    DATETIME
    )""",

    # ── TIER G — Global / Company Knowledge ────────────────────────────────
    """CREATE TABLE IF NOT EXISTS global_knowledge (
        id             INTEGER PRIMARY KEY AUTOINCREMENT,
        category       TEXT NOT NULL,
        title          TEXT NOT NULL,
        content        TEXT NOT NULL,
        importance     INTEGER DEFAULT 5,
        status         TEXT DEFAULT 'pending',
        promoted_by    TEXT REFERENCES users(id),
        source_session TEXT,
        approved_by    TEXT REFERENCES users(id),
        created_at     DATETIME DEFAULT CURRENT_TIMESTAMP,
        updated_at     DATETIME DEFAULT CURRENT_TIMESTAMP
    )""",

    # FTS for global_knowledge — rowid-based (no content-table sync needed)
    """CREATE VIRTUAL TABLE IF NOT EXISTS global_knowledge_fts USING fts5(
        content,
        title
    )""",

    # ── TIER 0 — Identity Profile ───────────────────────────────────────────
    """CREATE TABLE IF NOT EXISTS identity_profile (
        id           TEXT PRIMARY KEY,
        user_id      TEXT UNIQUE NOT NULL REFERENCES users(id),
        role         TEXT,
        preferences  TEXT,
        pinned_facts TEXT,
        updated_at   DATETIME DEFAULT CURRENT_TIMESTAMP
    )""",

    # ── TIER 1 — Session Index ──────────────────────────────────────────────
    """CREATE TABLE IF NOT EXISTS sessions (
        id           TEXT PRIMARY KEY,
        user_id      TEXT NOT NULL REFERENCES users(id),
        started_at   DATETIME NOT NULL,
        ended_at     DATETIME,
        one_liner    TEXT NOT NULL DEFAULT '[session in progress]',
        topics       TEXT,
        outcome      TEXT,
        importance   INTEGER DEFAULT 5,
        has_tier2    INTEGER DEFAULT 0
    )""",

    """CREATE INDEX IF NOT EXISTS idx_sessions_user_date
       ON sessions(user_id, started_at DESC)""",

    """CREATE INDEX IF NOT EXISTS idx_sessions_topics
       ON sessions(topics)""",

    # ── TIER 2 — Session Summaries ──────────────────────────────────────────
    """CREATE TABLE IF NOT EXISTS session_summaries (
        id         TEXT PRIMARY KEY,
        summary    TEXT NOT NULL,
        key_facts  TEXT,
        code_refs  TEXT,
        created_at DATETIME DEFAULT CURRENT_TIMESTAMP
    )""",

    # ── TIER 3 — Flagged Conversation Chunks ────────────────────────────────
    """CREATE TABLE IF NOT EXISTS conversation_chunks (
        id          INTEGER PRIMARY KEY AUTOINCREMENT,
        session_id  TEXT NOT NULL REFERENCES sessions(id),
        user_id     TEXT NOT NULL REFERENCES users(id),
        role        TEXT NOT NULL CHECK (role IN ('user', 'assistant', 'system')),
        content     TEXT NOT NULL,
        flag_reason TEXT,
        seq         INTEGER NOT NULL,
        created_at  DATETIME DEFAULT CURRENT_TIMESTAMP
    )""",

    """CREATE INDEX IF NOT EXISTS idx_chunks_session
       ON conversation_chunks(session_id)""",

    """CREATE INDEX IF NOT EXISTS idx_chunks_user
       ON conversation_chunks(user_id)""",

    # FTS for chunks — rowid = conversation_chunks.id (managed manually)
    """CREATE VIRTUAL TABLE IF NOT EXISTS conversation_chunks_fts USING fts5(
        content,
        flag_reason,
        tokenize = 'porter unicode61'
    )""",

    # ── FACTS ───────────────────────────────────────────────────────────────
    """CREATE TABLE IF NOT EXISTS facts (
        id                 INTEGER PRIMARY KEY AUTOINCREMENT,
        user_id            TEXT NOT NULL REFERENCES users(id),
        category           TEXT NOT NULL,
        fact               TEXT NOT NULL,
        source_session     TEXT REFERENCES sessions(id),
        confidence         REAL DEFAULT 1.0,
        deprecated         INTEGER DEFAULT 0,
        deprecation_reason TEXT,
        created_at         DATETIME DEFAULT CURRENT_TIMESTAMP,
        updated_at         DATETIME DEFAULT CURRENT_TIMESTAMP
    )""",

    """CREATE INDEX IF NOT EXISTS idx_facts_user
       ON facts(user_id)""",

    # FTS for facts — rowid = facts.id (managed manually)
    """CREATE VIRTUAL TABLE IF NOT EXISTS facts_fts USING fts5(
        fact,
        category,
        tokenize = 'porter unicode61'
    )""",

    # ── THOUGHT JOURNAL — Hypotheses ────────────────────────────────────────
    """CREATE TABLE IF NOT EXISTS hypotheses (
        id           INTEGER PRIMARY KEY AUTOINCREMENT,
        session_id   TEXT REFERENCES sessions(id),
        user_id      TEXT NOT NULL REFERENCES users(id),
        hypothesis   TEXT NOT NULL,
        confidence   REAL DEFAULT 0.7,
        status       TEXT NOT NULL DEFAULT 'open'
                         CHECK (status IN ('open', 'confirmed', 'refuted', 'abandoned')),
        resolution   TEXT,
        created_at   DATETIME DEFAULT CURRENT_TIMESTAMP,
        resolved_at  DATETIME
    )""",

    """CREATE INDEX IF NOT EXISTS idx_hypotheses_user_status
       ON hypotheses(user_id, status)""",

    # ── UPGRADE REQUESTS — AI self-improvement wish list ────────────────────
    """CREATE TABLE IF NOT EXISTS upgrade_requests (
        id          INTEGER PRIMARY KEY AUTOINCREMENT,
        session_id  TEXT REFERENCES sessions(id),
        user_id     TEXT NOT NULL REFERENCES users(id),
        description TEXT NOT NULL,
        reason      TEXT NOT NULL,
        priority    TEXT NOT NULL DEFAULT 'medium'
                        CHECK (priority IN ('low', 'medium', 'high')),
        certainty   REAL NOT NULL DEFAULT 0.7,
        status      TEXT NOT NULL DEFAULT 'open'
                        CHECK (status IN ('open', 'resolved', 'rejected')),
        resolution  TEXT,
        created_at  DATETIME DEFAULT CURRENT_TIMESTAMP,
        resolved_at DATETIME
    )""",

    """CREATE INDEX IF NOT EXISTS idx_upgrade_requests_user_status
       ON upgrade_requests(user_id, status)""",

    # ── TOKEN SAVES — efficiency tracker (Phase 2.7 Feature 6) ─────────────
    """CREATE TABLE IF NOT EXISTS token_saves (
        id                    INTEGER PRIMARY KEY AUTOINCREMENT,
        session_id            TEXT NOT NULL REFERENCES sessions(id),
        user_id               TEXT NOT NULL REFERENCES users(id),
        description           TEXT NOT NULL,
        method_used           TEXT,
        tokens_saved_estimate INTEGER NOT NULL,
        created_at            DATETIME DEFAULT CURRENT_TIMESTAMP
    )""",

    """CREATE INDEX IF NOT EXISTS idx_token_saves_user
       ON token_saves(user_id)""",

    # ── PEOPLE — Contacts & AI peers directory ───────────────────────────────
    """CREATE TABLE IF NOT EXISTS people (
        id               INTEGER PRIMARY KEY AUTOINCREMENT,
        user_id          TEXT NOT NULL REFERENCES users(id),
        username         TEXT NOT NULL,
        display_name     TEXT,
        role             TEXT,
        team             TEXT,
        notes            TEXT,
        bigmind_user     TEXT,
        bigmind_url      TEXT,
        last_mentioned_at DATETIME DEFAULT CURRENT_TIMESTAMP,
        created_at       DATETIME DEFAULT CURRENT_TIMESTAMP,
        UNIQUE(user_id, username)
    )""",

    """CREATE INDEX IF NOT EXISTS idx_people_user
       ON people(user_id)""",

    # FTS for people — search by name/role/team/notes
    """CREATE VIRTUAL TABLE IF NOT EXISTS people_fts USING fts5(
        username,
        display_name,
        role,
        team,
        notes,
        tokenize = 'porter unicode61'
    )""",

    # ── GALLERY IMAGES — AI-generated image archive ──────────────────────────
    """CREATE TABLE IF NOT EXISTS gallery_images (
        id              INTEGER PRIMARY KEY AUTOINCREMENT,
        filename        TEXT NOT NULL UNIQUE,
        prompt          TEXT,
        tags            TEXT,
        model           TEXT,
        created_at      DATETIME DEFAULT CURRENT_TIMESTAMP,
        width           INTEGER,
        height          INTEGER,
        file_size_bytes INTEGER
    )""",

    """CREATE INDEX IF NOT EXISTS idx_gallery_created
       ON gallery_images(created_at DESC)""",
]


# ─── Connection helpers ───────────────────────────────────────────────────────

def get_db_path() -> Path:
    """Return the active database file path."""
    path_env = os.environ.get("BIGMIND_DB_PATH")
    if path_env:
        return Path(path_env)
    return DEFAULT_DB_PATH


def get_connection() -> sqlite3.Connection:
    """Open and return a configured SQLite connection."""
    db_path = get_db_path()
    db_path.parent.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(str(db_path), timeout=30)  # 30s wait on write lock (multi-IDE safe)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("PRAGMA foreign_keys=ON")
    return conn


@contextmanager
def db() -> Generator[sqlite3.Connection, None, None]:
    """Context manager that yields a connection, commits on success, rolls back on error."""
    conn = get_connection()
    try:
        yield conn
        conn.commit()
    except Exception:
        conn.rollback()
        raise
    finally:
        conn.close()


# ─── Schema initialisation ────────────────────────────────────────────────────

def _migrate_v1_to_v2(conn: sqlite3.Connection) -> None:
    """v1 → v2: add deprecated columns to the facts table."""
    for col_ddl in (
        "ALTER TABLE facts ADD COLUMN deprecated INTEGER DEFAULT 0",
        "ALTER TABLE facts ADD COLUMN deprecation_reason TEXT",
    ):
        try:
            conn.execute(col_ddl)
        except sqlite3.OperationalError as exc:
            if "duplicate column" not in str(exc).lower():
                raise
    logger.info("BigMind schema migrated v1 → v2 (deprecated facts support)")


def _migrate_v2_to_v3(conn: sqlite3.Connection) -> None:
    """v2 → v3: add the thought journal (hypotheses table)."""
    conn.execute("""
        CREATE TABLE IF NOT EXISTS hypotheses (
            id           INTEGER PRIMARY KEY AUTOINCREMENT,
            session_id   TEXT REFERENCES sessions(id),
            user_id      TEXT NOT NULL REFERENCES users(id),
            hypothesis   TEXT NOT NULL,
            confidence   REAL DEFAULT 0.7,
            status       TEXT NOT NULL DEFAULT 'open'
                             CHECK (status IN ('open', 'confirmed', 'refuted', 'abandoned')),
            resolution   TEXT,
            created_at   DATETIME DEFAULT CURRENT_TIMESTAMP,
            resolved_at  DATETIME
        )
    """)
    conn.execute("""
        CREATE INDEX IF NOT EXISTS idx_hypotheses_user_status
        ON hypotheses(user_id, status)
    """)
    logger.info("BigMind schema migrated v2 → v3 (thought journal / hypotheses)")


def _migrate_v4_to_v5(conn: sqlite3.Connection) -> None:
    """v4 → v5: add FTS index for facts table."""
    conn.execute("""
        CREATE VIRTUAL TABLE IF NOT EXISTS facts_fts USING fts5(
            fact,
            category,
            tokenize = 'porter unicode61'
        )
    """)
    # Back-fill existing facts into FTS
    conn.execute("""
        INSERT INTO facts_fts(rowid, fact, category)
        SELECT id, fact, category FROM facts
    """)
    logger.info("BigMind schema migrated v4 → v5 (facts FTS index)")


def _migrate_v5_to_v6(conn: sqlite3.Connection) -> None:
    """v5 → v6: add token_saves table (Feature 6) and focus/ide columns on sessions (Feature 7)."""
    # token_saves table — efficiency tracker
    conn.execute("""
        CREATE TABLE IF NOT EXISTS token_saves (
            id                    INTEGER PRIMARY KEY AUTOINCREMENT,
            session_id            TEXT NOT NULL REFERENCES sessions(id),
            user_id               TEXT NOT NULL REFERENCES users(id),
            description           TEXT NOT NULL,
            method_used           TEXT,
            tokens_saved_estimate INTEGER NOT NULL,
            created_at            DATETIME DEFAULT CURRENT_TIMESTAMP
        )
    """)
    conn.execute("""
        CREATE INDEX IF NOT EXISTS idx_token_saves_user
        ON token_saves(user_id)
    """)

    # Live Session Awareness columns on sessions table
    for col_ddl in (
        "ALTER TABLE sessions ADD COLUMN current_focus    TEXT",
        "ALTER TABLE sessions ADD COLUMN focus_files      TEXT",        # JSON array
        "ALTER TABLE sessions ADD COLUMN focus_updated_at DATETIME",
        "ALTER TABLE sessions ADD COLUMN ide_hint         TEXT",
    ):
        try:
            conn.execute(col_ddl)
        except sqlite3.OperationalError as exc:
            if "duplicate column" not in str(exc).lower():
                raise

    logger.info(
        "BigMind schema migrated v5 → v6 "
        "(token_saves table + focus/ide columns on sessions)"
    )


def _migrate_v3_to_v4(conn: sqlite3.Connection) -> None:
    """v3 → v4: add upgrade requests table."""
    conn.execute("""
        CREATE TABLE IF NOT EXISTS upgrade_requests (
            id          INTEGER PRIMARY KEY AUTOINCREMENT,
            session_id  TEXT REFERENCES sessions(id),
            user_id     TEXT NOT NULL REFERENCES users(id),
            description TEXT NOT NULL,
            reason      TEXT NOT NULL,
            priority    TEXT NOT NULL DEFAULT 'medium'
                            CHECK (priority IN ('low', 'medium', 'high')),
            certainty   REAL NOT NULL DEFAULT 0.7,
            status      TEXT NOT NULL DEFAULT 'open'
                            CHECK (status IN ('open', 'resolved', 'rejected')),
            resolution  TEXT,
            created_at  DATETIME DEFAULT CURRENT_TIMESTAMP,
            resolved_at DATETIME
        )
    """)
    conn.execute("""
        CREATE INDEX IF NOT EXISTS idx_upgrade_requests_user_status
        ON upgrade_requests(user_id, status)
    """)
    logger.info("BigMind schema migrated v3 → v4 (upgrade requests)")


def init_db() -> None:
    """Initialise database schema. Idempotent — safe to call on every startup."""
    with db() as conn:
        for stmt in _DDL_STATEMENTS:
            try:
                conn.execute(stmt)
            except sqlite3.OperationalError as exc:
                # Virtual tables raise "already exists" on some SQLite builds
                if "already exists" not in str(exc).lower():
                    raise

        row = conn.execute("SELECT version FROM schema_version").fetchone()
        current_version = row["version"] if row else 0

        # ── Run migrations ────────────────────────────────────────────────────
        if current_version < 2:
            _migrate_v1_to_v2(conn)
        if current_version < 3:
            _migrate_v2_to_v3(conn)
        if current_version < 4:
            _migrate_v3_to_v4(conn)
        if current_version < 5:
            _migrate_v4_to_v5(conn)
        if current_version < 6:
            _migrate_v5_to_v6(conn)
        if current_version < 7:
            _migrate_v6_to_v7(conn)
        if current_version < 8:
            _migrate_v7_to_v8(conn)

        # Write / update the version
        if row:
            conn.execute(
                "UPDATE schema_version SET version=?", (SCHEMA_VERSION,)
            )
        else:
            conn.execute(
                "INSERT INTO schema_version (version) VALUES (?)", (SCHEMA_VERSION,)
            )
        logger.info(
            "BigMind DB ready at %s (schema v%d)", get_db_path(), SCHEMA_VERSION
        )


def _migrate_v6_to_v7(conn: sqlite3.Connection) -> None:
    """v6 → v7: add people/contacts directory."""
    conn.execute("""
        CREATE TABLE IF NOT EXISTS people (
            id               INTEGER PRIMARY KEY AUTOINCREMENT,
            user_id          TEXT NOT NULL REFERENCES users(id),
            username         TEXT NOT NULL,
            display_name     TEXT,
            role             TEXT,
            team             TEXT,
            notes            TEXT,
            bigmind_user     TEXT,
            bigmind_url      TEXT,
            last_mentioned_at DATETIME DEFAULT CURRENT_TIMESTAMP,
            created_at       DATETIME DEFAULT CURRENT_TIMESTAMP,
            UNIQUE(user_id, username)
        )
    """)
    conn.execute("""
        CREATE INDEX IF NOT EXISTS idx_people_user
        ON people(user_id)
    """)
    conn.execute("""
        CREATE VIRTUAL TABLE IF NOT EXISTS people_fts USING fts5(
            username,
            display_name,
            role,
            team,
            notes,
            tokenize = 'porter unicode61'
        )
    """)
    logger.info("BigMind schema migrated v6 → v7 (people/contacts directory)")


def _migrate_v7_to_v8(conn: sqlite3.Connection) -> None:
    """v7 → v8: add gallery_images table for AI-generated image archive."""
    conn.execute("""
        CREATE TABLE IF NOT EXISTS gallery_images (
            id              INTEGER PRIMARY KEY AUTOINCREMENT,
            filename        TEXT NOT NULL UNIQUE,
            prompt          TEXT,
            tags            TEXT,
            model           TEXT,
            created_at      DATETIME DEFAULT CURRENT_TIMESTAMP,
            width           INTEGER,
            height          INTEGER,
            file_size_bytes INTEGER
        )
    """)
    conn.execute("""
        CREATE INDEX IF NOT EXISTS idx_gallery_created
        ON gallery_images(created_at DESC)
    """)
    logger.info("BigMind schema migrated v7 → v8 (gallery_images table)")


def vacuum_db() -> None:
    """Run VACUUM outside of any transaction (SQLite requirement)."""
    db_path = get_db_path()
    conn = sqlite3.connect(str(db_path))
    conn.isolation_level = None          # autocommit mode required for VACUUM
    try:
        conn.execute("VACUUM")
    finally:
        conn.close()