flynn/src/memory/store.ts

import { readFileSync, writeFileSync, existsSync, mkdirSync, readdirSync, statSync } from 'fs';
import { join, relative, dirname } from 'path';
import { MEMORY_CATEGORIES, categoryNamespace, isMemoryCategory, type MemoryCategory } from './categories.js';

/**
 * Configuration for the MemoryStore.
 */
export interface MemoryStoreConfig {
  /** Base directory for memory files. */
  dir: string;
  /** Maximum tokens to inject into system prompt per turn. */
  maxContextTokens: number;
}

export interface PromptMemorySection {
  title: string;
  content: string;
}

/**
 * A single search result from scanning memory files.
 */
export interface SearchResult {
  /** Namespace the match was found in (e.g. 'user', 'sessions/abc123'). */
  namespace: string;
  /** 1-based line number of the match. */
  line: number;
  /** The matched line content. */
  content: string;
  /** Lines of context around the match (1 line above + match + 1 line below). */
  context: string;
}

export interface SearchOptions {
  categories?: MemoryCategory[];
  baseNamespacePrefix?: string;
}

/**
 * Manages persistent markdown memory files on disk.
 *
 * Directory layout:
 *   {dir}/
 *   ├── global.md          # Cross-session knowledge
 *   ├── user.md            # User preferences and facts
 *   └── sessions/
 *       └── {session_id}.md  # Per-session notes
 *
 * Namespaces map directly to filenames: 'user' -> user.md,
 * 'sessions/abc123' -> sessions/abc123.md.
 */
export class MemoryStore {
  private _config: MemoryStoreConfig;
  private _dirtyNamespaces: Set<string> = new Set();

  constructor(config: MemoryStoreConfig) {
    this._config = config;
  }

  // ---------------------------------------------------------------------------
  // Public API
  // ---------------------------------------------------------------------------

  /**
   * Read the contents of a memory file by namespace.
   * Returns an empty string if the file does not exist.
   */
  read(namespace: string): string {
    const filePath = this._namespacePath(namespace);
    if (!existsSync(filePath)) {
      return '';
    }
    return readFileSync(filePath, 'utf-8');
  }

  /**
   * Write content to a memory file.
   *
   * @param namespace - Target namespace (e.g. 'user', 'sessions/abc123').
   * @param content   - Markdown content to write.
   * @param mode      - 'append' adds content after a newline separator;
   *                    'replace' overwrites the file entirely.
   */
  write(namespace: string, content: string, mode: 'append' | 'replace'): void {
    const filePath = this._namespacePath(namespace);

    // Ensure parent directories exist on first write
    const dir = dirname(filePath);
    if (!existsSync(dir)) {
      mkdirSync(dir, { recursive: true });
    }

    if (mode === 'replace') {
      writeFileSync(filePath, content, 'utf-8');
    } else {
      // Append: read existing content, add a newline separator, then the new content
      const existing = existsSync(filePath)
        ? readFileSync(filePath, 'utf-8')
        : '';
      const separator = existing.length > 0 ? '\n' : '';
      writeFileSync(filePath, existing + separator + content, 'utf-8');
    }

    // Mark namespace as needing re-indexing for vector search
    this._dirtyNamespaces.add(namespace);
  }

  /** Read content for a category under a base namespace. */
  readCategory(baseNamespace: string, category: MemoryCategory): string {
    return this.read(categoryNamespace(baseNamespace, category));
  }

  /** Write content for a category under a base namespace. */
  writeCategory(baseNamespace: string, category: MemoryCategory, content: string, mode: 'append' | 'replace'): void {
    this.write(categoryNamespace(baseNamespace, category), content, mode);
  }

  /** List categories that currently exist under a base namespace. */
  listCategories(baseNamespace: string): MemoryCategory[] {
    const categorySet = new Set<MemoryCategory>();
    const prefix = `${baseNamespace}/`;

    for (const namespace of this.listNamespaces()) {
      if (!namespace.startsWith(prefix)) {
        continue;
      }

      const suffix = namespace.slice(prefix.length);
      if (suffix.includes('/')) {
        continue;
      }

      if (isMemoryCategory(suffix)) {
        categorySet.add(suffix);
      }
    }

    return MEMORY_CATEGORIES.filter(category => categorySet.has(category));
  }

  /** Read all category files under a base namespace. */
  readAllCategories(baseNamespace: string): Partial<Record<MemoryCategory, string>> {
    const result: Partial<Record<MemoryCategory, string>> = {};

    for (const category of this.listCategories(baseNamespace)) {
      const content = this.readCategory(baseNamespace, category);
      if (content.length > 0) {
        result[category] = content;
      }
    }

    return result;
  }

  /**
   * Search across all memory files for a keyword or phrase.
   * Performs case-insensitive line-by-line matching.
   * Returns matching lines with 1 line of context above and below.
   */
  search(query: string, opts?: SearchOptions): SearchResult[] {
    const results: SearchResult[] = [];
    const namespaces = this.listNamespaces().filter((namespace) => {
      if (opts?.baseNamespacePrefix && !namespace.startsWith(opts.baseNamespacePrefix)) {
        return false;
      }

      if (opts?.categories && opts.categories.length > 0) {
        const suffix = namespace.split('/').pop() ?? '';
        return isMemoryCategory(suffix) && opts.categories.includes(suffix);
      }

      return true;
    });
    const lowerQuery = query.toLowerCase();

    for (const namespace of namespaces) {
      const content = this.read(namespace);
      if (content.length === 0) {
        continue;
      }

      const lines = content.split('\n');

      for (let i = 0; i < lines.length; i++) {
        if (lines[i].toLowerCase().includes(lowerQuery)) {
          // Gather 1 line of context above and below
          const contextLines: string[] = [];
          if (i > 0) {
            contextLines.push(lines[i - 1]);
          }
          contextLines.push(lines[i]);
          if (i < lines.length - 1) {
            contextLines.push(lines[i + 1]);
          }

          results.push({
            namespace,
            line: i + 1, // 1-based
            content: lines[i],
            context: contextLines.join('\n'),
          });
        }
      }
    }

    return results;
  }

  /**
   * List all available namespaces by scanning the memory directory
   * recursively for `.md` files.
   *
   * Returns namespace strings (file path relative to the base dir, without
   * the `.md` extension). E.g. ['global', 'user', 'sessions/abc123'].
   */
  listNamespaces(): string[] {
    if (!existsSync(this._config.dir)) {
      return [];
    }
    return this._scanDir(this._config.dir);
  }

  /**
   * Return namespaces that have been modified since last call, then clear
   * the dirty set. Used by the background indexer to re-embed changed content.
   */
  getDirtyNamespaces(): string[] {
    const dirty = Array.from(this._dirtyNamespaces);
    this._dirtyNamespaces.clear();
    return dirty;
  }

  /**
   * Mark all existing namespaces as dirty (e.g. for initial full indexing).
   */
  markAllDirty(): void {
    for (const ns of this.listNamespaces()) {
      this._dirtyNamespaces.add(ns);
    }
  }

  /**
   * Build memory context suitable for injection into a system prompt.
   *
   * Reads `user.md` and `global.md`, formats them under markdown headings,
   * and truncates to stay within {@link MemoryStoreConfig.maxContextTokens}
   * (estimated at 4 characters per token).
   */
  getContextForPrompt(): string {
    const sections = this.getPromptSections().map((section) => `## ${section.title}\n\n${section.content}`);

    // Nothing to inject
    if (sections.length === 0) {
      return '';
    }

    const full = sections.join('\n\n');

    // Truncate to fit within the token budget (estimate: 4 chars ≈ 1 token)
    const maxChars = this._config.maxContextTokens * 4;
    if (full.length <= maxChars) {
      return full;
    }
    return full.slice(0, maxChars);
  }

  /** Build memory sections used by prompt injectors. */
  getPromptSections(): PromptMemorySection[] {
    const userMemory = this.read('user');
    const globalMemory = this.read('global');
    const userCategoryMemory = this.readAllCategories('user');
    const globalCategoryMemory = this.readAllCategories('global');

    const sections: PromptMemorySection[] = [];

    if (userMemory.length > 0) {
      sections.push({ title: 'User Memory', content: userMemory });
    }
    if (globalMemory.length > 0) {
      sections.push({ title: 'Global Memory', content: globalMemory });
    }

    for (const category of MEMORY_CATEGORIES) {
      const content = userCategoryMemory[category];
      if (content) {
        sections.push({
          title: `User ${this._categoryLabel(category)}`,
          content,
        });
      }
    }

    for (const category of MEMORY_CATEGORIES) {
      const content = globalCategoryMemory[category];
      if (content) {
        sections.push({
          title: `Global ${this._categoryLabel(category)}`,
          content,
        });
      }
    }

    return sections;
  }

  // ---------------------------------------------------------------------------
  // Private helpers
  // ---------------------------------------------------------------------------

  /** Resolve a namespace string to an absolute file path. */
  private _namespacePath(namespace: string): string {
    return join(this._config.dir, `${namespace}.md`);
  }

  /** Recursively scan a directory for .md files, returning namespace strings. */
  private _scanDir(dir: string): string[] {
    const namespaces: string[] = [];
    const entries = readdirSync(dir);

    for (const entry of entries) {
      const fullPath = join(dir, entry);
      const stat = statSync(fullPath);

      if (stat.isDirectory()) {
        // Recurse into subdirectories
        namespaces.push(...this._scanDir(fullPath));
      } else if (stat.isFile() && entry.endsWith('.md')) {
        // Convert absolute path back to namespace: strip base dir and .md extension
        const rel = relative(this._config.dir, fullPath);
        const namespace = rel.slice(0, -3); // remove '.md'
        namespaces.push(namespace);
      }
    }

    return namespaces;
  }

  private _categoryLabel(category: MemoryCategory): string {
    return `${category.charAt(0).toUpperCase()}${category.slice(1)}`;
  }
}