persistent-memory

/**
 * @file storage/db.ts
 * SQLite storage layer using sql.js (pure WASM — no native bindings).
 *
 * sql.js runs SQLite compiled to WebAssembly via Emscripten.
 * - Zero native compilation needed (works with Bun, Node, any runtime)
 * - Async initialization (must call init() before use)
 * - Synchronous queries after init
 * - Manual persistence: db lives in memory, we flush to disk on writes
 *
 * Auto-persist strategy: after every write operation, we export the
 * database to a Buffer and write it to disk. This is fast for our
 * use case (small DB, infrequent writes) and guarantees durability.
 */

import * as path from "path";
import * as os from "os";
import * as fs from "fs";
import * as crypto from "crypto";
import {
  DB_FILENAME,
  MAX_MEMORIES_TOTAL,
  MAX_MEMORY_CONTENT_LENGTH,
  MAX_TAGS_PER_MEMORY,
  VALID_CATEGORIES,
  VALID_SCOPES,
  MAX_PROJECT_NAME_LENGTH,
} from "../constants";
import type { MemoryRecord, MemoryRow, MemoryStats } from "../types";
import type { MemoryCategory, MemoryScope } from "../constants";

const SCHEMA_VERSION = 2;

function defaultMemoryDir(): string {
  const home = os.homedir();
  const dir = path.join(home, ".lmstudio", "plugin-data", "persistent-memory");
  fs.mkdirSync(dir, { recursive: true });
  return dir;
}

function generateId(): string {
  return crypto.randomBytes(12).toString("base64url");
}

function rowToRecord(row: Record<string, unknown>): MemoryRecord {
  let tags: string[] = [];
  try {
    tags = JSON.parse(String(row.tags ?? "[]"));
  } catch {
    tags = [];
  }
  const cat = String(row.category ?? "note");
  const sc = String(row.scope ?? "global");
  return {
    id: String(row.id),
    content: String(row.content),
    category: (VALID_CATEGORIES.includes(cat as MemoryCategory)
      ? cat
      : "note") as MemoryCategory,
    tags,
    confidence: Number(row.confidence ?? 1),
    source: String(row.source ?? "user"),
    scope: (VALID_SCOPES.includes(sc as MemoryScope)
      ? sc
      : "global") as MemoryScope,
    project: row.project ? String(row.project) : null,
    createdAt: Number(row.created_at ?? 0),
    updatedAt: Number(row.updated_at ?? 0),
    lastAccessedAt: Number(row.last_accessed_at ?? 0),
    accessCount: Number(row.access_count ?? 0),
    supersedes: row.supersedes ? String(row.supersedes) : null,
  };
}

function escapeLike(input: string): string {
  return input.replace(/[%_\\]/g, (c) => `\\${c}`);
}

/** sql.js Database type (loaded dynamically). */
type SqlJsDatabase = any;
type SqlJs = any;

export class MemoryDatabase {
  private db!: SqlJsDatabase;
  private readonly dbPath: string;
  private readonly dbDir: string;
  private initialized = false;

  constructor(storagePath?: string) {
    this.dbDir = storagePath || defaultMemoryDir();
    fs.mkdirSync(this.dbDir, { recursive: true });
    this.dbPath = path.join(this.dbDir, DB_FILENAME);
  }

  /** Must be called before any other method. Loads WASM + opens/creates DB. */
  async init(): Promise<void> {
    if (this.initialized) return;

    const initSqlJs = require("sql.js") as (config?: any) => Promise<SqlJs>;

    const wasmPath = path.join(
      path.dirname(require.resolve("sql.js")),
      "sql-wasm.wasm",
    );

    const SQL = await initSqlJs({
      locateFile: () => wasmPath,
    });

    if (fs.existsSync(this.dbPath)) {
      const fileBuffer = fs.readFileSync(this.dbPath);
      this.db = new SQL.Database(fileBuffer);
    } else {
      this.db = new SQL.Database();
    }

    this.setupSchema();
    this.initialized = true;
  }

  private setupSchema(): void {
    const verResult = this.db.exec("PRAGMA user_version");
    const currentVersion =
      verResult.length > 0 ? Number(verResult[0].values[0][0]) : 0;

    if (currentVersion < SCHEMA_VERSION) {
      this.migrate(currentVersion);
      this.db.run(`PRAGMA user_version = ${SCHEMA_VERSION}`);
    }

    this.db.run(`
      CREATE TABLE IF NOT EXISTS memories (
        id               TEXT PRIMARY KEY,
        content          TEXT NOT NULL,
        category         TEXT NOT NULL DEFAULT 'note',
        tags             TEXT NOT NULL DEFAULT '[]',
        confidence       REAL NOT NULL DEFAULT 1.0,
        source           TEXT NOT NULL DEFAULT 'user',
        scope            TEXT NOT NULL DEFAULT 'global',
        project          TEXT,
        created_at       INTEGER NOT NULL,
        updated_at       INTEGER NOT NULL,
        last_accessed_at INTEGER NOT NULL,
        access_count     INTEGER NOT NULL DEFAULT 0,
        supersedes       TEXT
      )
    `);

    this.db.run("CREATE INDEX IF NOT EXISTS idx_cat ON memories(category)");
    this.db.run(
      "CREATE INDEX IF NOT EXISTS idx_created ON memories(created_at)",
    );
    this.db.run(
      "CREATE INDEX IF NOT EXISTS idx_accessed ON memories(last_accessed_at)",
    );
    this.db.run("CREATE INDEX IF NOT EXISTS idx_scope ON memories(scope)");
    this.db.run("CREATE INDEX IF NOT EXISTS idx_project ON memories(project)");

    this.persist();
  }

  private migrate(oldVersion: number): void {
    if (oldVersion >= 1 && oldVersion < 2) {
      try {
        this.db.run(
          "ALTER TABLE memories ADD COLUMN scope TEXT NOT NULL DEFAULT 'global'",
        );
        this.db.run("ALTER TABLE memories ADD COLUMN project TEXT");
      } catch {
      }
    }
  }

  /** Write in-memory DB to disk. Called after every write operation. */
  private persist(): void {
    try {
      const data = this.db.export();
      const buffer = Buffer.from(data);
      fs.writeFileSync(this.dbPath, buffer);
    } catch {
    }
  }

  /** Execute a SELECT and return rows as plain objects. */
  private query(
    sql: string,
    params: unknown[] = [],
  ): Record<string, unknown>[] {
    const stmt = this.db.prepare(sql);
    if (params.length > 0) stmt.bind(params);
    const rows: Record<string, unknown>[] = [];
    while (stmt.step()) {
      rows.push(stmt.getAsObject());
    }
    stmt.free();
    return rows;
  }

  /** Execute a single-row SELECT. */
  private queryOne(
    sql: string,
    params: unknown[] = [],
  ): Record<string, unknown> | null {
    const rows = this.query(sql, params);
    return rows.length > 0 ? rows[0] : null;
  }

  /** Execute a write statement. */
  private exec(sql: string, params: unknown[] = []): void {
    this.db.run(sql, params);
  }

  store(
    content: string,
    category: MemoryCategory,
    tags: string[],
    confidence: number = 1.0,
    source: string = "user",
    supersedes?: string | null,
    scope: MemoryScope = "global",
    project?: string | null,
  ): string {
    if (content.length > MAX_MEMORY_CONTENT_LENGTH)
      content = content.slice(0, MAX_MEMORY_CONTENT_LENGTH);
    const safeTags = tags
      .slice(0, MAX_TAGS_PER_MEMORY)
      .map((t) => t.slice(0, 50).toLowerCase().trim())
      .filter(Boolean);
    const safeProject = project
      ? project.slice(0, MAX_PROJECT_NAME_LENGTH).trim()
      : null;
    const safeConfidence = Math.max(0, Math.min(1, confidence));
    const now = Date.now();
    const id = generateId();

    const countRow = this.queryOne("SELECT COUNT(*) as count FROM memories");
    if (countRow && Number(countRow.count) >= MAX_MEMORIES_TOTAL) {
      this.evictLeastValuable();
    }

    this.exec(
      `INSERT INTO memories (id,content,category,tags,confidence,source,scope,project,
        created_at,updated_at,last_accessed_at,access_count,supersedes)
       VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?)`,
      [
        id,
        content,
        category,
        JSON.stringify(safeTags),
        safeConfidence,
        source,
        scope,
        safeProject,
        now,
        now,
        now,
        0,
        supersedes ?? null,
      ],
    );
    this.persist();
    return id;
  }

  getById(id: string): MemoryRecord | null {
    const row = this.queryOne("SELECT * FROM memories WHERE id = ?", [id]);
    return row ? rowToRecord(row) : null;
  }

  getByIds(ids: string[]): MemoryRecord[] {
    if (ids.length === 0) return [];
    const ph = ids.map(() => "?").join(",");
    return this.query(`SELECT * FROM memories WHERE id IN (${ph})`, ids).map(
      rowToRecord,
    );
  }

  getAll(limit: number = 100): MemoryRecord[] {
    return this.query(
      "SELECT * FROM memories ORDER BY last_accessed_at DESC LIMIT ?",
      [limit],
    ).map(rowToRecord);
  }

  getByCategory(category: MemoryCategory, limit: number = 50): MemoryRecord[] {
    return this.query(
      "SELECT * FROM memories WHERE category = ? ORDER BY last_accessed_at DESC LIMIT ?",
      [category, limit],
    ).map(rowToRecord);
  }

  getByTag(tag: string, limit: number = 50): MemoryRecord[] {
    const escaped = escapeLike(tag.toLowerCase().trim());
    return this.query(
      `SELECT * FROM memories WHERE tags LIKE ? ESCAPE '\\' ORDER BY last_accessed_at DESC LIMIT ?`,
      [`%"${escaped}"%`, limit],
    ).map(rowToRecord);
  }

  getByProject(project: string, limit: number = 50): MemoryRecord[] {
    return this.query(
      "SELECT * FROM memories WHERE project = ? ORDER BY last_accessed_at DESC LIMIT ?",
      [project.trim(), limit],
    ).map(rowToRecord);
  }

  getByScope(scope: MemoryScope, limit: number = 50): MemoryRecord[] {
    return this.query(
      "SELECT * FROM memories WHERE scope = ? ORDER BY last_accessed_at DESC LIMIT ?",
      [scope, limit],
    ).map(rowToRecord);
  }

  /** Text search using LIKE (sql.js doesn't support FTS5). TF-IDF handles semantic search. */
  ftsSearch(query: string, limit: number = 20): MemoryRecord[] {
    const escaped = escapeLike(query);
    return this.query(
      `SELECT * FROM memories WHERE content LIKE ? ESCAPE '\\' ORDER BY last_accessed_at DESC LIMIT ?`,
      [`%${escaped}%`, limit],
    ).map(rowToRecord);
  }

  getRecent(limit: number = 10): MemoryRecord[] {
    return this.query(
      "SELECT * FROM memories ORDER BY created_at DESC LIMIT ?",
      [limit],
    ).map(rowToRecord);
  }

  touchAccess(id: string): void {
    this.exec(
      "UPDATE memories SET last_accessed_at = ?, access_count = access_count + 1 WHERE id = ?",
      [Date.now(), id],
    );
    this.persist();
  }

  touchAccessBatch(ids: string[]): void {
    if (ids.length === 0) return;
    const now = Date.now();
    for (const id of ids) {
      this.exec(
        "UPDATE memories SET last_accessed_at = ?, access_count = access_count + 1 WHERE id = ?",
        [now, id],
      );
    }
    this.persist();
  }

  update(
    id: string,
    content: string,
    confidence: number,
    tags: string[],
  ): boolean {
    const safeTags = tags
      .slice(0, MAX_TAGS_PER_MEMORY)
      .map((t) => t.slice(0, 50).toLowerCase().trim())
      .filter(Boolean);
    this.exec(
      "UPDATE memories SET content = ?, updated_at = ?, confidence = ?, tags = ? WHERE id = ?",
      [
        content.slice(0, MAX_MEMORY_CONTENT_LENGTH),
        Date.now(),
        Math.max(0, Math.min(1, confidence)),
        JSON.stringify(safeTags),
        id,
      ],
    );
    this.persist();
    return this.getById(id) !== null;
  }

  delete(id: string): boolean {
    const existed = this.getById(id) !== null;
    if (existed) {
      this.exec("DELETE FROM memories WHERE id = ?", [id]);
      this.persist();
    }
    return existed;
  }

  deleteByPattern(pattern: string): number {
    const escaped = escapeLike(pattern);
    const before = this.queryOne("SELECT COUNT(*) as c FROM memories") as {
      c: number;
    } | null;
    this.exec(`DELETE FROM memories WHERE content LIKE ? ESCAPE '\\'`, [
      `%${escaped}%`,
    ]);
    const after = this.queryOne("SELECT COUNT(*) as c FROM memories") as {
      c: number;
    } | null;
    this.persist();
    return Number(before?.c ?? 0) - Number(after?.c ?? 0);
  }

  deleteAll(): number {
    const countRow = this.queryOne("SELECT COUNT(*) as count FROM memories");
    const count = Number(countRow?.count ?? 0);
    this.exec("DELETE FROM memories");
    this.persist();
    return count;
  }

  getStats(): MemoryStats {
    const totalRow = this.queryOne("SELECT COUNT(*) as count FROM memories");
    const totalMemories = Number(totalRow?.count ?? 0);

    const catRows = this.query(
      "SELECT category, COUNT(*) as count FROM memories GROUP BY category",
    );
    const byCategory: Record<string, number> = {};
    for (const r of catRows) byCategory[String(r.category)] = Number(r.count);

    const oldest = this.queryOne(
      "SELECT content, created_at FROM memories ORDER BY created_at ASC LIMIT 1",
    );
    const newest = this.queryOne(
      "SELECT content, created_at FROM memories ORDER BY created_at DESC LIMIT 1",
    );
    const top = this.queryOne(
      "SELECT content, access_count FROM memories ORDER BY access_count DESC LIMIT 1",
    );

    let totalTags = 0;
    try {
      const tagRow = this.queryOne(
        "SELECT COUNT(DISTINCT value) as count FROM memories, json_each(memories.tags)",
      );
      totalTags = Number(tagRow?.count ?? 0);
    } catch {
    }

    let dbSizeBytes = 0;
    try {
      dbSizeBytes = fs.statSync(this.dbPath).size;
    } catch {
    }

    return {
      totalMemories,
      byCategory,
      oldestMemory: oldest
        ? new Date(Number(oldest.created_at)).toISOString()
        : null,
      newestMemory: newest
        ? new Date(Number(newest.created_at)).toISOString()
        : null,
      mostAccessed:
        top && Number(top.access_count) > 0
          ? {
              content: String(top.content),
              accessCount: Number(top.access_count),
            }
          : null,
      totalTags,
      dbSizeBytes,
    };
  }

  private evictLeastValuable(): void {
    const victim = this.queryOne(
      "SELECT id FROM memories ORDER BY (access_count*0.3+confidence*0.3+(last_accessed_at/86400000.0)*0.4) ASC LIMIT 1",
    );
    if (victim)
      this.exec("DELETE FROM memories WHERE id = ?", [String(victim.id)]);
  }

  close(): void {
    if (this.db) {
      this.persist();
      this.db.close();
    }
  }
}