open-brainy/src/plugin.ts

/**
 * Brainy Plugin System
 *
 * Simple plugin architecture for two use cases:
 * 1. Native acceleration (@soulcraft/cortex)
 * 2. Custom storage adapters (e.g., Redis, DynamoDB, custom backends)
 *
 * Plugins are auto-detected by package name or registered manually.
 */

import type {
  StorageAdapter,
  Vector,
  VectorDocument,
  GraphVerb,
} from './coreTypes.js'
import type { NounType, VerbType } from './types/graphTypes.js'
import type { MetadataIndexStats } from './utils/metadataIndex.js'
import type { GraphIndexStats } from './graph/graphAdjacencyIndex.js'

// Re-export the provider contracts that already live closer to their
// implementations so a plugin author (Cortex) can import the *entire*
// provider surface from one stable entrypoint: `@soulcraft/brainy/plugin`.
export type { ColumnStoreProvider } from './indexes/columnStore/types.js'
export type {
  AggregationProvider,
  AggregateDefinition,
  AggregateGroupState,
  AggregateResult,
  AggregateQueryParams,
  GroupByDimension,
} from './types/brainy.types.js'

/**
 * Plugin interface — all brainy plugins must implement this.
 */
export interface BrainyPlugin {
  /** Unique plugin name (typically the npm package name) */
  name: string

  /**
   * Called by brainy during init() to activate the plugin.
   * Return true if activation succeeded, false to skip.
   */
  activate(context: BrainyPluginContext): Promise<boolean>

  /**
   * Called when brainy.close() is invoked. Optional cleanup.
   */
  deactivate?(): Promise<void>
}

/**
 * Context passed to plugins during activation.
 */
export interface BrainyPluginContext {
  /**
   * Register a provider for a named subsystem.
   *
   * Well-known provider keys (used by cortex):
   * - 'metadataIndex'      — MetadataIndexManager replacement
   * - 'graphIndex'         — GraphAdjacencyIndex replacement
   * - 'entityIdMapper'     — EntityIdMapper replacement
   * - 'cache'              — UnifiedCache replacement
   * - 'vector'             — JsHnswVectorIndex replacement (vector index engine)
   * - 'roaring'            — RoaringBitmap32 replacement
   * - 'embeddings'         — Embedding engine replacement (single text)
   * - 'embedBatch'         — Batch embedding engine (texts[] → vectors[])
   * - 'distance'           — Distance function overrides
   * - 'msgpack'            — Msgpack encode/decode
   * - 'aggregation'        — AggregationIndex replacement (incremental aggregates)
   *
   * Storage adapter keys:
   * - 'storage:<name>'     — Custom storage adapter factory
   */
  registerProvider(key: string, implementation: unknown): void

  /** Brainy version for compatibility checks */
  readonly version: string
}

// ===========================================================================
// Provider contracts — the EXACT surface Brainy calls on each registered
// provider.
//
// These interfaces are the type-level half of the provider-parity guarantee.
// They capture only what Brainy actually invokes, so a native accelerator
// (Cortex) that declares `implements MetadataIndexProvider` gets a compile
// error the moment a method Brainy depends on is dropped or its signature
// drifts. Brainy's own baseline classes (`MetadataIndexManager`,
// `GraphAdjacencyIndex`, `JsHnswVectorIndex`, `EntityIdMapper`, `UnifiedCache`)
// implement them too, so the contract can never silently diverge from the
// thing Brainy ships.
//
// Keep these in lockstep with the call sites in `brainy.ts` and the index
// classes. When Brainy starts calling a new member, add it here — every
// implementation then fails to compile until it provides the member.
// ===========================================================================

/**
 * The `'metadataIndex'` provider — a drop-in for `MetadataIndexManager`.
 * Brainy calls this surface via `this.metadataIndex.*` (see `brainy.ts`) and
 * the transactional add/remove operations.
 */
export interface MetadataIndexProvider {
  init(): Promise<void>
  flush(): Promise<void>
  rebuild(): Promise<void>

  addToIndex(id: string, entityOrMetadata: any, skipFlush?: boolean, deferWrites?: boolean): Promise<void>
  removeFromIndex(id: string, metadata?: any): Promise<void>

  getIds(field: string, value: any): Promise<string[]>
  getIdsForFilter(filter: any): Promise<string[]>
  getIdsForTextQuery(query: string): Promise<Array<{ id: string; matchCount: number }>>
  getSortedIdsForFilter(filter: any, orderBy: string, order?: 'asc' | 'desc'): Promise<string[]>
  getFilterValues(field: string): Promise<string[]>
  getFilterFields(): Promise<string[]>
  getFieldValueForEntity(entityId: string, field: string): Promise<any>
  getFieldsForType(nounType: NounType): Promise<Array<{ field: string; affinity: number; occurrences: number; totalEntities: number }>>
  getFieldStatistics(): Promise<Map<string, unknown>>
  getFieldsWithCardinality(): Promise<Array<{ field: string; cardinality: number; distribution: string }>>
  getOptimalQueryPlan(filters: Record<string, any>): Promise<unknown>

  /** Report which index path a `where` clause on `field` will hit (drives `brain.explain()`). */
  explainField(field: string): Promise<{ path: 'column-store' | 'sparse-chunked' | 'none'; notes?: string }>

  getCountForCriteria(field: string, value: any): Promise<number>
  getEntityCountByType(type: string): number
  getEntityCountByTypeEnum(type: NounType): number
  getTotalEntityCount(): number
  getAllEntityCounts(): Map<string, number>
  getTopNounTypes(n: number): NounType[]
  getTopVerbTypes(n: number): VerbType[]
  getAllNounTypeCounts(): Map<NounType, number>
  getAllVerbTypeCounts(): Map<VerbType, number>
  getAllVFSEntityCounts(): Promise<Map<string, number>>

  detectAndRepairCorruption(): Promise<void>
  validateConsistency(): Promise<{
    healthy: boolean
    avgEntriesPerEntity: number
    entityCount: number
    indexEntryCount: number
    recommendation: string | null
  }>

  tokenize(text: string): string[]
  extractTextContent(data: any): string

  getStats(): Promise<MetadataIndexStats>

  /**
   * The shared UUID ↔ int mapper — the single source of truth for entity-int
   * resolution at the provider boundary. The coordinator (`brainy.ts`) resolves
   * UUID → int exactly once before every graph-index call (`getOrAssign` on
   * writes, `getInt` on reads — `undefined` means "never mapped", i.e. the
   * entity has no relations) and converts provider-returned ints back with
   * `getUuid`. Ints are u32 today (the `EntityIdSpaceExceeded` guard enforces
   * the ceiling on the JS path), so `Number(bigint)` narrowing is lossless.
   */
  getIdMapper(): {
    getOrAssign(uuid: string): number
    getInt(uuid: string): number | undefined
    getUuid(intId: number): string | undefined
  }

  /** The column store the coordinator delegates `where`/`orderBy` to. */
  readonly columnStore: import('./indexes/columnStore/types.js').ColumnStoreProvider
}

/**
 * The `'graphIndex'` provider — a drop-in for `GraphAdjacencyIndex`.
 * Brainy calls this surface via `this.graphIndex.*` (some optional-chained).
 *
 * **8.0 u64 contract — BigInt at the boundary.** Reads take entity ints
 * (from the metadata index's idMapper) and return entity/verb ints as
 * `bigint[]`. The coordinator owns ALL UUID ↔ int conversion: it resolves
 * UUIDs to ints once at the `brainy.ts` boundary (`getOrAssign` on writes,
 * `getInt` on reads, returning empty results for unmapped UUIDs without
 * calling the provider) and maps returned ints back (`getUuid` for entities,
 * {@link GraphIndexProvider.verbIntsToIds} for verbs). Implementations may
 * stay u32 internally — `Number(bigint)` narrowing is lossless under the
 * shipped `EntityIdSpaceExceeded` u32 guard — but must speak the bigint
 * contract at this surface.
 */
export interface GraphIndexProvider {
  /** `false` until the index has loaded; Brainy probes this before fast paths. */
  readonly isInitialized: boolean

  /**
   * @description Entity ints reachable from `id` (1 hop), deduped.
   * @param id - The entity's interned int (from the shared idMapper).
   * @param options - Direction (`'both'` default) and limit/offset pagination.
   * @returns Neighbor entity ints. Empty when the entity has no edges.
   */
  getNeighbors(
    id: bigint,
    options?: { direction?: 'in' | 'out' | 'both'; limit?: number; offset?: number }
  ): Promise<bigint[]>
  /**
   * @description Verb ints for all edges originating at `sourceInt`.
   * @param sourceInt - The source entity's interned int.
   * @param options - Optional limit/offset pagination.
   * @returns Verb ints, resolvable via {@link GraphIndexProvider.verbIntsToIds}.
   */
  getVerbIdsBySource(sourceInt: bigint, options?: { limit?: number; offset?: number }): Promise<bigint[]>
  /**
   * @description Verb ints for all edges pointing at `targetInt`.
   * @param targetInt - The target entity's interned int.
   * @param options - Optional limit/offset pagination.
   * @returns Verb ints, resolvable via {@link GraphIndexProvider.verbIntsToIds}.
   */
  getVerbIdsByTarget(targetInt: bigint, options?: { limit?: number; offset?: number }): Promise<bigint[]>
  /**
   * @description Batch reverse resolver: verb ints → verb-id strings. REQUIRED —
   * the provider owns the durable verb-int interning (Brainy keeps only a
   * bounded in-memory warm cache fed by `addVerb` returns and this resolver;
   * pure optimization, no durability role).
   * @param verbInts - Verb ints as returned by the read methods.
   * @returns One entry per input, order-preserving; `null` for unknown ints.
   */
  verbIntsToIds(verbInts: bigint[]): Promise<(string | null)[]>
  getVerbsBatchCached(verbIds: string[]): Promise<Map<string, GraphVerb>>

  /**
   * @description Index one verb. The coordinator resolves both endpoint ints
   * via `idMapper.getOrAssign` and mirrors them onto `verb.sourceInt` /
   * `verb.targetInt` before the call.
   * @param verb - The verb to index (endpoint UUIDs + derived `sourceInt`/`targetInt`).
   * @param sourceInt - The source entity's interned int.
   * @param targetInt - The target entity's interned int.
   * @returns The interned verb int for `verb.id` (feeds Brainy's warm cache).
   */
  addVerb(verb: GraphVerb, sourceInt: bigint, targetInt: bigint): Promise<bigint>
  /**
   * @description Remove one verb from the index by its id string. The verb's
   * interned int stays reserved (ints are never recycled within a generation).
   * @param verbId - The verb's UUID string.
   * @returns Resolves once the verb no longer appears in reads.
   */
  removeVerb(verbId: string): Promise<void>

  rebuild(): Promise<void>
  flush(): Promise<void>
  close(): Promise<void>
  size(): number

  getStats(): GraphIndexStats
  getRelationshipStats(): {
    totalRelationships: number
    relationshipsByType: Record<string, number>
    uniqueSourceNodes: number
    uniqueTargetNodes: number
    totalNodes: number
  }
  getRelationshipCountByType(type: string): number
  getTotalRelationshipCount(): number
  getAllRelationshipCounts(): Map<string, number>
}

/**
 * Optional capability interface for index providers that maintain versioned
 * (generation-aware) internal state — the provider-side half of Brainy 8.0's
 * generational MVCC contract. Implemented by native providers whose storage
 * engines keep immutable versions (e.g. LSM snapshots); Brainy's own JS
 * indexes do not implement it (they are rebuilt from the storage records,
 * which carry the versioning).
 *
 * **Detection.** Brainy feature-detects this interface on every registered
 * index provider (`'vector'`, `'metadataIndex'`, `'graphIndex'`): when a
 * provider exposes `pin`/`release` functions, Brainy calls them in lockstep
 * with `Db` lifecycle — `pin(g)` when a `Db` value pins generation `g`
 * (`brain.now()`, `brain.transact()`, `brain.asOf()`), `release(g)` when that
 * `Db` is released (explicitly or via the GC backstop). Pins are refcounted
 * on the Brainy side; a provider may receive multiple `pin(g)` calls for the
 * same generation and will receive exactly one matching `release(g)` per pin.
 *
 * **Consistency model (locked cross-team design).** Index providers are
 * *post-commit appliers*: the storage-record commit (atomic manifest rename)
 * is the source of truth, and provider index state is derived, applied after
 * the commit point. On open, a provider compares its own persisted
 * `generation()` against the store's committed generation and replays the
 * gap from the storage records (or requests a rebuild) — there are no
 * provider rollback hooks, because an uncommitted transaction is repaired at
 * the storage layer before any index is opened. The explicit pin/release
 * lifetime OVERRIDES any time-based snapshot retention the provider has
 * (e.g. an LSM snapshot TTL): a pinned generation must stay readable until
 * released, regardless of age.
 *
 * **Speculative reads.** `db.with(txData)` overlays are Brainy-side only —
 * providers are never asked to read uncommitted/speculative state; they
 * always serve committed generations.
 */
export interface VersionedIndexProvider {
  /**
   * @description The newest generation this provider's persisted index state
   * reflects. Brainy compares it against the storage layer's committed
   * generation on open to detect a replay gap (index behind storage after a
   * crash between commit and index apply).
   * @returns The provider's current generation as a `bigint` (u64 at the
   *   native boundary; Brainy's generation counter is a safe integer today).
   */
  generation(): bigint

  /**
   * @description True ⇒ the provider can serve consistent reads at
   * `generation` (segments retained). Combined with pin: pinning a VISIBLE
   * generation guarantees it stays servable until release. Pinning an
   * invisible generation is permitted (refcount-only) — Brainy serves that
   * generation from canonical storage instead.
   *
   * Brainy consults this at pin time (the read-routing rule: a `Db` at
   * generation `g` uses provider-accelerated reads when
   * `isGenerationVisible(g)` was true at pin time; otherwise canonical
   * generation records).
   * @param generation - The generation a `Db` value is about to pin.
   */
  isGenerationVisible(generation: bigint): boolean

  /**
   * @description Pin a generation: the provider must keep index state for
   * `generation` readable until the matching {@link VersionedIndexProvider.release}
   * call, overriding any time-based snapshot retention. Called once per
   * Brainy-side pin (refcounted upstream — expect balanced pin/release pairs).
   * @param generation - The generation a live `Db` value just pinned.
   */
  pin(generation: bigint): void

  /**
   * @description Release one pin on `generation`. After the last release the
   * provider may reclaim resources for that generation at its discretion.
   * @param generation - The generation being released (matches a prior
   *   {@link VersionedIndexProvider.pin} call).
   */
  release(generation: bigint): void
}

/**
 * @description Feature-detection guard for {@link VersionedIndexProvider}.
 * Brainy applies it to every registered index provider (vector, metadata,
 * graph) when a `Db` value pins or releases a generation: providers exposing
 * all four capability methods receive lockstep `pin`/`release` calls;
 * everything else (including Brainy's own JS indexes) is skipped.
 *
 * @param candidate - A registered index provider instance.
 * @returns Whether the candidate implements the versioned capability.
 */
export function isVersionedIndexProvider(
  candidate: unknown
): candidate is VersionedIndexProvider {
  if (candidate === null || typeof candidate !== 'object') {
    return false
  }
  const c = candidate as Record<string, unknown>
  return (
    typeof c.generation === 'function' &&
    typeof c.isGenerationVisible === 'function' &&
    typeof c.pin === 'function' &&
    typeof c.release === 'function'
  )
}

/**
 * The object returned by the `'vector'` provider factory — Brainy's vector
 * index contract. Implementations include Brainy's own JS HNSW index and any
 * native acceleration provider (e.g. cortex's Adaptive DiskANN).
 *
 * Brainy calls this surface via `this.index.*` plus the transactional add/remove
 * operations. `enableCOW`, `getItem`, and `setPersistMode` are intentionally
 * absent: Brainy guards each with feature-detection (`typeof x === 'function'`),
 * so they are optional and not part of the required contract (Brainy's own JS
 * HNSW index omits `setPersistMode`, for instance).
 *
 * **Provider key:** registered under `'vector'` — the only key Brainy
 * consults for the vector index. The pre-8.0 `'hnsw'` and `'diskann'` keys
 * are retired and never looked up.
 */
export interface VectorIndexProvider {
  addItem(item: VectorDocument): Promise<string>
  removeItem(id: string): Promise<boolean>
  search(
    queryVector: Vector,
    k?: number,
    filter?: (id: string) => Promise<boolean>,
    options?: { rerank?: { multiplier: number }; candidateIds?: string[] }
  ): Promise<Array<[string, number]>>
  size(): number
  clear(): void
  rebuild(options?: any): Promise<void>
  flush(): Promise<number>
  getPersistMode(): 'immediate' | 'deferred'
}


/**
 * The `'entityIdMapper'` provider — a drop-in for `EntityIdMapper`. Injected
 * into the TypeScript `MetadataIndexManager` when a native metadata index is
 * not also registered; that coordinator calls this full surface (incl.
 * `getAllIntIds`, the all-ids universe for negation / `exists:false` filters).
 */
export interface EntityIdMapperProvider {
  init(): Promise<void>
  getOrAssign(uuid: string): number
  getUuid(intId: number): string | undefined
  getInt(uuid: string): number | undefined
  remove(uuid: string): boolean
  flush(): Promise<void>
  clear(): Promise<void>
  getAllIntIds(): number[]
  intsIterableToUuids(ints: Iterable<number>): string[]
  readonly size: number
}

/**
 * The `'cache'` provider — a drop-in for `UnifiedCache`. Brainy installs it as
 * the global cache (`setGlobalCache`) and calls this surface via
 * `getGlobalCache()`.
 */
export interface CacheProvider {
  getSync(key: string): any | undefined
  set(key: string, data: any, type: 'vectors' | 'metadata' | 'embedding' | 'other', size: number, rebuildCost?: number): void
  delete(key: string): boolean
  deleteByPrefix(prefix: string): number
  clear(type?: 'vectors' | 'metadata' | 'embedding' | 'other'): void
}

// The `'embeddings'` / `'embedBatch'` providers are function-shaped and already
// typed by the existing `EmbeddingFunction` (see `coreTypes.ts`), which Brainy
// uses at the `getProvider('embeddings')` call site. No separate interface is
// added here to avoid a duplicate, unwired contract.

/**
 * The `'graph:compression'` provider — pure-function encode/decode for HNSW
 * connection lists as compact delta-varint byte sequences (cortex's
 * `encodeConnections` / `decodeConnections`).
 *
 * Brainy's `JsHnswVectorIndex` consumes this via a `ConnectionsCodec` that translates
 * UUIDs to stable int slots via the `EntityIdMapper`, encodes, and persists
 * the compressed bytes through the binary-blob primitive. On load, the blob
 * is fetched + decoded back into UUID sets — `setConnectionsCodec()` on
 * `JsHnswVectorIndex` is the injection point. Read path is dual-format: when no blob
 * exists for a node, the connections fall back to the legacy JSON-array path
 * embedded in `saveVectorIndexData`, so pre-2.4.0 indexes keep loading unchanged
 * and convergence to the compressed form happens lazily on next save.
 *
 * Activated only when the storage adapter exposes the binary-blob primitive
 * AND the metadata index resolves a stable idMapper. Cloud adapters that
 * lack a real local-path resolution still benefit, since the blob primitive
 * itself works across every adapter as of brainy 7.25.0.
 */
export interface GraphCompressionProvider {
  /** Encode a list of u32 ints to compact delta-varint bytes. Sorts internally. */
  encode(ids: number[]): Buffer
  /** Decode delta-varint bytes back to a u32 list. */
  decode(data: Buffer): number[]
}

/**
 * Storage adapter factory — plugins register these to provide
 * new storage backends that users reference by name.
 *
 * Example: A Redis plugin registers 'storage:redis', then users
 * can use `new Brainy({ storage: 'redis', redis: { host: '...' } })`
 */
export interface StorageAdapterFactory {
  create(config: Record<string, unknown>): StorageAdapter | Promise<StorageAdapter>
  name: string
}

/**
 * Plugin registry — manages plugin lifecycle and provider resolution.
 */
export class PluginRegistry {
  private plugins: Map<string, BrainyPlugin> = new Map()
  private providers: Map<string, unknown> = new Map()
  private activated: Set<string> = new Set()

  /**
   * Register a plugin manually.
   */
  register(plugin: BrainyPlugin): void {
    this.plugins.set(plugin.name, plugin)
  }


  /**
   * Activate all registered plugins.
   */
  async activateAll(context: BrainyPluginContext): Promise<string[]> {
    const activated: string[] = []

    for (const [name, plugin] of this.plugins) {
      if (this.activated.has(name)) continue

      try {
        const success = await plugin.activate(context)
        if (success) {
          this.activated.add(name)
          activated.push(name)
        }
      } catch (error) {
        console.warn(`[brainy] Plugin ${name} failed to activate:`, error)
      }
    }

    return activated
  }

  /**
   * Deactivate all plugins (called during close()).
   */
  async deactivateAll(): Promise<void> {
    for (const [name, plugin] of this.plugins) {
      if (!this.activated.has(name)) continue
      try {
        await plugin.deactivate?.()
        this.activated.delete(name)
      } catch {
        // Non-fatal
      }
    }
  }

  /**
   * Get a registered provider by key.
   */
  getProvider<T = unknown>(key: string): T | undefined {
    return this.providers.get(key) as T | undefined
  }

  /**
   * Check if a provider is registered.
   */
  hasProvider(key: string): boolean {
    return this.providers.has(key)
  }

  /**
   * Register a provider (called by plugins via BrainyPluginContext).
   */
  registerProvider(key: string, implementation: unknown): void {
    this.providers.set(key, implementation)
  }

  /**
   * Get a storage adapter factory by name.
   */
  getStorageFactory(name: string): StorageAdapterFactory | undefined {
    return this.providers.get(`storage:${name}`) as StorageAdapterFactory | undefined
  }

  /** Get active plugin names */
  getActivePlugins(): string[] {
    return [...this.activated]
  }

  /** Check if any plugins are active */
  hasActivePlugins(): boolean {
    return this.activated.size > 0
  }
}