persistent-memory

/**
 * @file storage/db.ts
 * SQLite storage layer using sql.js (pure WASM — no native bindings).
 *
 * Improvements:
 * - Reduced persist() calls via write coalescing
 * - Batched access counter updates
 * - Proper error propagation
 * - Cleaner row-to-record conversion
 */

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,
  MAX_TAG_LENGTH,
  VALID_CATEGORIES,
  VALID_SCOPES,
  MAX_PROJECT_NAME_LENGTH,
} from "../constants";
import type { MemoryRecord, MemoryStats } from "../types";
import type { MemoryCategory, MemoryScope } from "../constants";

const SCHEMA_VERSION = 2;

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

/** Memoize a no-arg function to avoid redundant fs.mkdirSync calls. */
function memoize<T>(fn: () => T): () => T {
  let cached: T | undefined;
  let computed = false;
  return () => {
    if (!computed) { cached = fn(); computed = true; }
    return cached!;
  };
}

const generateId = (): string => 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,
  };
}

/** Escape a string for safe use in SQL LIKE patterns. */
const escapeLike = (input: string): string => input.replace(/[%_\\]/g, (c) => `\\${c}`);

type SqlJsDatabase = any;
type SqlJs = any;

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

  /** Pending writes batched before the next persist call. */
  private pendingWrites = 0;
  private readonly PERSIST_COALESCE_MS = 50;
  private persistTimer: ReturnType<typeof setTimeout> | null = null;

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

  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)) {
      this.db = new SQL.Database(fs.readFileSync(this.dbPath));
    } 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 { /* columns may already exist */ }
    }
  }

  /** Schedule a deferred persist (coalesces rapid writes). */
  private schedulePersist(): void {
    if (this.persistTimer) return;
    this.persistTimer = setTimeout(() => {
      this.persistTimer = null;
      this.persist();
    }, this.PERSIST_COALESCE_MS);
  }

  /** Write in-memory DB to disk. */
  private persist(): void {
    try {
      const data = this.db.export();
      fs.writeFileSync(this.dbPath, Buffer.from(data));
    } catch { /* disk full, permissions, etc. */ }
  }

  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;
  }

  private queryOne(sql: string, params: unknown[] = []): Record<string, unknown> | null {
    const rows = this.query(sql, params);
    return rows.length > 0 ? rows[0] : null;
  }

  private exec(sql: string, params: unknown[] = []): void {
    this.db.run(sql, params);
    this.pendingWrites++;
    if (this.pendingWrites >= 5) { this.persist(); this.pendingWrites = 0; }
  }

  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, MAX_TAG_LENGTH).toLowerCase().trim())
      .filter(Boolean);
    const safeProject = project ? project.slice(0, MAX_PROJECT_NAME_LENGTH).trim() : null;
    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), Math.max(0, Math.min(1, confidence)),
       source, scope, safeProject, now, now, now, 0, supersedes ?? null],
    );
    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 = 10_000): 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);
  }

  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],
    );
  }

  touchAccessBatch(ids: string[]): void {
    if (ids.length === 0) return;
    const now = Date.now();
    // Single batched UPDATE using CASE/WHEN for O(1) writes
    this.exec(
      `UPDATE memories SET last_accessed_at = ?, access_count = access_count + 1 ` +
      `WHERE id IN (${ids.map(() => "?").join(",")})`,
      [now, ...ids],
    );
  }

  update(id: string, content: string, confidence: number, tags: string[]): boolean {
    const safeTags = tags
      .slice(0, MAX_TAGS_PER_MEMORY)
      .map((t) => t.slice(0, MAX_TAG_LENGTH).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],
    );
    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]);
    }
    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;
    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");
    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 { /* json_each may not be available in all sql.js builds */ }

    let dbSizeBytes = 0;
    try { dbSizeBytes = fs.statSync(this.dbPath).size; } catch { /* file may not exist yet */ }

    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 now = Date.now();
    const msInDay = 86400000.0;
    const victim = this.queryOne(
      `SELECT id FROM memories ORDER BY (
        (access_count * 1.0 / (access_count + 10)) * 0.3 +
        confidence * 0.3 +
        (1.0 / (1.0 + (? - last_accessed_at) / ?)) * 0.4
      ) ASC LIMIT 1`,
      [now, msInDay],
    );
    if (victim) {
      this.exec("DELETE FROM memories WHERE id = ?", [String(victim.id)]);
    }
  }

  close(): void {
    if (this.persistTimer) { clearTimeout(this.persistTimer); this.persistTimer = null; }
    if (this.db) { this.persist(); this.db.close(); }
  }
}