brainy/src/hnsw/typeAwareHNSWIndex.ts

/**
 * Type-Aware HNSW Index - Phase 2 Billion-Scale Optimization
 *
 * Maintains separate HNSW graphs per entity type for massive memory savings:
 * - Memory @ 1B scale: PROJECTED 384GB → 50GB (-87% from architectural analysis, not yet benchmarked)
 * - Query speed: PROJECTED 10x faster for single-type queries (not yet benchmarked)
 * - Storage: Already type-first from Phase 1a
 *
 * Architecture:
 * - One HNSWIndex per NounType (42 total)
 * - Lazy initialization (indexes created on first use)
 * - Type routing for optimal performance
 * - Falls back to multi-type search when type unknown
 */

import { HNSWIndex } from './hnswIndex.js'
import {
  DistanceFunction,
  HNSWConfig,
  Vector,
  VectorDocument
} from '../coreTypes.js'
import { NounType, NOUN_TYPE_COUNT, TypeUtils } from '../types/graphTypes.js'
import { euclideanDistance } from '../utils/index.js'
import type { BaseStorage } from '../storage/baseStorage.js'
import { prodLog } from '../utils/logger.js'

// Default HNSW parameters (same as HNSWIndex)
const DEFAULT_CONFIG: HNSWConfig = {
  M: 16,
  efConstruction: 200,
  efSearch: 50,
  ml: 16
}

/**
 * Type-aware HNSW statistics
 */
export interface TypeAwareHNSWStats {
  totalNodes: number
  totalMemoryMB: number
  typeCount: number
  typeStats: Map<NounType, {
    nodeCount: number
    memoryMB: number
    maxLevel: number
    entryPointId: string | null
  }>
  memoryReductionPercent: number
  estimatedMonolithicMemoryMB: number
}

/**
 * TypeAwareHNSWIndex - Separate HNSW graphs per entity type
 *
 * Phase 2 of billion-scale optimization roadmap.
 * PROJECTED: Reduces HNSW memory by 87% @ billion scale (calculated from architecture, not yet benchmarked).
 */
export class TypeAwareHNSWIndex {
  // One HNSW index per noun type (lazy initialization)
  private indexes: Map<NounType, HNSWIndex> = new Map()

  // Configuration
  private config: HNSWConfig
  private distanceFunction: DistanceFunction
  private storage: BaseStorage | null
  private useParallelization: boolean
  private persistMode: 'immediate' | 'deferred'

  /**
   * Create a new TypeAwareHNSWIndex
   *
   * @param config HNSW configuration (M, efConstruction, efSearch, ml)
   * @param distanceFunction Distance function (default: euclidean)
   * @param options Additional options (storage, parallelization, persistMode)
   */
  constructor(
    config: Partial<HNSWConfig> = {},
    distanceFunction: DistanceFunction = euclideanDistance,
    options: { useParallelization?: boolean; storage?: BaseStorage; persistMode?: 'immediate' | 'deferred' } = {}
  ) {
    this.config = { ...DEFAULT_CONFIG, ...config }
    this.distanceFunction = distanceFunction
    this.storage = options.storage || null
    this.useParallelization =
      options.useParallelization !== undefined
        ? options.useParallelization
        : true
    this.persistMode = options.persistMode || 'immediate'

    prodLog.info('TypeAwareHNSWIndex initialized (Phase 2: Type-Aware HNSW)')
  }

  /**
   * Enable COW (Copy-on-Write) mode - Instant fork via shallow copy
   *
   * Propagates enableCOW() to all underlying type-specific HNSW indexes.
   * Each index performs O(1) shallow copy of its own data structures.
   *
   * @param parent - Parent TypeAwareHNSWIndex to copy from
   */
  public enableCOW(parent: TypeAwareHNSWIndex): void {
    // Shallow copy indexes Map
    this.indexes = new Map(parent.indexes)

    // Enable COW on each underlying type-specific index
    for (const [type, parentIndex] of parent.indexes.entries()) {
      const childIndex = new HNSWIndex(this.config, this.distanceFunction, {
        useParallelization: this.useParallelization,
        storage: this.storage || undefined,
        persistMode: this.persistMode
      })
      childIndex.enableCOW(parentIndex)
      this.indexes.set(type, childIndex)
    }

    prodLog.info(`TypeAwareHNSWIndex COW enabled: ${parent.indexes.size} type-specific indexes shallow copied`)
  }

  /**
   * Flush dirty HNSW data to storage for all type-specific indexes
   *
   * In deferred persistence mode, HNSW connections are tracked as dirty but not
   * immediately persisted. Call flush() to persist all pending changes across
   * all type-specific indexes.
   *
   * @returns Total number of nodes flushed across all indexes
   */
  public async flush(): Promise<number> {
    if (this.indexes.size === 0) {
      return 0
    }

    const flushPromises = Array.from(this.indexes.values()).map(index =>
      index.flush()
    )

    const results = await Promise.all(flushPromises)
    const totalFlushed = results.reduce((sum, count) => sum + count, 0)

    if (totalFlushed > 0) {
      prodLog.info(`[TypeAwareHNSW] Flushed ${totalFlushed} dirty nodes across ${this.indexes.size} type indexes`)
    }

    return totalFlushed
  }

  /**
   * Get the total number of dirty (unpersisted) nodes across all type-specific indexes
   */
  public getDirtyNodeCount(): number {
    let total = 0
    for (const index of this.indexes.values()) {
      total += index.getDirtyNodeCount()
    }
    return total
  }

  /**
   * Get the current persist mode
   */
  public getPersistMode(): 'immediate' | 'deferred' {
    return this.persistMode
  }

  /**
   * Get or create HNSW index for a specific type (lazy initialization)
   *
   * Indexes are created on-demand to save memory.
   * Only types with entities get an index.
   *
   * @param type The noun type
   * @returns HNSWIndex for this type
   */
  private getIndexForType(type: NounType): HNSWIndex {
    // Validate type is a valid NounType
    const typeIndex = TypeUtils.getNounIndex(type)
    if (typeIndex === undefined || typeIndex === null || typeIndex < 0) {
      throw new Error(
        `Invalid NounType: ${type}. Must be one of the 42 defined types.`
      )
    }

    if (!this.indexes.has(type)) {
      prodLog.info(`Creating HNSW index for type: ${type}`)

      const index = new HNSWIndex(this.config, this.distanceFunction, {
        useParallelization: this.useParallelization,
        storage: this.storage || undefined,
        persistMode: this.persistMode
      })

      this.indexes.set(type, index)
    }

    const index = this.indexes.get(type)
    if (!index) {
      throw new Error(
        `Unexpected: Index for type ${type} not found after creation`
      )
    }
    return index
  }

  /**
   * Add a vector to the type-aware index
   *
   * Routes to the correct type's HNSW graph.
   *
   * @param item Vector document to add
   * @param type The noun type (required for routing)
   * @returns The item ID
   */
  public async addItem(item: VectorDocument, type: NounType): Promise<string> {
    if (!item || !item.vector) {
      throw new Error(
        'Invalid VectorDocument: item or vector is null/undefined'
      )
    }
    if (!type) {
      throw new Error('Type is required for type-aware indexing')
    }

    const index = this.getIndexForType(type)
    return await index.addItem(item)
  }

  /**
   * Search for nearest neighbors (type-aware)
   *
   * **Single-type search** (fast path):
   * ```typescript
   * await index.search(queryVector, 10, 'person')
   * // Searches only person graph (100M nodes instead of 1B)
   * ```
   *
   * **Multi-type search**:
   * ```typescript
   * await index.search(queryVector, 10, ['person', 'organization'])
   * // Searches person + organization, merges results
   * ```
   *
   * **All-types search** (fallback):
   * ```typescript
   * await index.search(queryVector, 10)
   * // Searches all 42 graphs (slower but comprehensive)
   * ```
   *
   * @param queryVector Query vector
   * @param k Number of results
   * @param type Type or types to search (undefined = all types)
   * @param filter Optional filter function
   * @returns Array of [id, distance] tuples sorted by distance
   */
  public async search(
    queryVector: Vector,
    k: number = 10,
    type?: NounType | NounType[],
    filter?: (id: string) => Promise<boolean>,
    options?: { rerank?: { multiplier: number }; candidateIds?: string[] }
  ): Promise<Array<[string, number]>> {
    // Single-type search (fast path)
    if (type && typeof type === 'string') {
      const index = this.getIndexForType(type)
      return await index.search(queryVector, k, filter, options)
    }

    // Multi-type search (handle empty array edge case)
    if (type && Array.isArray(type) && type.length > 0) {
      return await this.searchMultipleTypes(queryVector, k, type, filter, options)
    }

    // All-types search (slowest path + empty array fallback)
    return await this.searchAllTypes(queryVector, k, filter, options)
  }

  /**
   * Search across multiple specific types
   *
   * @param queryVector Query vector
   * @param k Number of results
   * @param types Array of types to search
   * @param filter Optional filter function
   * @returns Merged and sorted results
   */
  private async searchMultipleTypes(
    queryVector: Vector,
    k: number,
    types: NounType[],
    filter?: (id: string) => Promise<boolean>,
    options?: { rerank?: { multiplier: number }; candidateIds?: string[] }
  ): Promise<Array<[string, number]>> {
    const allResults: Array<[string, number]> = []

    // Search each specified type
    for (const type of types) {
      if (this.indexes.has(type)) {
        const index = this.indexes.get(type)!
        const results = await index.search(queryVector, k, filter, options)
        allResults.push(...results)
      }
    }

    // Merge and sort by distance
    allResults.sort((a, b) => a[1] - b[1])

    // Return top k
    return allResults.slice(0, k)
  }

  /**
   * Search across all types (fallback for type-agnostic queries)
   *
   * This is the slowest path, but provides comprehensive results.
   * Used when type cannot be inferred from query.
   *
   * @param queryVector Query vector
   * @param k Number of results
   * @param filter Optional filter function
   * @returns Merged and sorted results from all types
   */
  private async searchAllTypes(
    queryVector: Vector,
    k: number,
    filter?: (id: string) => Promise<boolean>,
    options?: { rerank?: { multiplier: number }; candidateIds?: string[] }
  ): Promise<Array<[string, number]>> {
    const allResults: Array<[string, number]> = []

    // Search each type's graph
    for (const [type, index] of this.indexes.entries()) {
      const results = await index.search(queryVector, k, filter, options)
      allResults.push(...results)
    }

    // Merge and sort by distance
    allResults.sort((a, b) => a[1] - b[1])

    // Return top k
    return allResults.slice(0, k)
  }

  /**
   * Remove an item from the index
   *
   * @param id Item ID to remove
   * @param type The noun type (required for routing)
   * @returns True if item was removed, false if not found
   */
  public async removeItem(id: string, type: NounType): Promise<boolean> {
    const index = this.indexes.get(type)
    if (!index) {
      return false // Type has no index (no items ever added)
    }

    return await index.removeItem(id)
  }

  /**
   * Get total number of items across all types
   *
   * @returns Total item count
   */
  public size(): number {
    let total = 0
    for (const index of this.indexes.values()) {
      total += index.size()
    }
    return total
  }

  /**
   * Get number of items for a specific type
   *
   * @param type The noun type
   * @returns Item count for this type
   */
  public sizeForType(type: NounType): number {
    const index = this.indexes.get(type)
    return index ? index.size() : 0
  }

  /**
   * Clear all indexes
   */
  public clear(): void {
    for (const index of this.indexes.values()) {
      index.clear()
    }
    this.indexes.clear()
  }

  /**
   * Clear index for a specific type
   *
   * @param type The noun type to clear
   */
  public clearType(type: NounType): void {
    const index = this.indexes.get(type)
    if (index) {
      index.clear()
      this.indexes.delete(type)
    }
  }

  /**
   * Get configuration
   *
   * @returns HNSW configuration
   */
  public getConfig(): HNSWConfig {
    return { ...this.config }
  }

  /**
   * Get distance function
   *
   * @returns Distance function
   */
  public getDistanceFunction(): DistanceFunction {
    return this.distanceFunction
  }

  /**
   * Set parallelization (applies to all indexes)
   *
   * @param useParallelization Whether to use parallelization
   */
  public setUseParallelization(useParallelization: boolean): void {
    this.useParallelization = useParallelization
    for (const index of this.indexes.values()) {
      index.setUseParallelization(useParallelization)
    }
  }

  /**
   * Get parallelization setting
   *
   * @returns Whether parallelization is enabled
   */
  public getUseParallelization(): boolean {
    return this.useParallelization
  }

  /**
   * Rebuild HNSW indexes from storage (type-aware)
   *
   * CRITICAL: This implementation uses type-filtered pagination to avoid
   * loading ALL entities for each type (which would be 42 billion reads @ 1B scale).
   *
   * Can rebuild all types or specific types.
   * Much faster than rebuilding a monolithic index.
   *
   * @param options Rebuild options
   */
  public async rebuild(
    options: {
      types?: NounType[] // Rebuild specific types (undefined = all types)
      batchSize?: number // Entities per batch
      onProgress?: (type: NounType, loaded: number, total: number) => void
    } = {}
  ): Promise<void> {
    if (!this.storage) {
      prodLog.warn('TypeAwareHNSW rebuild skipped: no storage adapter')
      return
    }

    const batchSize = options.batchSize || 1000

    // Determine which types to rebuild
    const typesToRebuild = options.types || this.getAllNounTypes()

    prodLog.info(
      `Rebuilding ${typesToRebuild.length} type-aware HNSW indexes from persisted data...`
    )

    // Clear all indexes we're rebuilding
    for (const type of typesToRebuild) {
      const index = this.getIndexForType(type)
      ;(index as any).nouns.clear()
    }

    // Determine preloading strategy (adaptive caching) for entire dataset
    const stats = await this.storage.getStatistics()
    const entityCount = stats?.totalNodes || 0
    const vectorMemory = entityCount * 1536 // 384 dims × 4 bytes

    // Use first index's cache (they all share the same UnifiedCache)
    const firstIndex = this.getIndexForType(typesToRebuild[0])
    const cacheStats = (firstIndex as any).unifiedCache.getStats()
    const availableCache = cacheStats.maxSize * 0.80
    const shouldPreload = vectorMemory < availableCache

    if (shouldPreload) {
      prodLog.info(
        `HNSW: Preloading ${entityCount.toLocaleString()} vectors at init ` +
        `(${(vectorMemory / 1024 / 1024).toFixed(1)}MB < ${(availableCache / 1024 / 1024).toFixed(1)}MB cache)`
      )
    } else {
      prodLog.info(
        `HNSW: Adaptive caching for ${entityCount.toLocaleString()} vectors ` +
        `(${(vectorMemory / 1024 / 1024).toFixed(1)}MB > ${(availableCache / 1024 / 1024).toFixed(1)}MB cache) - loading on-demand`
      )
    }

    // Load ALL nouns ONCE and route to correct type indexes
    // This is O(N) instead of O(42*N) from the previous parallel approach
    let offset = 0  // Use offset-based pagination instead of cursor (bug fix for infinite loop)
    let hasMore = true
    let totalLoaded = 0
    const loadedByType = new Map<NounType, number>()

    while (hasMore) {
      const result: {
        items: Array<{ id: string; vector: number[]; type?: NounType }>
        hasMore: boolean
        nextCursor?: string
        totalCount?: number
      } = await (this.storage as any).getNounsWithPagination({
        limit: batchSize,
        offset  // Pass offset for proper pagination (previously passed cursor which was ignored)
      })

      // Route each noun to its type index
      for (const nounData of result.items) {
        try {
          // Use 'type' property from HNSWNounWithMetadata (not 'nounType')
          // Previously accessed wrong property, causing N+1 metadata fetches
          // getNounsWithPagination already includes type in response
          const nounType = nounData.type

          // Skip if type not in rebuild list
          if (!nounType || !typesToRebuild.includes(nounType as NounType)) {
            continue
          }

          // Get the index for this type
          const index = this.getIndexForType(nounType as NounType)

          // Load HNSW graph data
          const hnswData = await (this.storage as any).getHNSWData(nounData.id)
          if (!hnswData) {
            continue // No HNSW data
          }

          // Create noun with restored connections
          const noun = {
            id: nounData.id,
            vector: shouldPreload ? nounData.vector : [],
            connections: new Map(),
            level: hnswData.level
          }

          // Restore connections from storage
          for (const [levelStr, nounIds] of Object.entries(hnswData.connections)) {
            const level = parseInt(levelStr, 10)
            noun.connections.set(level, new Set<string>(nounIds as string[]))
          }

          // Add to type-specific index
          ;(index as any).nouns.set(nounData.id, noun)

          // Track high-level nodes
          if (noun.level >= 2 && noun.level <= (index as any).MAX_TRACKED_LEVELS) {
            if (!(index as any).highLevelNodes.has(noun.level)) {
              ;(index as any).highLevelNodes.set(noun.level, new Set())
            }
            ;(index as any).highLevelNodes.get(noun.level).add(nounData.id)
          }

          // Track progress
          loadedByType.set(nounType as NounType, (loadedByType.get(nounType as NounType) || 0) + 1)
          totalLoaded++

          if (options.onProgress && totalLoaded % 100 === 0) {
            options.onProgress(nounType as NounType, loadedByType.get(nounType as NounType) || 0, totalLoaded)
          }
        } catch (error) {
          prodLog.error(`Failed to restore HNSW data for ${nounData.id}:`, error)
        }
      }

      hasMore = result.hasMore
      offset += batchSize  // Increment offset for next page

      // Progress logging
      if (totalLoaded % 1000 === 0) {
        prodLog.info(`Progress: ${totalLoaded.toLocaleString()} entities loaded...`)
      }
    }

    // Restore entry points for each type
    for (const type of typesToRebuild) {
      const index = this.getIndexForType(type)
      let maxLevel = 0
      let entryPointId: string | null = null

      for (const [id, noun] of (index as any).nouns.entries()) {
        if (entryPointId === null || noun.level > maxLevel) {
          maxLevel = noun.level
          entryPointId = id
        }
      }

      // Recovery: if still null after loop but nouns exist
      if ((index as any).nouns.size > 0 && !entryPointId) {
        const { id: recoveredId, level: recoveredLevel } = (index as any).recoverEntryPointO1()
        entryPointId = recoveredId
        maxLevel = recoveredLevel
      }

      // Validation: if entry point doesn't exist in loaded nouns
      if (entryPointId && !(index as any).nouns.has(entryPointId)) {
        const { id: recoveredId, level: recoveredLevel } = (index as any).recoverEntryPointO1()
        entryPointId = recoveredId
        maxLevel = recoveredLevel
      }

      ;(index as any).entryPointId = entryPointId
      ;(index as any).maxLevel = maxLevel

      const loaded = loadedByType.get(type) || 0
      const cacheInfo = shouldPreload ? ' (vectors preloaded)' : ' (adaptive caching)'

      prodLog.info(
        `✅ Rebuilt ${type} index: ${loaded.toLocaleString()} entities, ` +
        `${maxLevel + 1} levels, entry point: ${entryPointId || 'none'}${cacheInfo}`
      )
    }

    prodLog.info(
      `✅ TypeAwareHNSW rebuild complete: ${this.size().toLocaleString()} total entities across ${this.indexes.size} types (loaded from persisted graph structure)`
    )
  }

  /**
   * Get comprehensive statistics
   *
   * Shows memory reduction compared to monolithic approach.
   *
   * @returns Type-aware HNSW statistics
   */
  public getStats(): TypeAwareHNSWStats {
    const typeStats = new Map<
      NounType,
      {
        nodeCount: number
        memoryMB: number
        maxLevel: number
        entryPointId: string | null
      }
    >()

    let totalNodes = 0
    let totalMemoryMB = 0

    // Collect stats from each type's index
    for (const [type, index] of this.indexes.entries()) {
      const cacheStats = index.getCacheStats()
      const nodeCount = index.size()
      const memoryMB = cacheStats.hnswCache.estimatedMemoryMB

      typeStats.set(type, {
        nodeCount,
        memoryMB,
        maxLevel: index.getMaxLevel(),
        entryPointId: index.getEntryPointId()
      })

      totalNodes += nodeCount
      totalMemoryMB += memoryMB
    }

    // Estimate monolithic memory (for comparison)
    // Monolithic would use ~384 bytes per entity @ 1B scale
    const estimatedMonolithicMemoryMB = (totalNodes * 384) / (1024 * 1024)

    // Calculate memory reduction
    const memoryReductionPercent =
      estimatedMonolithicMemoryMB > 0
        ? ((estimatedMonolithicMemoryMB - totalMemoryMB) /
            estimatedMonolithicMemoryMB) *
          100
        : 0

    return {
      totalNodes,
      totalMemoryMB: parseFloat(totalMemoryMB.toFixed(2)),
      typeCount: this.indexes.size,
      typeStats,
      memoryReductionPercent: parseFloat(memoryReductionPercent.toFixed(2)),
      estimatedMonolithicMemoryMB: parseFloat(
        estimatedMonolithicMemoryMB.toFixed(2)
      )
    }
  }

  /**
   * Get statistics for a specific type
   *
   * @param type The noun type
   * @returns Statistics for this type's index (null if no index)
   */
  public getStatsForType(
    type: NounType
  ): {
    nodeCount: number
    memoryMB: number
    maxLevel: number
    entryPointId: string | null
    cacheStats: any
  } | null {
    const index = this.indexes.get(type)
    if (!index) {
      return null
    }

    const cacheStats = index.getCacheStats()

    return {
      nodeCount: index.size(),
      memoryMB: cacheStats.hnswCache.estimatedMemoryMB,
      maxLevel: index.getMaxLevel(),
      entryPointId: index.getEntryPointId(),
      cacheStats
    }
  }

  /**
   * Get all noun types (for iteration)
   *
   * @returns Array of all noun types
   */
  private getAllNounTypes(): NounType[] {
    const types: NounType[] = []
    for (let i = 0; i < NOUN_TYPE_COUNT; i++) {
      types.push(TypeUtils.getNounFromIndex(i))
    }
    return types
  }

  /**
   * Get list of types that have indexes (have entities)
   *
   * @returns Array of types with indexes
   */
  public getActiveTypes(): NounType[] {
    return Array.from(this.indexes.keys())
  }
}