open-brainy/src/brainyData.ts

/**
 * BrainyData
 * Main class that provides the vector database functionality
 */

import { v4 as uuidv4 } from 'uuid'
import { HNSWIndex } from './hnsw/hnswIndex.js'
import {
  HNSWIndexOptimized,
  HNSWOptimizedConfig
} from './hnsw/hnswIndexOptimized.js'
import { createStorage } from './storage/opfsStorage.js'
import {
  DistanceFunction,
  GraphVerb,
  EmbeddingFunction,
  HNSWConfig,
  HNSWNoun,
  SearchResult,
  StorageAdapter,
  Vector,
  VectorDocument
} from './coreTypes.js'
import {
  cosineDistance,
  defaultEmbeddingFunction,
  defaultBatchEmbeddingFunction,
  euclideanDistance,
  cleanupWorkerPools
} from './utils/index.js'
import { NounType, VerbType, GraphNoun } from './types/graphTypes.js'
import {
  ServerSearchConduitAugmentation,
  createServerSearchAugmentations
} from './augmentations/serverSearchAugmentations.js'
import { WebSocketConnection } from './types/augmentations.js'
import { BrainyDataInterface } from './types/brainyDataInterface.js'

export interface BrainyDataConfig {
  /**
   * Vector dimensions (required if not using an embedding function that auto-detects dimensions)
   */
  dimensions?: number

  /**
   * HNSW index configuration
   */
  hnsw?: Partial<HNSWConfig>

  /**
   * Optimized HNSW index configuration
   * If provided, will use the optimized HNSW index instead of the standard one
   */
  hnswOptimized?: Partial<HNSWOptimizedConfig>

  /**
   * Distance function to use for similarity calculations
   */
  distanceFunction?: DistanceFunction

  /**
   * Custom storage adapter (if not provided, will use OPFS or memory storage)
   */
  storageAdapter?: StorageAdapter

  /**
   * Storage configuration options
   * These will be passed to createStorage if storageAdapter is not provided
   */
  storage?: {
    requestPersistentStorage?: boolean
    r2Storage?: {
      bucketName?: string
      accountId?: string
      accessKeyId?: string
      secretAccessKey?: string
    }
    s3Storage?: {
      bucketName?: string
      accessKeyId?: string
      secretAccessKey?: string
      region?: string
    }
    gcsStorage?: {
      bucketName?: string
      accessKeyId?: string
      secretAccessKey?: string
      endpoint?: string
    }
    customS3Storage?: {
      bucketName?: string
      accessKeyId?: string
      secretAccessKey?: string
      endpoint?: string
      region?: string
    }
    forceFileSystemStorage?: boolean
    forceMemoryStorage?: boolean
  }

  /**
   * Embedding function to convert data to vectors
   */
  embeddingFunction?: EmbeddingFunction

  /**
   * Set the database to read-only mode
   * When true, all write operations will throw an error
   */
  readOnly?: boolean

  /**
   * Remote server configuration for search operations
   */
  remoteServer?: {
    /**
     * WebSocket URL of the remote Brainy server
     */
    url: string

    /**
     * WebSocket protocols to use for the connection
     */
    protocols?: string | string[]

    /**
     * Whether to automatically connect to the remote server on initialization
     */
    autoConnect?: boolean
  }
}

export class BrainyData<T = any> implements BrainyDataInterface<T> {
  private index: HNSWIndex | HNSWIndexOptimized
  private storage: StorageAdapter | null = null
  private isInitialized = false
  private isInitializing = false
  private embeddingFunction: EmbeddingFunction
  private distanceFunction: DistanceFunction
  private requestPersistentStorage: boolean
  private readOnly: boolean
  private storageConfig: BrainyDataConfig['storage'] = {}
  private useOptimizedIndex: boolean = false
  private _dimensions: number

  // Remote server properties
  private remoteServerConfig: BrainyDataConfig['remoteServer'] | null = null
  private serverSearchConduit: ServerSearchConduitAugmentation | null = null
  private serverConnection: WebSocketConnection | null = null

  /**
   * Get the vector dimensions
   */
  public get dimensions(): number {
    return this._dimensions
  }

  /**
   * Get the maximum connections parameter from HNSW configuration
   */
  public get maxConnections(): number {
    const config = this.index.getConfig()
    return config.M || 16
  }

  /**
   * Get the efConstruction parameter from HNSW configuration
   */
  public get efConstruction(): number {
    const config = this.index.getConfig()
    return config.efConstruction || 200
  }

  /**
   * Create a new vector database
   */
  constructor(config: BrainyDataConfig = {}) {
    // Validate dimensions
    if (config.dimensions !== undefined && config.dimensions <= 0) {
      throw new Error('Dimensions must be a positive number')
    }

    // Set dimensions (default to 512 for embedding functions, or require explicit config)
    this._dimensions = config.dimensions || 512

    // Set distance function
    this.distanceFunction = config.distanceFunction || cosineDistance

    // Check if optimized HNSW index configuration is provided
    if (config.hnswOptimized) {
      // Initialize optimized HNSW index
      this.index = new HNSWIndexOptimized(
        config.hnswOptimized,
        this.distanceFunction,
        config.storageAdapter || null
      )
      this.useOptimizedIndex = true
    } else {
      // Initialize standard HNSW index
      this.index = new HNSWIndex(config.hnsw, this.distanceFunction)
    }

    // Set storage if provided, otherwise it will be initialized in init()
    this.storage = config.storageAdapter || null

    // Set embedding function if provided, otherwise use default
    this.embeddingFunction =
      config.embeddingFunction || defaultEmbeddingFunction

    // Set persistent storage request flag
    this.requestPersistentStorage =
      config.storage?.requestPersistentStorage || false

    // Set read-only flag
    this.readOnly = config.readOnly || false

    // Store storage configuration for later use in init()
    this.storageConfig = config.storage || {}

    // Store remote server configuration if provided
    if (config.remoteServer) {
      this.remoteServerConfig = config.remoteServer
    }
  }

  /**
   * Check if the database is in read-only mode and throw an error if it is
   * @throws Error if the database is in read-only mode
   */
  private checkReadOnly(): void {
    if (this.readOnly) {
      throw new Error(
        'Cannot perform write operation: database is in read-only mode'
      )
    }
  }

  /**
   * Initialize the database
   * Loads existing data from storage if available
   */
  public async init(): Promise<void> {
    if (this.isInitialized) {
      return
    }

    // Prevent recursive initialization
    if (this.isInitializing) {
      return
    }

    this.isInitializing = true

    try {
      // Pre-load the embedding model early to ensure it's always available
      // This helps prevent issues with the Universal Sentence Encoder not being loaded
      try {
        console.log('Pre-loading Universal Sentence Encoder model...')
        // Call embedding function directly to avoid circular dependency with embed()
        await this.embeddingFunction('')
        console.log('Universal Sentence Encoder model loaded successfully')
      } catch (embedError) {
        console.warn(
          'Failed to pre-load Universal Sentence Encoder:',
          embedError
        )

        // Try again with a retry mechanism
        console.log('Retrying Universal Sentence Encoder initialization...')
        try {
          // Wait a moment before retrying
          await new Promise((resolve) => setTimeout(resolve, 1000))

          // Try again with a different approach - use the non-threaded version
          // This is a fallback in case the threaded version fails
          const { createTensorFlowEmbeddingFunction } = await import(
            './utils/embedding.js'
          )
          const fallbackEmbeddingFunction = createTensorFlowEmbeddingFunction()

          // Test the fallback embedding function
          await fallbackEmbeddingFunction('')

          // If successful, replace the embedding function
          console.log(
            'Successfully loaded Universal Sentence Encoder with fallback method'
          )
          this.embeddingFunction = fallbackEmbeddingFunction
        } catch (retryError) {
          console.error(
            'All attempts to load Universal Sentence Encoder failed:',
            retryError
          )
          // Continue initialization even if embedding model fails to load
          // The application will need to handle missing embedding functionality
        }
      }

      // Initialize storage if not provided in constructor
      if (!this.storage) {
        // Combine storage config with requestPersistentStorage for backward compatibility
        const storageOptions = {
          ...this.storageConfig,
          requestPersistentStorage: this.requestPersistentStorage
        }

        this.storage = await createStorage(storageOptions)
      }

      // Initialize storage
      await this.storage!.init()

      // If using optimized index, set the storage adapter
      if (this.useOptimizedIndex && this.index instanceof HNSWIndexOptimized) {
        this.index.setStorage(this.storage!)
      }

      // Load all nouns from storage
      const nouns: HNSWNoun[] = await this.storage!.getAllNouns()

      // Clear the index and add all nouns
      this.index.clear()
      for (const noun of nouns) {
        // Check if the vector dimensions match the expected dimensions
        if (noun.vector.length !== this._dimensions) {
          console.warn(
            `Skipping noun ${noun.id} due to dimension mismatch: expected ${this._dimensions}, got ${noun.vector.length}`
          )
          // Optionally, you could delete the mismatched noun from storage
          // await this.storage!.deleteNoun(noun.id)
          continue
        }

        // Add to index
        await this.index.addItem({
          id: noun.id,
          vector: noun.vector
        })
      }

      // Connect to remote server if configured with autoConnect
      if (this.remoteServerConfig && this.remoteServerConfig.autoConnect) {
        try {
          await this.connectToRemoteServer(
            this.remoteServerConfig.url,
            this.remoteServerConfig.protocols
          )
        } catch (remoteError) {
          console.warn('Failed to auto-connect to remote server:', remoteError)
          // Continue initialization even if remote connection fails
        }
      }

      this.isInitialized = true
      this.isInitializing = false
    } catch (error) {
      console.error('Failed to initialize BrainyData:', error)
      this.isInitializing = false
      throw new Error(`Failed to initialize BrainyData: ${error}`)
    }
  }

  /**
   * Connect to a remote Brainy server for search operations
   * @param serverUrl WebSocket URL of the remote Brainy server
   * @param protocols Optional WebSocket protocols to use
   * @returns The connection object
   */
  public async connectToRemoteServer(
    serverUrl: string,
    protocols?: string | string[]
  ): Promise<WebSocketConnection> {
    await this.ensureInitialized()

    try {
      // Create server search augmentations
      const { conduit, connection } = await createServerSearchAugmentations(
        serverUrl,
        {
          protocols,
          localDb: this
        }
      )

      // Store the conduit and connection
      this.serverSearchConduit = conduit
      this.serverConnection = connection

      return connection
    } catch (error) {
      console.error('Failed to connect to remote server:', error)
      throw new Error(`Failed to connect to remote server: ${error}`)
    }
  }

  /**
   * Add a vector or data to the database
   * If the input is not a vector, it will be converted using the embedding function
   * @param vectorOrData Vector or data to add
   * @param metadata Optional metadata to associate with the vector
   * @param options Additional options
   * @returns The ID of the added vector
   */
  public async add(
    vectorOrData: Vector | any,
    metadata?: T,
    options: {
      forceEmbed?: boolean // Force using the embedding function even if input is a vector
      addToRemote?: boolean // Whether to also add to the remote server if connected
      id?: string // Optional ID to use instead of generating a new one
    } = {}
  ): Promise<string> {
    await this.ensureInitialized()

    // Check if database is in read-only mode
    this.checkReadOnly()

    try {
      let vector: Vector

      // Check if input is already a vector
      if (
        Array.isArray(vectorOrData) &&
        vectorOrData.every((item) => typeof item === 'number') &&
        !options.forceEmbed
      ) {
        // Input is already a vector
        vector = vectorOrData
      } else {
        // Input needs to be vectorized
        try {
          vector = await this.embeddingFunction(vectorOrData)
        } catch (embedError) {
          throw new Error(`Failed to vectorize data: ${embedError}`)
        }
      }

      // Check if vector is defined
      if (!vector) {
        throw new Error('Vector is undefined or null')
      }

      // Validate vector dimensions
      if (vector.length !== this._dimensions) {
        throw new Error(`Vector dimension mismatch: expected ${this._dimensions}, got ${vector.length}`)
      }

      // Use ID from options if it exists, otherwise from metadata, otherwise generate a new UUID
      const id =
        options.id ||
        (metadata && typeof metadata === 'object' && 'id' in metadata
          ? (metadata as any).id
          : uuidv4())

      // Add to index
      await this.index.addItem({ id, vector })

      // Get the noun from the index
      const noun = this.index.getNouns().get(id)

      if (!noun) {
        throw new Error(`Failed to retrieve newly created noun with ID ${id}`)
      }

      // Save noun to storage
      await this.storage!.saveNoun(noun)

      // Save metadata if provided
      if (metadata !== undefined) {
        // Validate noun type if metadata is for a GraphNoun
        if (metadata && typeof metadata === 'object' && 'noun' in metadata) {
          const nounType = (metadata as unknown as GraphNoun).noun

          // Check if the noun type is valid
          const isValidNounType = Object.values(NounType).includes(nounType)

          if (!isValidNounType) {
            console.warn(
              `Invalid noun type: ${nounType}. Falling back to GraphNoun.`
            )
            // Set a default noun type
            ;(metadata as unknown as GraphNoun).noun = NounType.Concept
          }
        }

        // Ensure metadata has the correct id field
        let metadataToSave = metadata
        if (metadata && typeof metadata === 'object') {
          metadataToSave = { ...metadata, id }
        }

        await this.storage!.saveMetadata(id, metadataToSave)
      }

      // If addToRemote is true and we're connected to a remote server, add to remote as well
      if (options.addToRemote && this.isConnectedToRemoteServer()) {
        try {
          await this.addToRemote(id, vector, metadata)
        } catch (remoteError) {
          console.warn(
            `Failed to add to remote server: ${remoteError}. Continuing with local add.`
          )
        }
      }

      return id
    } catch (error) {
      console.error('Failed to add vector:', error)
      throw new Error(`Failed to add vector: ${error}`)
    }
  }

  /**
   * Add a text item to the database with automatic embedding
   * This is a convenience method for adding text data with metadata
   * @param text Text data to add
   * @param metadata Metadata to associate with the text
   * @param options Additional options
   * @returns The ID of the added item
   */
  public async addItem(
    text: string,
    metadata?: T,
    options: {
      addToRemote?: boolean // Whether to also add to the remote server if connected
      id?: string // Optional ID to use instead of generating a new one
    } = {}
  ): Promise<string> {
    // Use the existing add method with forceEmbed to ensure text is embedded
    return this.add(text, metadata, { ...options, forceEmbed: true })
  }

  /**
   * Add data to both local and remote Brainy instances
   * @param vectorOrData Vector or data to add
   * @param metadata Optional metadata to associate with the vector
   * @param options Additional options
   * @returns The ID of the added vector
   */
  public async addToBoth(
    vectorOrData: Vector | any,
    metadata?: T,
    options: {
      forceEmbed?: boolean // Force using the embedding function even if input is a vector
    } = {}
  ): Promise<string> {
    // Check if connected to a remote server
    if (!this.isConnectedToRemoteServer()) {
      throw new Error(
        'Not connected to a remote server. Call connectToRemoteServer() first.'
      )
    }

    // Add to local with addToRemote option
    return this.add(vectorOrData, metadata, { ...options, addToRemote: true })
  }

  /**
   * Add a vector to the remote server
   * @param id ID of the vector to add
   * @param vector Vector to add
   * @param metadata Optional metadata to associate with the vector
   * @returns True if successful, false otherwise
   * @private
   */
  private async addToRemote(
    id: string,
    vector: Vector,
    metadata?: T
  ): Promise<boolean> {
    if (!this.isConnectedToRemoteServer()) {
      return false
    }

    try {
      if (!this.serverSearchConduit || !this.serverConnection) {
        throw new Error(
          'Server search conduit or connection is not initialized'
        )
      }

      // Add to remote server
      const addResult = await this.serverSearchConduit.addToBoth(
        this.serverConnection.connectionId,
        vector,
        metadata
      )

      if (!addResult.success) {
        throw new Error(`Remote add failed: ${addResult.error}`)
      }

      return true
    } catch (error) {
      console.error('Failed to add to remote server:', error)
      throw new Error(`Failed to add to remote server: ${error}`)
    }
  }

  /**
   * Add multiple vectors or data items to the database
   * @param items Array of items to add
   * @param options Additional options
   * @returns Array of IDs for the added items
   */
  public async addBatch(
    items: Array<{
      vectorOrData: Vector | any
      metadata?: T
    }>,
    options: {
      forceEmbed?: boolean // Force using the embedding function even if input is a vector
      addToRemote?: boolean // Whether to also add to the remote server if connected
      concurrency?: number // Maximum number of concurrent operations (default: 4)
      batchSize?: number // Maximum number of items to process in a single batch (default: 50)
    } = {}
  ): Promise<string[]> {
    await this.ensureInitialized()

    // Check if database is in read-only mode
    this.checkReadOnly()

    // Default concurrency to 4 if not specified
    const concurrency = options.concurrency || 4

    // Default batch size to 50 if not specified
    const batchSize = options.batchSize || 50

    try {
      // Process items in batches to control concurrency and memory usage
      const ids: string[] = []
      const itemsToProcess = [...items] // Create a copy to avoid modifying the original array

      while (itemsToProcess.length > 0) {
        // Take up to 'batchSize' items to process in a batch
        const batch = itemsToProcess.splice(0, batchSize)

        // Separate items that are already vectors from those that need embedding
        const vectorItems: Array<{
          vectorOrData: Vector
          metadata?: T
          index: number
        }> = []

        const textItems: Array<{
          text: string
          metadata?: T
          index: number
        }> = []

        // Categorize items
        batch.forEach((item, index) => {
          if (
            Array.isArray(item.vectorOrData) &&
            item.vectorOrData.every((val) => typeof val === 'number') &&
            !options.forceEmbed
          ) {
            // Item is already a vector
            vectorItems.push({
              vectorOrData: item.vectorOrData,
              metadata: item.metadata,
              index
            })
          } else if (typeof item.vectorOrData === 'string') {
            // Item is text that needs embedding
            textItems.push({
              text: item.vectorOrData,
              metadata: item.metadata,
              index
            })
          } else {
            // For now, treat other types as text
            // In a more complete implementation, we might handle other types differently
            const textRepresentation = String(item.vectorOrData)
            textItems.push({
              text: textRepresentation,
              metadata: item.metadata,
              index
            })
          }
        })

        // Process vector items (already embedded)
        const vectorPromises = vectorItems.map((item) =>
          this.add(item.vectorOrData, item.metadata, options)
        )

        // Process text items in a single batch embedding operation
        let textPromises: Promise<string>[] = []
        if (textItems.length > 0) {
          // Extract just the text for batch embedding
          const texts = textItems.map((item) => item.text)

          // Perform batch embedding
          const embeddings = await defaultBatchEmbeddingFunction(texts)

          // Add each item with its embedding
          textPromises = textItems.map((item, i) =>
            this.add(embeddings[i], item.metadata, {
              ...options,
              forceEmbed: false
            })
          )
        }

        // Combine all promises
        const batchResults = await Promise.all([
          ...vectorPromises,
          ...textPromises
        ])

        // Add the results to our ids array
        ids.push(...batchResults)
      }

      return ids
    } catch (error) {
      console.error('Failed to add batch of items:', error)
      throw new Error(`Failed to add batch of items: ${error}`)
    }
  }

  /**
   * Add multiple vectors or data items to both local and remote databases
   * @param items Array of items to add
   * @param options Additional options
   * @returns Array of IDs for the added items
   */
  public async addBatchToBoth(
    items: Array<{
      vectorOrData: Vector | any
      metadata?: T
    }>,
    options: {
      forceEmbed?: boolean // Force using the embedding function even if input is a vector
      concurrency?: number // Maximum number of concurrent operations (default: 4)
    } = {}
  ): Promise<string[]> {
    // Check if connected to a remote server
    if (!this.isConnectedToRemoteServer()) {
      throw new Error(
        'Not connected to a remote server. Call connectToRemoteServer() first.'
      )
    }

    // Add to local with addToRemote option
    return this.addBatch(items, { ...options, addToRemote: true })
  }

  /**
   * Search for similar vectors within specific noun types
   * @param queryVectorOrData Query vector or data to search for
   * @param k Number of results to return
   * @param nounTypes Array of noun types to search within, or null to search all
   * @param options Additional options
   * @returns Array of search results
   */
  public async searchByNounTypes(
    queryVectorOrData: Vector | any,
    k: number = 10,
    nounTypes: string[] | null = null,
    options: {
      forceEmbed?: boolean // Force using the embedding function even if input is a vector
    } = {}
  ): Promise<SearchResult<T>[]> {
    if (!this.isInitialized) {
      throw new Error('BrainyData must be initialized before searching. Call init() first.')
    }

    try {
      let queryVector: Vector

      // Check if input is already a vector
      if (
        Array.isArray(queryVectorOrData) &&
        queryVectorOrData.every((item) => typeof item === 'number') &&
        !options.forceEmbed
      ) {
        // Input is already a vector
        queryVector = queryVectorOrData
      } else {
        // Input needs to be vectorized
        try {
          queryVector = await this.embeddingFunction(queryVectorOrData)
        } catch (embedError) {
          throw new Error(`Failed to vectorize query data: ${embedError}`)
        }
      }

      // Check if query vector is defined
      if (!queryVector) {
        throw new Error('Query vector is undefined or null')
      }

      // If no noun types specified, search all nouns
      if (!nounTypes || nounTypes.length === 0) {
        // Search in the index
        const results = await this.index.search(queryVector, k)

        // Get metadata for each result
        const searchResults: SearchResult<T>[] = []

        for (const [id, score] of results) {
          const noun = this.index.getNouns().get(id)
          if (!noun) {
            continue
          }

          const metadata = await this.storage!.getMetadata(id)

          searchResults.push({
            id,
            score,
            vector: noun.vector,
            metadata: metadata as T | undefined
          })
        }

        return searchResults
      } else {
        // Get nouns for each noun type in parallel
        const nounPromises = nounTypes.map((nounType) =>
          this.storage!.getNounsByNounType(nounType)
        )
        const nounArrays = await Promise.all(nounPromises)

        // Combine all nouns
        const nouns: HNSWNoun[] = []
        for (const nounArray of nounArrays) {
          nouns.push(...nounArray)
        }

        // Calculate distances for each noun
        const results: Array<[string, number]> = []
        for (const noun of nouns) {
          const distance = this.index.getDistanceFunction()(
            queryVector,
            noun.vector
          )
          results.push([noun.id, distance])
        }

        // Sort by distance (ascending)
        results.sort((a, b) => a[1] - b[1])

        // Take top k results
        const topResults = results.slice(0, k)

        // Get metadata for each result
        const searchResults: SearchResult<T>[] = []

        for (const [id, score] of topResults) {
          const noun = nouns.find((n) => n.id === id)
          if (!noun) {
            continue
          }

          const metadata = await this.storage!.getMetadata(id)

          searchResults.push({
            id,
            score,
            vector: noun.vector,
            metadata: metadata as T | undefined
          })
        }

        return searchResults
      }
    } catch (error) {
      console.error('Failed to search vectors by noun types:', error)
      throw new Error(`Failed to search vectors by noun types: ${error}`)
    }
  }

  /**
   * Search for similar vectors
   * @param queryVectorOrData Query vector or data to search for
   * @param k Number of results to return
   * @param options Additional options
   * @returns Array of search results
   */
  public async search(
    queryVectorOrData: Vector | any,
    k: number = 10,
    options: {
      forceEmbed?: boolean // Force using the embedding function even if input is a vector
      nounTypes?: string[] // Optional array of noun types to search within
      includeVerbs?: boolean // Whether to include associated GraphVerbs in the results
      searchMode?: 'local' | 'remote' | 'combined' // Where to search: local, remote, or both
      searchVerbs?: boolean // Whether to search for verbs directly instead of nouns
      verbTypes?: string[] // Optional array of verb types to search within or filter by
      searchConnectedNouns?: boolean // Whether to search for nouns connected by verbs
      verbDirection?: 'outgoing' | 'incoming' | 'both' // Direction of verbs to consider when searching connected nouns
    } = {}
  ): Promise<SearchResult<T>[]> {
    if (!this.isInitialized) {
      throw new Error('BrainyData must be initialized before searching. Call init() first.')
    }
    // If searching for verbs directly
    if (options.searchVerbs) {
      const verbResults = await this.searchVerbs(queryVectorOrData, k, {
        forceEmbed: options.forceEmbed,
        verbTypes: options.verbTypes
      })

      // Convert verb results to SearchResult format
      return verbResults.map((verb) => ({
        id: verb.id,
        score: verb.similarity,
        vector: verb.embedding || [],
        metadata: {
          verb: verb.verb,
          source: verb.source,
          target: verb.target,
          ...verb.data
        } as unknown as T
      }))
    }

    // If searching for nouns connected by verbs
    if (options.searchConnectedNouns) {
      return this.searchNounsByVerbs(queryVectorOrData, k, {
        forceEmbed: options.forceEmbed,
        verbTypes: options.verbTypes,
        direction: options.verbDirection
      })
    }

    // If a specific search mode is specified, use the appropriate search method
    if (options.searchMode === 'local') {
      return this.searchLocal(queryVectorOrData, k, options)
    } else if (options.searchMode === 'remote') {
      return this.searchRemote(queryVectorOrData, k, options)
    } else if (options.searchMode === 'combined') {
      return this.searchCombined(queryVectorOrData, k, options)
    }

    // Default behavior (backward compatible): search locally
    return this.searchLocal(queryVectorOrData, k, options)
  }

  /**
   * Search the local database for similar vectors
   * @param queryVectorOrData Query vector or data to search for
   * @param k Number of results to return
   * @param options Additional options
   * @returns Array of search results
   */
  public async searchLocal(
    queryVectorOrData: Vector | any,
    k: number = 10,
    options: {
      forceEmbed?: boolean // Force using the embedding function even if input is a vector
      nounTypes?: string[] // Optional array of noun types to search within
      includeVerbs?: boolean // Whether to include associated GraphVerbs in the results
    } = {}
  ): Promise<SearchResult<T>[]> {
    if (!this.isInitialized) {
      throw new Error('BrainyData must be initialized before searching. Call init() first.')
    }
    // If input is a string and not a vector, automatically vectorize it
    let queryToUse = queryVectorOrData
    if (typeof queryVectorOrData === 'string' && !options.forceEmbed) {
      queryToUse = await this.embed(queryVectorOrData)
      options.forceEmbed = false // Already embedded, don't force again
    }

    // If noun types are specified, use searchByNounTypes
    let searchResults
    if (options.nounTypes && options.nounTypes.length > 0) {
      searchResults = await this.searchByNounTypes(
        queryToUse,
        k,
        options.nounTypes,
        {
          forceEmbed: options.forceEmbed
        }
      )
    } else {
      // Otherwise, search all GraphNouns
      searchResults = await this.searchByNounTypes(queryToUse, k, null, {
        forceEmbed: options.forceEmbed
      })
    }

    // If includeVerbs is true, retrieve associated GraphVerbs for each result
    if (options.includeVerbs && this.storage) {
      for (const result of searchResults) {
        try {
          // Get outgoing verbs for this noun
          const outgoingVerbs = await this.storage.getVerbsBySource(result.id)

          // Get incoming verbs for this noun
          const incomingVerbs = await this.storage.getVerbsByTarget(result.id)

          // Combine all verbs
          const allVerbs = [...outgoingVerbs, ...incomingVerbs]

          // Add verbs to the result metadata
          if (!result.metadata) {
            result.metadata = {} as T
          }

          // Add the verbs to the metadata
          ;(result.metadata as Record<string, any>).associatedVerbs = allVerbs
        } catch (error) {
          console.warn(`Failed to retrieve verbs for noun ${result.id}:`, error)
        }
      }
    }

    return searchResults
  }

  /**
   * Find entities similar to a given entity ID
   * @param id ID of the entity to find similar entities for
   * @param options Additional options
   * @returns Array of search results with similarity scores
   */
  public async findSimilar(
    id: string,
    options: {
      limit?: number // Number of results to return
      nounTypes?: string[] // Optional array of noun types to search within
      includeVerbs?: boolean // Whether to include associated GraphVerbs in the results
      searchMode?: 'local' | 'remote' | 'combined' // Where to search: local, remote, or both
    } = {}
  ): Promise<SearchResult<T>[]> {
    await this.ensureInitialized()

    // Get the entity by ID
    const entity = await this.get(id)
    if (!entity) {
      throw new Error(`Entity with ID ${id} not found`)
    }

    // Use the entity's vector to search for similar entities
    const k = (options.limit || 10) + 1 // Add 1 to account for the original entity
    const searchResults = await this.search(entity.vector, k, {
      forceEmbed: false,
      nounTypes: options.nounTypes,
      includeVerbs: options.includeVerbs,
      searchMode: options.searchMode
    })

    // Filter out the original entity and limit to the requested number
    return searchResults
      .filter((result) => result.id !== id)
      .slice(0, options.limit || 10)
  }

  /**
   * Get a vector by ID
   */
  public async get(id: string): Promise<VectorDocument<T> | null> {
    await this.ensureInitialized()

    try {
      // Get noun from index
      const noun = this.index.getNouns().get(id)
      if (!noun) {
        return null
      }

      // Get metadata
      const metadata = await this.storage!.getMetadata(id)

      return {
        id,
        vector: noun.vector,
        metadata: metadata as T | undefined
      }
    } catch (error) {
      console.error(`Failed to get vector ${id}:`, error)
      throw new Error(`Failed to get vector ${id}: ${error}`)
    }
  }

  /**
   * Get all nouns in the database
   * @returns Array of vector documents
   */
  public async getAllNouns(): Promise<VectorDocument<T>[]> {
    await this.ensureInitialized()

    try {
      const nouns = this.index.getNouns()
      const result: VectorDocument<T>[] = []

      for (const [id, noun] of nouns.entries()) {
        const metadata = await this.storage!.getMetadata(id)
        result.push({
          id,
          vector: noun.vector,
          metadata: metadata as T | undefined
        })
      }

      return result
    } catch (error) {
      console.error('Failed to get all nouns:', error)
      throw new Error(`Failed to get all nouns: ${error}`)
    }
  }

  /**
   * Delete a vector by ID
   */
  public async delete(id: string): Promise<boolean> {
    await this.ensureInitialized()

    // Check if database is in read-only mode
    this.checkReadOnly()

    try {
      // Remove from index
      const removed = this.index.removeItem(id)
      if (!removed) {
        return false
      }

      // Remove from storage
      await this.storage!.deleteNoun(id)

      // Try to remove metadata (ignore errors)
      try {
        await this.storage!.saveMetadata(id, null)
      } catch (error) {
        // Ignore
      }

      return true
    } catch (error) {
      console.error(`Failed to delete vector ${id}:`, error)
      throw new Error(`Failed to delete vector ${id}: ${error}`)
    }
  }

  /**
   * Update metadata for a vector
   */
  public async updateMetadata(id: string, metadata: T): Promise<boolean> {
    await this.ensureInitialized()

    // Check if database is in read-only mode
    this.checkReadOnly()

    try {
      // Check if a vector exists
      const noun = this.index.getNouns().get(id)
      if (!noun) {
        return false
      }

      // Validate noun type if metadata is for a GraphNoun
      if (metadata && typeof metadata === 'object' && 'noun' in metadata) {
        const nounType = (metadata as unknown as GraphNoun).noun

        // Check if the noun type is valid
        const isValidNounType = Object.values(NounType).includes(nounType)

        if (!isValidNounType) {
          console.warn(
            `Invalid noun type: ${nounType}. Falling back to GraphNoun.`
          )
          // Set a default noun type
          ;(metadata as unknown as GraphNoun).noun = NounType.Concept
        }
      }

      // Update metadata
      await this.storage!.saveMetadata(id, metadata)

      return true
    } catch (error) {
      console.error(`Failed to update metadata for vector ${id}:`, error)
      throw new Error(`Failed to update metadata for vector ${id}: ${error}`)
    }
  }

  /**
   * Create a relationship between two entities
   * This is a convenience wrapper around addVerb
   */
  public async relate(
    sourceId: string,
    targetId: string,
    relationType: string,
    metadata?: any
  ): Promise<string> {
    return this.addVerb(sourceId, targetId, undefined, {
      type: relationType,
      metadata: metadata
    })
  }

  /**
   * Add a verb between two nouns
   * If metadata is provided and vector is not, the metadata will be vectorized using the embedding function
   */
  public async addVerb(
    sourceId: string,
    targetId: string,
    vector?: Vector,
    options: {
      type?: string
      weight?: number
      metadata?: any
      forceEmbed?: boolean // Force using the embedding function for metadata even if vector is provided
      id?: string // Optional ID to use instead of generating a new one
    } = {}
  ): Promise<string> {
    await this.ensureInitialized()

    // Check if database is in read-only mode
    this.checkReadOnly()

    try {
      // Check if source and target nouns exist
      const sourceNoun = this.index.getNouns().get(sourceId)
      const targetNoun = this.index.getNouns().get(targetId)

      if (!sourceNoun) {
        throw new Error(`Source noun with ID ${sourceId} not found`)
      }

      if (!targetNoun) {
        throw new Error(`Target noun with ID ${targetId} not found`)
      }

      // Use provided ID or generate a new one
      const id = options.id || uuidv4()

      let verbVector: Vector

      // If metadata is provided and no vector is provided or forceEmbed is true, vectorize the metadata
      if (options.metadata && (!vector || options.forceEmbed)) {
        try {
          // Extract a string representation from metadata for embedding
          let textToEmbed: string
          if (typeof options.metadata === 'string') {
            textToEmbed = options.metadata
          } else if (
            options.metadata.description &&
            typeof options.metadata.description === 'string'
          ) {
            textToEmbed = options.metadata.description
          } else {
            // Convert to JSON string as fallback
            textToEmbed = JSON.stringify(options.metadata)
          }

          // Ensure textToEmbed is a string
          if (typeof textToEmbed !== 'string') {
            textToEmbed = String(textToEmbed)
          }

          verbVector = await this.embeddingFunction(textToEmbed)
        } catch (embedError) {
          throw new Error(`Failed to vectorize verb metadata: ${embedError}`)
        }
      } else {
        // Use a provided vector or average of source and target vectors
        if (vector) {
          verbVector = vector
        } else {
          // Ensure both source and target vectors have the same dimension
          if (
            !sourceNoun.vector ||
            !targetNoun.vector ||
            sourceNoun.vector.length === 0 ||
            targetNoun.vector.length === 0 ||
            sourceNoun.vector.length !== targetNoun.vector.length
          ) {
            throw new Error(
              `Cannot average vectors: source or target vector is invalid or dimensions don't match`
            )
          }

          // Average the vectors
          verbVector = sourceNoun.vector.map(
            (val, i) => (val + targetNoun.vector[i]) / 2
          )
        }
      }

      // Validate verb type if provided
      let verbType = options.type
      if (verbType) {
        // Check if the verb type is valid
        const isValidVerbType = Object.values(VerbType).includes(
          verbType as VerbType
        )

        if (!isValidVerbType) {
          console.warn(
            `Invalid verb type: ${verbType}. Using RelatedTo as default.`
          )
          // Set a default verb type
          verbType = VerbType.RelatedTo
        }
      }

      // Create verb
      const verb: GraphVerb = {
        id,
        vector: verbVector,
        connections: new Map(),
        sourceId,
        targetId,
        type: verbType,
        weight: options.weight,
        metadata: options.metadata
      }

      // Add to index
      await this.index.addItem({ id, vector: verbVector })

      // Get the noun from the index
      const indexNoun = this.index.getNouns().get(id)

      if (!indexNoun) {
        throw new Error(
          `Failed to retrieve newly created verb noun with ID ${id}`
        )
      }

      // Update verb connections from index
      verb.connections = indexNoun.connections

      // Save verb to storage
      await this.storage!.saveVerb(verb)

      return id
    } catch (error) {
      console.error('Failed to add verb:', error)
      throw new Error(`Failed to add verb: ${error}`)
    }
  }

  /**
   * Get a verb by ID
   */
  public async getVerb(id: string): Promise<GraphVerb | null> {
    await this.ensureInitialized()

    try {
      return await this.storage!.getVerb(id)
    } catch (error) {
      console.error(`Failed to get verb ${id}:`, error)
      throw new Error(`Failed to get verb ${id}: ${error}`)
    }
  }

  /**
   * Get all verbs
   */
  public async getAllVerbs(): Promise<GraphVerb[]> {
    await this.ensureInitialized()

    try {
      return await this.storage!.getAllVerbs()
    } catch (error) {
      console.error('Failed to get all verbs:', error)
      throw new Error(`Failed to get all verbs: ${error}`)
    }
  }

  /**
   * Get verbs by source noun ID
   */
  public async getVerbsBySource(sourceId: string): Promise<GraphVerb[]> {
    await this.ensureInitialized()

    try {
      return await this.storage!.getVerbsBySource(sourceId)
    } catch (error) {
      console.error(`Failed to get verbs by source ${sourceId}:`, error)
      throw new Error(`Failed to get verbs by source ${sourceId}: ${error}`)
    }
  }

  /**
   * Get verbs by target noun ID
   */
  public async getVerbsByTarget(targetId: string): Promise<GraphVerb[]> {
    await this.ensureInitialized()

    try {
      return await this.storage!.getVerbsByTarget(targetId)
    } catch (error) {
      console.error(`Failed to get verbs by target ${targetId}:`, error)
      throw new Error(`Failed to get verbs by target ${targetId}: ${error}`)
    }
  }

  /**
   * Get verbs by type
   */
  public async getVerbsByType(type: string): Promise<GraphVerb[]> {
    await this.ensureInitialized()

    try {
      return await this.storage!.getVerbsByType(type)
    } catch (error) {
      console.error(`Failed to get verbs by type ${type}:`, error)
      throw new Error(`Failed to get verbs by type ${type}: ${error}`)
    }
  }

  /**
   * Delete a verb
   */
  public async deleteVerb(id: string): Promise<boolean> {
    await this.ensureInitialized()

    // Check if database is in read-only mode
    this.checkReadOnly()

    try {
      // Remove from index
      const removed = this.index.removeItem(id)
      if (!removed) {
        return false
      }

      // Remove from storage
      await this.storage!.deleteVerb(id)

      return true
    } catch (error) {
      console.error(`Failed to delete verb ${id}:`, error)
      throw new Error(`Failed to delete verb ${id}: ${error}`)
    }
  }

  /**
   * Clear the database
   */
  public async clear(): Promise<void> {
    await this.ensureInitialized()

    // Check if database is in read-only mode
    this.checkReadOnly()

    try {
      // Clear index
      this.index.clear()

      // Clear storage
      await this.storage!.clear()
    } catch (error) {
      console.error('Failed to clear vector database:', error)
      throw new Error(`Failed to clear vector database: ${error}`)
    }
  }

  /**
   * Get the number of vectors in the database
   */
  public size(): number {
    return this.index.size()
  }

  /**
   * Check if the database is in read-only mode
   * @returns True if the database is in read-only mode, false otherwise
   */
  public isReadOnly(): boolean {
    return this.readOnly
  }

  /**
   * Set the database to read-only mode
   * @param readOnly True to set the database to read-only mode, false to allow writes
   */
  public setReadOnly(readOnly: boolean): void {
    this.readOnly = readOnly
  }

  /**
   * Embed text or data into a vector using the same embedding function used by this instance
   * This allows clients to use the same TensorFlow Universal Sentence Encoder throughout their application
   *
   * @param data Text or data to embed
   * @returns A promise that resolves to the embedded vector
   */
  public async embed(data: string | string[]): Promise<Vector> {
    await this.ensureInitialized()

    try {
      return await this.embeddingFunction(data)
    } catch (error) {
      console.error('Failed to embed data:', error)
      throw new Error(`Failed to embed data: ${error}`)
    }
  }

  /**
   * Search for verbs by type and/or vector similarity
   * @param queryVectorOrData Query vector or data to search for
   * @param k Number of results to return
   * @param options Additional options
   * @returns Array of verbs with similarity scores
   */
  public async searchVerbs(
    queryVectorOrData: Vector | any,
    k: number = 10,
    options: {
      forceEmbed?: boolean // Force using the embedding function even if input is a vector
      verbTypes?: string[] // Optional array of verb types to search within
    } = {}
  ): Promise<Array<GraphVerb & { similarity: number }>> {
    await this.ensureInitialized()

    try {
      let queryVector: Vector

      // Check if input is already a vector
      if (
        Array.isArray(queryVectorOrData) &&
        queryVectorOrData.every((item) => typeof item === 'number') &&
        !options.forceEmbed
      ) {
        // Input is already a vector
        queryVector = queryVectorOrData
      } else {
        // Input needs to be vectorized
        try {
          queryVector = await this.embeddingFunction(queryVectorOrData)
        } catch (embedError) {
          throw new Error(`Failed to vectorize query data: ${embedError}`)
        }
      }

      // Get verbs to search through
      let verbs: GraphVerb[] = []

      // If verb types are specified, get verbs of those types
      if (options.verbTypes && options.verbTypes.length > 0) {
        // Get verbs for each verb type in parallel
        const verbPromises = options.verbTypes.map((verbType) =>
          this.getVerbsByType(verbType)
        )
        const verbArrays = await Promise.all(verbPromises)

        // Combine all verbs
        for (const verbArray of verbArrays) {
          verbs.push(...verbArray)
        }
      } else {
        // Get all verbs
        verbs = await this.storage!.getAllVerbs()
      }

      // Filter out verbs without embeddings
      verbs = verbs.filter(
        (verb) => verb.embedding && verb.embedding.length > 0
      )

      // Calculate similarity for each verb
      const results: Array<GraphVerb & { similarity: number }> = []
      for (const verb of verbs) {
        if (verb.embedding) {
          const distance = this.index.getDistanceFunction()(
            queryVector,
            verb.embedding
          )
          results.push({
            ...verb,
            similarity: distance
          })
        }
      }

      // Sort by similarity (ascending distance)
      results.sort((a, b) => a.similarity - b.similarity)

      // Take top k results
      return results.slice(0, k)
    } catch (error) {
      console.error('Failed to search verbs:', error)
      throw new Error(`Failed to search verbs: ${error}`)
    }
  }

  /**
   * Search for nouns connected by specific verb types
   * @param queryVectorOrData Query vector or data to search for
   * @param k Number of results to return
   * @param options Additional options
   * @returns Array of search results
   */
  public async searchNounsByVerbs(
    queryVectorOrData: Vector | any,
    k: number = 10,
    options: {
      forceEmbed?: boolean // Force using the embedding function even if input is a vector
      verbTypes?: string[] // Optional array of verb types to filter by
      direction?: 'outgoing' | 'incoming' | 'both' // Direction of verbs to consider
    } = {}
  ): Promise<SearchResult<T>[]> {
    await this.ensureInitialized()

    try {
      // First, search for nouns
      const nounResults = await this.searchByNounTypes(
        queryVectorOrData,
        k * 2, // Get more results initially to account for filtering
        null,
        { forceEmbed: options.forceEmbed }
      )

      // If no verb types specified, return the noun results directly
      if (!options.verbTypes || options.verbTypes.length === 0) {
        return nounResults.slice(0, k)
      }

      // For each noun, get connected nouns through specified verb types
      const connectedNounIds = new Set<string>()
      const direction = options.direction || 'both'

      for (const result of nounResults) {
        // Get verbs connected to this noun
        let connectedVerbs: GraphVerb[] = []

        if (direction === 'outgoing' || direction === 'both') {
          // Get outgoing verbs
          const outgoingVerbs = await this.storage!.getVerbsBySource(result.id)
          connectedVerbs.push(...outgoingVerbs)
        }

        if (direction === 'incoming' || direction === 'both') {
          // Get incoming verbs
          const incomingVerbs = await this.storage!.getVerbsByTarget(result.id)
          connectedVerbs.push(...incomingVerbs)
        }

        // Filter by verb types if specified
        if (options.verbTypes && options.verbTypes.length > 0) {
          connectedVerbs = connectedVerbs.filter(
            (verb) => verb.verb && options.verbTypes!.includes(verb.verb)
          )
        }

        // Add connected noun IDs to the set
        for (const verb of connectedVerbs) {
          if (verb.source && verb.source !== result.id) {
            connectedNounIds.add(verb.source)
          }
          if (verb.target && verb.target !== result.id) {
            connectedNounIds.add(verb.target)
          }
        }
      }

      // Get the connected nouns
      const connectedNouns: SearchResult<T>[] = []
      for (const id of connectedNounIds) {
        try {
          const noun = this.index.getNouns().get(id)
          if (noun) {
            const metadata = await this.storage!.getMetadata(id)

            // Calculate similarity score
            let queryVector: Vector
            if (
              Array.isArray(queryVectorOrData) &&
              queryVectorOrData.every((item) => typeof item === 'number') &&
              !options.forceEmbed
            ) {
              queryVector = queryVectorOrData
            } else {
              queryVector = await this.embeddingFunction(queryVectorOrData)
            }

            const distance = this.index.getDistanceFunction()(
              queryVector,
              noun.vector
            )

            connectedNouns.push({
              id,
              score: distance,
              vector: noun.vector,
              metadata: metadata as T | undefined
            })
          }
        } catch (error) {
          console.warn(`Failed to retrieve noun ${id}:`, error)
        }
      }

      // Sort by similarity score
      connectedNouns.sort((a, b) => a.score - b.score)

      // Return top k results
      return connectedNouns.slice(0, k)
    } catch (error) {
      console.error('Failed to search nouns by verbs:', error)
      throw new Error(`Failed to search nouns by verbs: ${error}`)
    }
  }

  /**
   * Search for similar documents using a text query
   * This is a convenience method that embeds the query text and performs a search
   *
   * @param query Text query to search for
   * @param k Number of results to return
   * @param options Additional options
   * @returns Array of search results
   */
  public async searchText(
    query: string,
    k: number = 10,
    options: {
      nounTypes?: string[]
      includeVerbs?: boolean
      searchMode?: 'local' | 'remote' | 'combined'
    } = {}
  ): Promise<SearchResult<T>[]> {
    await this.ensureInitialized()

    try {
      // Embed the query text
      const queryVector = await this.embed(query)

      // Search using the embedded vector
      return await this.search(queryVector, k, {
        nounTypes: options.nounTypes,
        includeVerbs: options.includeVerbs,
        searchMode: options.searchMode
      })
    } catch (error) {
      console.error('Failed to search with text query:', error)
      throw new Error(`Failed to search with text query: ${error}`)
    }
  }

  /**
   * Search a remote Brainy server for similar vectors
   * @param queryVectorOrData Query vector or data to search for
   * @param k Number of results to return
   * @param options Additional options
   * @returns Array of search results
   */
  public async searchRemote(
    queryVectorOrData: Vector | any,
    k: number = 10,
    options: {
      forceEmbed?: boolean // Force using the embedding function even if input is a vector
      nounTypes?: string[] // Optional array of noun types to search within
      includeVerbs?: boolean // Whether to include associated GraphVerbs in the results
      storeResults?: boolean // Whether to store the results in the local database (default: true)
    } = {}
  ): Promise<SearchResult<T>[]> {
    await this.ensureInitialized()

    // Check if connected to a remote server
    if (!this.isConnectedToRemoteServer()) {
      throw new Error(
        'Not connected to a remote server. Call connectToRemoteServer() first.'
      )
    }

    try {
      // If input is a string, convert it to a query string for the server
      let query: string
      if (typeof queryVectorOrData === 'string') {
        query = queryVectorOrData
      } else {
        // For vectors, we need to embed them as a string query
        // This is a simplification - ideally we would send the vector directly
        query = 'vector-query' // Placeholder, would need a better approach for vector queries
      }

      if (!this.serverSearchConduit || !this.serverConnection) {
        throw new Error(
          'Server search conduit or connection is not initialized'
        )
      }

      // Search the remote server
      const searchResult = await this.serverSearchConduit.searchServer(
        this.serverConnection.connectionId,
        query,
        k
      )

      if (!searchResult.success) {
        throw new Error(`Remote search failed: ${searchResult.error}`)
      }

      return searchResult.data as SearchResult<T>[]
    } catch (error) {
      console.error('Failed to search remote server:', error)
      throw new Error(`Failed to search remote server: ${error}`)
    }
  }

  /**
   * Search both local and remote Brainy instances, combining the results
   * @param queryVectorOrData Query vector or data to search for
   * @param k Number of results to return
   * @param options Additional options
   * @returns Array of search results
   */
  public async searchCombined(
    queryVectorOrData: Vector | any,
    k: number = 10,
    options: {
      forceEmbed?: boolean // Force using the embedding function even if input is a vector
      nounTypes?: string[] // Optional array of noun types to search within
      includeVerbs?: boolean // Whether to include associated GraphVerbs in the results
      localFirst?: boolean // Whether to search local first (default: true)
    } = {}
  ): Promise<SearchResult<T>[]> {
    await this.ensureInitialized()

    // Check if connected to a remote server
    if (!this.isConnectedToRemoteServer()) {
      // If not connected to a remote server, just search locally
      return this.searchLocal(queryVectorOrData, k, options)
    }

    try {
      // Default to searching local first
      const localFirst = options.localFirst !== false

      if (localFirst) {
        // Search local first
        const localResults = await this.searchLocal(
          queryVectorOrData,
          k,
          options
        )

        // If we have enough local results, return them
        if (localResults.length >= k) {
          return localResults
        }

        // Otherwise, search remote for additional results
        const remoteResults = await this.searchRemote(
          queryVectorOrData,
          k - localResults.length,
          { ...options, storeResults: true }
        )

        // Combine results, removing duplicates
        const combinedResults = [...localResults]
        const localIds = new Set(localResults.map((r) => r.id))

        for (const result of remoteResults) {
          if (!localIds.has(result.id)) {
            combinedResults.push(result)
          }
        }

        return combinedResults
      } else {
        // Search remote first
        const remoteResults = await this.searchRemote(queryVectorOrData, k, {
          ...options,
          storeResults: true
        })

        // If we have enough remote results, return them
        if (remoteResults.length >= k) {
          return remoteResults
        }

        // Otherwise, search local for additional results
        const localResults = await this.searchLocal(
          queryVectorOrData,
          k - remoteResults.length,
          options
        )

        // Combine results, removing duplicates
        const combinedResults = [...remoteResults]
        const remoteIds = new Set(remoteResults.map((r) => r.id))

        for (const result of localResults) {
          if (!remoteIds.has(result.id)) {
            combinedResults.push(result)
          }
        }

        return combinedResults
      }
    } catch (error) {
      console.error('Failed to perform combined search:', error)
      throw new Error(`Failed to perform combined search: ${error}`)
    }
  }

  /**
   * Check if the instance is connected to a remote server
   * @returns True if connected to a remote server, false otherwise
   */
  public isConnectedToRemoteServer(): boolean {
    return !!(this.serverSearchConduit && this.serverConnection)
  }

  /**
   * Disconnect from the remote server
   * @returns True if successfully disconnected, false if not connected
   */
  public async disconnectFromRemoteServer(): Promise<boolean> {
    if (!this.isConnectedToRemoteServer()) {
      return false
    }

    try {
      if (!this.serverSearchConduit || !this.serverConnection) {
        throw new Error(
          'Server search conduit or connection is not initialized'
        )
      }

      // Close the WebSocket connection
      await this.serverSearchConduit.closeWebSocket(
        this.serverConnection.connectionId
      )

      // Clear the connection information
      this.serverSearchConduit = null
      this.serverConnection = null

      return true
    } catch (error) {
      console.error('Failed to disconnect from remote server:', error)
      throw new Error(`Failed to disconnect from remote server: ${error}`)
    }
  }

  /**
   * Ensure the database is initialized
   */
  private async ensureInitialized(): Promise<void> {
    if (this.isInitialized) {
      return
    }

    if (this.isInitializing) {
      // If initialization is already in progress, wait for it to complete
      // by polling the isInitialized flag
      let attempts = 0
      const maxAttempts = 100 // Prevent infinite loop
      const delay = 50 // ms

      while (
        this.isInitializing &&
        !this.isInitialized &&
        attempts < maxAttempts
      ) {
        await new Promise((resolve) => setTimeout(resolve, delay))
        attempts++
      }

      if (!this.isInitialized) {
        // If still not initialized after waiting, try to initialize again
        await this.init()
      }
    } else {
      // Normal case - not initialized and not initializing
      await this.init()
    }
  }

  /**
   * Get information about the current storage usage and capacity
   * @returns Object containing the storage type, used space, quota, and additional details
   */
  public async status(): Promise<{
    type: string
    used: number
    quota: number | null
    details?: Record<string, any>
  }> {
    await this.ensureInitialized()

    if (!this.storage) {
      return {
        type: 'any',
        used: 0,
        quota: null,
        details: { error: 'Storage not initialized' }
      }
    }

    try {
      // Check if the storage adapter has a getStorageStatus method
      if (typeof this.storage.getStorageStatus !== 'function') {
        // If not, determine the storage type based on the constructor name
        const storageType = this.storage.constructor.name
          .toLowerCase()
          .replace('storage', '')
        return {
          type: storageType || 'any',
          used: 0,
          quota: null,
          details: {
            error: 'Storage adapter does not implement getStorageStatus method',
            storageAdapter: this.storage.constructor.name,
            indexSize: this.size()
          }
        }
      }

      // Get storage status from the storage adapter
      const storageStatus = await this.storage.getStorageStatus()

      // Add index information to the details
      let indexInfo: Record<string, any> = {
        indexSize: this.size()
      }

      // Add optimized index information if using optimized index
      if (this.useOptimizedIndex && this.index instanceof HNSWIndexOptimized) {
        const optimizedIndex = this.index as HNSWIndexOptimized
        indexInfo = {
          ...indexInfo,
          optimized: true,
          memoryUsage: optimizedIndex.getMemoryUsage(),
          productQuantization: optimizedIndex.getUseProductQuantization(),
          diskBasedIndex: optimizedIndex.getUseDiskBasedIndex()
        }
      } else {
        indexInfo.optimized = false
      }

      // Ensure all required fields are present
      return {
        type: storageStatus.type || 'any',
        used: storageStatus.used || 0,
        quota: storageStatus.quota || null,
        details: {
          ...(storageStatus.details || {}),
          index: indexInfo
        }
      }
    } catch (error) {
      console.error('Failed to get storage status:', error)

      // Determine the storage type based on the constructor name
      const storageType = this.storage.constructor.name
        .toLowerCase()
        .replace('storage', '')

      return {
        type: storageType || 'any',
        used: 0,
        quota: null,
        details: {
          error: String(error),
          storageAdapter: this.storage.constructor.name,
          indexSize: this.size()
        }
      }
    }
  }

  /**
   * Shut down the database and clean up resources
   * This should be called when the database is no longer needed
   */
  public async shutDown(): Promise<void> {
    try {
      // Disconnect from remote server if connected
      if (this.isConnectedToRemoteServer()) {
        await this.disconnectFromRemoteServer()
      }

      // Clean up worker pools to release resources
      cleanupWorkerPools()

      // Additional cleanup could be added here in the future

      this.isInitialized = false
    } catch (error) {
      console.error('Failed to shut down BrainyData:', error)
      throw new Error(`Failed to shut down BrainyData: ${error}`)
    }
  }

  /**
   * Backup all data from the database to a JSON-serializable format
   * @returns Object containing all nouns, verbs, noun types, verb types, HNSW index, and other related data
   *
   * The HNSW index data includes:
   * - entryPointId: The ID of the entry point for the graph
   * - maxLevel: The maximum level in the hierarchical structure
   * - dimension: The dimension of the vectors
   * - config: Configuration parameters for the HNSW algorithm
   * - connections: A serialized representation of the connections between nouns
   */
  public async backup(): Promise<{
    nouns: VectorDocument<T>[]
    verbs: GraphVerb[]
    nounTypes: string[]
    verbTypes: string[]
    version: string
    hnswIndex?: {
      entryPointId: string | null
      maxLevel: number
      dimension: number | null
      config: HNSWConfig
      connections: Record<string, Record<string, string[]>>
    }
  }> {
    await this.ensureInitialized()

    try {
      // Get all nouns
      const nouns = await this.getAllNouns()

      // Get all verbs
      const verbs = await this.getAllVerbs()

      // Get all noun types
      const nounTypes = Object.values(NounType)

      // Get all verb types
      const verbTypes = Object.values(VerbType)

      // Get HNSW index data
      const hnswIndexData = {
        entryPointId: this.index.getEntryPointId(),
        maxLevel: this.index.getMaxLevel(),
        dimension: this.index.getDimension(),
        config: this.index.getConfig(),
        connections: {} as Record<string, Record<string, string[]>>
      }

      // Convert Map<number, Set<string>> to a serializable format
      const indexNouns = this.index.getNouns()
      for (const [id, noun] of indexNouns.entries()) {
        hnswIndexData.connections[id] = {}
        for (const [level, connections] of noun.connections.entries()) {
          hnswIndexData.connections[id][level] = Array.from(connections)
        }
      }

      // Return the data with version information
      return {
        nouns,
        verbs,
        nounTypes,
        verbTypes,
        hnswIndex: hnswIndexData,
        version: '1.0.0' // Version of the backup format
      }
    } catch (error) {
      console.error('Failed to backup data:', error)
      throw new Error(`Failed to backup data: ${error}`)
    }
  }

  /**
   * Import sparse data into the database
   * @param data The sparse data to import
   *             If vectors are not present for nouns, they will be created using the embedding function
   * @param options Import options
   * @returns Object containing counts of imported items
   */
  public async importSparseData(
    data: {
      nouns: VectorDocument<T>[]
      verbs: GraphVerb[]
      nounTypes?: string[]
      verbTypes?: string[]
      hnswIndex?: {
        entryPointId: string | null
        maxLevel: number
        dimension: number | null
        config: HNSWConfig
        connections: Record<string, Record<string, string[]>>
      }
      version: string
    },
    options: {
      clearExisting?: boolean
    } = {}
  ): Promise<{
    nounsRestored: number
    verbsRestored: number
  }> {
    return this.restore(data, options)
  }

  /**
   * Restore data into the database from a previously backed up format
   * @param data The data to restore, in the format returned by backup()
   *             This can include HNSW index data if it was included in the backup
   *             If vectors are not present for nouns, they will be created using the embedding function
   * @param options Restore options
   * @returns Object containing counts of restored items
   */
  public async restore(
    data: {
      nouns: VectorDocument<T>[]
      verbs: GraphVerb[]
      nounTypes?: string[]
      verbTypes?: string[]
      hnswIndex?: {
        entryPointId: string | null
        maxLevel: number
        dimension: number | null
        config: HNSWConfig
        connections: Record<string, Record<string, string[]>>
      }
      version: string
    },
    options: {
      clearExisting?: boolean
    } = {}
  ): Promise<{
    nounsRestored: number
    verbsRestored: number
  }> {
    await this.ensureInitialized()

    // Check if database is in read-only mode
    this.checkReadOnly()

    try {
      // Clear existing data if requested
      if (options.clearExisting) {
        await this.clear()
      }

      // Validate the data format
      if (!data || !data.nouns || !data.verbs || !data.version) {
        throw new Error('Invalid restore data format')
      }

      // Log additional data if present
      if (data.nounTypes) {
        console.log(`Found ${data.nounTypes.length} noun types in restore data`)
      }

      if (data.verbTypes) {
        console.log(`Found ${data.verbTypes.length} verb types in restore data`)
      }

      if (data.hnswIndex) {
        console.log('Found HNSW index data in backup')
      }

      // Restore nouns
      let nounsRestored = 0
      for (const noun of data.nouns) {
        try {
          // Check if the noun has a vector
          if (!noun.vector || noun.vector.length === 0) {
            // If no vector, create one using the embedding function
            if (
              noun.metadata &&
              typeof noun.metadata === 'object' &&
              'text' in noun.metadata
            ) {
              // If the metadata has a text field, use it for embedding
              noun.vector = await this.embeddingFunction(noun.metadata.text)
            } else {
              // Otherwise, use the entire metadata for embedding
              noun.vector = await this.embeddingFunction(noun.metadata)
            }
          }

          // Add the noun with its vector and metadata
          await this.add(noun.vector, noun.metadata, { id: noun.id })
          nounsRestored++
        } catch (error) {
          console.error(`Failed to restore noun ${noun.id}:`, error)
          // Continue with other nouns
        }
      }

      // Restore verbs
      let verbsRestored = 0
      for (const verb of data.verbs) {
        try {
          // Check if the verb has a vector
          if (!verb.vector || verb.vector.length === 0) {
            // If no vector, create one using the embedding function
            if (
              verb.metadata &&
              typeof verb.metadata === 'object' &&
              'text' in verb.metadata
            ) {
              // If the metadata has a text field, use it for embedding
              verb.vector = await this.embeddingFunction(verb.metadata.text)
            } else {
              // Otherwise, use the entire metadata for embedding
              verb.vector = await this.embeddingFunction(verb.metadata)
            }
          }

          // Add the verb
          await this.addVerb(verb.sourceId, verb.targetId, verb.vector, {
            id: verb.id,
            type: verb.metadata?.verb || VerbType.RelatedTo,
            metadata: verb.metadata
          })
          verbsRestored++
        } catch (error) {
          console.error(`Failed to restore verb ${verb.id}:`, error)
          // Continue with other verbs
        }
      }

      // If HNSW index data is provided and we've restored nouns, reconstruct the index
      if (data.hnswIndex && nounsRestored > 0) {
        try {
          console.log('Reconstructing HNSW index from backup data...')

          // Create a new index with the restored configuration
          this.index = new HNSWIndex(
            data.hnswIndex.config,
            this.distanceFunction
          )

          // Re-add all nouns to the index
          for (const noun of data.nouns) {
            if (noun.vector && noun.vector.length > 0) {
              await this.index.addItem({ id: noun.id, vector: noun.vector })
            }
          }

          console.log('HNSW index reconstruction complete')
        } catch (error) {
          console.error('Failed to reconstruct HNSW index:', error)
          console.log('Continuing with standard restore process...')
        }
      }

      return {
        nounsRestored,
        verbsRestored
      }
    } catch (error) {
      console.error('Failed to restore data:', error)
      throw new Error(`Failed to restore data: ${error}`)
    }
  }

  /**
   * Generate a random graph of data with typed nouns and verbs for testing and experimentation
   * @param options Configuration options for the random graph
   * @returns Object containing the IDs of the generated nouns and verbs
   */
  public async generateRandomGraph(
    options: {
      nounCount?: number // Number of nouns to generate (default: 10)
      verbCount?: number // Number of verbs to generate (default: 20)
      nounTypes?: NounType[] // Types of nouns to generate (default: all types)
      verbTypes?: VerbType[] // Types of verbs to generate (default: all types)
      clearExisting?: boolean // Whether to clear existing data before generating (default: false)
      seed?: string // Seed for random generation (default: random)
    } = {}
  ): Promise<{
    nounIds: string[]
    verbIds: string[]
  }> {
    await this.ensureInitialized()

    // Check if database is in read-only mode
    this.checkReadOnly()

    // Set default options
    const nounCount = options.nounCount || 10
    const verbCount = options.verbCount || 20
    const nounTypes = options.nounTypes || Object.values(NounType)
    const verbTypes = options.verbTypes || Object.values(VerbType)
    const clearExisting = options.clearExisting || false

    // Clear existing data if requested
    if (clearExisting) {
      await this.clear()
    }

    try {
      // Generate random nouns
      const nounIds: string[] = []
      const nounDescriptions: Record<string, string> = {
        [NounType.Person]: 'A person with unique characteristics',
        [NounType.Place]: 'A location with specific attributes',
        [NounType.Thing]: 'An object with distinct properties',
        [NounType.Event]: 'An occurrence with temporal aspects',
        [NounType.Concept]: 'An abstract idea or notion',
        [NounType.Content]: 'A piece of content or information',
        [NounType.Group]: 'A collection of related entities',
        [NounType.List]: 'An ordered sequence of items',
        [NounType.Category]: 'A classification or grouping'
      }

      for (let i = 0; i < nounCount; i++) {
        // Select a random noun type
        const nounType = nounTypes[Math.floor(Math.random() * nounTypes.length)]

        // Generate a random label
        const label = `Random ${nounType} ${i + 1}`

        // Create metadata
        const metadata = {
          noun: nounType,
          label,
          description: nounDescriptions[nounType] || `A random ${nounType}`,
          randomAttributes: {
            value: Math.random() * 100,
            priority: Math.floor(Math.random() * 5) + 1,
            tags: [`tag-${i % 5}`, `category-${i % 3}`]
          }
        }

        // Add the noun
        const id = await this.add(metadata.description, metadata as T)
        nounIds.push(id)
      }

      // Generate random verbs between nouns
      const verbIds: string[] = []
      const verbDescriptions: Record<string, string> = {
        [VerbType.AttributedTo]: 'Attribution relationship',
        [VerbType.Controls]: 'Control relationship',
        [VerbType.Created]: 'Creation relationship',
        [VerbType.Earned]: 'Achievement relationship',
        [VerbType.Owns]: 'Ownership relationship',
        [VerbType.MemberOf]: 'Membership relationship',
        [VerbType.RelatedTo]: 'General relationship',
        [VerbType.WorksWith]: 'Collaboration relationship',
        [VerbType.FriendOf]: 'Friendship relationship',
        [VerbType.ReportsTo]: 'Reporting relationship',
        [VerbType.Supervises]: 'Supervision relationship',
        [VerbType.Mentors]: 'Mentorship relationship'
      }

      for (let i = 0; i < verbCount; i++) {
        // Select random source and target nouns
        const sourceIndex = Math.floor(Math.random() * nounIds.length)
        let targetIndex = Math.floor(Math.random() * nounIds.length)

        // Ensure source and target are different
        while (targetIndex === sourceIndex && nounIds.length > 1) {
          targetIndex = Math.floor(Math.random() * nounIds.length)
        }

        const sourceId = nounIds[sourceIndex]
        const targetId = nounIds[targetIndex]

        // Select a random verb type
        const verbType = verbTypes[Math.floor(Math.random() * verbTypes.length)]

        // Create metadata
        const metadata = {
          verb: verbType,
          description:
            verbDescriptions[verbType] || `A random ${verbType} relationship`,
          weight: Math.random(),
          confidence: Math.random(),
          randomAttributes: {
            strength: Math.random() * 100,
            duration: Math.floor(Math.random() * 365) + 1,
            tags: [`relation-${i % 5}`, `strength-${i % 3}`]
          }
        }

        // Add the verb
        const id = await this.addVerb(sourceId, targetId, undefined, {
          type: verbType,
          weight: metadata.weight,
          metadata
        })

        verbIds.push(id)
      }

      return {
        nounIds,
        verbIds
      }
    } catch (error) {
      console.error('Failed to generate random graph:', error)
      throw new Error(`Failed to generate random graph: ${error}`)
    }
  }
}

// Export distance functions for convenience
export {
  euclideanDistance,
  cosineDistance,
  manhattanDistance,
  dotProductDistance
} from './utils/index.js'