brainy/src/storage/adapters/batchS3Operations.ts

/**
 * Enhanced Batch S3 Operations for High-Performance Vector Retrieval
 * Implements optimized batch operations to reduce S3 API calls and latency
 */

import { HNSWNoun, HNSWVerb } from '../../coreTypes.js'

// S3 client types - dynamically imported
type S3Client = any
type GetObjectCommand = any
type ListObjectsV2Command = any

export interface BatchRetrievalOptions {
  maxConcurrency?: number
  prefetchSize?: number
  useS3Select?: boolean
  compressionEnabled?: boolean
}

export interface BatchResult<T> {
  items: Map<string, T>
  errors: Map<string, Error>
  statistics: {
    totalRequested: number
    totalRetrieved: number
    totalErrors: number
    duration: number
    apiCalls: number
  }
}

/**
 * High-performance batch operations for S3-compatible storage
 * Optimizes retrieval patterns for HNSW search operations
 */
export class BatchS3Operations {
  private s3Client: S3Client
  private bucketName: string
  private options: BatchRetrievalOptions

  constructor(
    s3Client: S3Client,
    bucketName: string,
    options: BatchRetrievalOptions = {}
  ) {
    this.s3Client = s3Client
    this.bucketName = bucketName
    this.options = {
      maxConcurrency: 50, // AWS S3 rate limit friendly
      prefetchSize: 100,
      useS3Select: false,
      compressionEnabled: false,
      ...options
    }
  }

  /**
   * Batch retrieve HNSW nodes with intelligent prefetching
   */
  public async batchGetNodes(
    nodeIds: string[],
    prefix: string = 'nodes/'
  ): Promise<BatchResult<HNSWNoun>> {
    const startTime = Date.now()
    const result: BatchResult<HNSWNoun> = {
      items: new Map(),
      errors: new Map(),
      statistics: {
        totalRequested: nodeIds.length,
        totalRetrieved: 0,
        totalErrors: 0,
        duration: 0,
        apiCalls: 0
      }
    }

    if (nodeIds.length === 0) {
      result.statistics.duration = Date.now() - startTime
      return result
    }

    // Use different strategies based on request size
    if (nodeIds.length <= 10) {
      // Small batch - use parallel GetObject
      await this.parallelGetObjects(nodeIds, prefix, result)
    } else if (nodeIds.length <= 1000) {
      // Medium batch - use chunked parallel with prefetching
      await this.chunkedParallelGet(nodeIds, prefix, result)
    } else {
      // Large batch - use S3 list-based approach with filtering
      await this.listBasedBatchGet(nodeIds, prefix, result)
    }

    result.statistics.duration = Date.now() - startTime
    return result
  }

  /**
   * Parallel GetObject operations for small batches
   */
  private async parallelGetObjects<T>(
    ids: string[],
    prefix: string,
    result: BatchResult<T>
  ): Promise<void> {
    const { GetObjectCommand } = await import('@aws-sdk/client-s3')
    
    const semaphore = new Semaphore(this.options.maxConcurrency!)
    
    const promises = ids.map(async (id) => {
      await semaphore.acquire()
      try {
        result.statistics.apiCalls++
        
        const response = await this.s3Client.send(
          new GetObjectCommand({
            Bucket: this.bucketName,
            Key: `${prefix}${id}.json`
          })
        )

        if (response.Body) {
          const content = await response.Body.transformToString()
          const item = this.parseStoredObject(content)
          if (item) {
            result.items.set(id, item)
            result.statistics.totalRetrieved++
          }
        }
      } catch (error) {
        result.errors.set(id, error as Error)
        result.statistics.totalErrors++
      } finally {
        semaphore.release()
      }
    })

    await Promise.all(promises)
  }

  /**
   * Chunked parallel retrieval with intelligent batching
   */
  private async chunkedParallelGet<T>(
    ids: string[],
    prefix: string,
    result: BatchResult<T>
  ): Promise<void> {
    const chunkSize = Math.min(50, Math.ceil(ids.length / 10))
    const chunks = this.chunkArray(ids, chunkSize)
    
    // Process chunks with controlled concurrency
    const semaphore = new Semaphore(Math.min(5, chunks.length))
    
    const chunkPromises = chunks.map(async (chunk) => {
      await semaphore.acquire()
      try {
        await this.parallelGetObjects(chunk, prefix, result)
      } finally {
        semaphore.release()
      }
    })

    await Promise.all(chunkPromises)
  }

  /**
   * List-based batch retrieval for large datasets
   * Uses S3 ListObjects to reduce API calls
   */
  private async listBasedBatchGet<T>(
    ids: string[],
    prefix: string,
    result: BatchResult<T>
  ): Promise<void> {
    const { ListObjectsV2Command, GetObjectCommand } = await import('@aws-sdk/client-s3')
    
    // Create a set for O(1) lookup
    const idSet = new Set(ids)
    
    // List objects with the prefix
    let continuationToken: string | undefined
    const maxKeys = 1000
    
    do {
      result.statistics.apiCalls++
      
      const listResponse = await this.s3Client.send(
        new ListObjectsV2Command({
          Bucket: this.bucketName,
          Prefix: prefix,
          MaxKeys: maxKeys,
          ContinuationToken: continuationToken
        })
      )

      if (listResponse.Contents) {
        // Filter objects that match our requested IDs
        const matchingObjects = listResponse.Contents.filter((obj: any) => {
          if (!obj.Key) return false
          const id = obj.Key.replace(prefix, '').replace('.json', '')
          return idSet.has(id)
        })

        // Batch retrieve matching objects
        const semaphore = new Semaphore(this.options.maxConcurrency!)
        
        const retrievalPromises = matchingObjects.map(async (obj: any) => {
          if (!obj.Key) return

          await semaphore.acquire()
          try {
            result.statistics.apiCalls++
            
            const response = await this.s3Client.send(
              new GetObjectCommand({
                Bucket: this.bucketName,
                Key: obj.Key
              })
            )

            if (response.Body) {
              const content = await response.Body.transformToString()
              const item = this.parseStoredObject(content)
              if (item) {
                const id = obj.Key.replace(prefix, '').replace('.json', '')
                result.items.set(id, item)
                result.statistics.totalRetrieved++
              }
            }
          } catch (error) {
            const id = obj.Key.replace(prefix, '').replace('.json', '')
            result.errors.set(id, error as Error)
            result.statistics.totalErrors++
          } finally {
            semaphore.release()
          }
        })

        await Promise.all(retrievalPromises)
      }

      continuationToken = listResponse.NextContinuationToken
    } while (continuationToken && result.items.size < ids.length)
  }

  /**
   * Intelligent prefetch based on HNSW graph connectivity
   */
  public async prefetchConnectedNodes(
    currentNodeIds: string[],
    connectionMap: Map<string, Set<string>>,
    prefix: string = 'nodes/'
  ): Promise<BatchResult<HNSWNoun>> {
    // Analyze connection patterns to predict next nodes
    const predictedNodes = new Set<string>()
    
    for (const nodeId of currentNodeIds) {
      const connections = connectionMap.get(nodeId)
      if (connections) {
        // Add immediate neighbors
        connections.forEach(connId => predictedNodes.add(connId))
        
        // Add second-degree neighbors (limited)
        let count = 0
        for (const connId of connections) {
          if (count >= 5) break // Limit prefetch scope
          const secondDegree = connectionMap.get(connId)
          if (secondDegree) {
            secondDegree.forEach(id => {
              if (count < 20) {
                predictedNodes.add(id)
                count++
              }
            })
          }
        }
      }
    }

    // Remove nodes we already have
    const nodesToPrefetch = Array.from(predictedNodes).filter(
      id => !currentNodeIds.includes(id)
    )

    return this.batchGetNodes(nodesToPrefetch.slice(0, this.options.prefetchSize!), prefix)
  }

  /**
   * S3 Select-based retrieval for filtered queries
   */
  public async selectiveRetrieve(
    prefix: string,
    filter: {
      vectorDimension?: number
      metadataKey?: string
      metadataValue?: any
    }
  ): Promise<BatchResult<HNSWNoun>> {
    // This would use S3 Select to filter objects server-side
    // Reducing data transfer for large-scale operations
    
    const startTime = Date.now()
    const result: BatchResult<HNSWNoun> = {
      items: new Map(),
      errors: new Map(),
      statistics: {
        totalRequested: 0,
        totalRetrieved: 0,
        totalErrors: 0,
        duration: 0,
        apiCalls: 0
      }
    }

    // S3 Select implementation would go here
    // For now, fall back to list-based approach
    console.warn('S3 Select not implemented, falling back to list-based retrieval')
    
    result.statistics.duration = Date.now() - startTime
    return result
  }

  /**
   * Parse stored object from JSON string
   */
  private parseStoredObject(content: string): any {
    try {
      const parsed = JSON.parse(content)
      
      // Reconstruct HNSW node structure
      if (parsed.connections && typeof parsed.connections === 'object') {
        const connections = new Map<number, Set<string>>()
        for (const [level, nodeIds] of Object.entries(parsed.connections)) {
          connections.set(Number(level), new Set(nodeIds as string[]))
        }
        parsed.connections = connections
      }
      
      return parsed
    } catch (error) {
      console.error('Failed to parse stored object:', error)
      return null
    }
  }

  /**
   * Utility function to chunk arrays
   */
  private chunkArray<T>(array: T[], chunkSize: number): T[][] {
    const chunks: T[][] = []
    for (let i = 0; i < array.length; i += chunkSize) {
      chunks.push(array.slice(i, i + chunkSize))
    }
    return chunks
  }
}

/**
 * Simple semaphore implementation for concurrency control
 */
class Semaphore {
  private permits: number
  private waiting: Array<() => void> = []

  constructor(permits: number) {
    this.permits = permits
  }

  async acquire(): Promise<void> {
    if (this.permits > 0) {
      this.permits--
      return Promise.resolve()
    }

    return new Promise<void>((resolve) => {
      this.waiting.push(resolve)
    })
  }

  release(): void {
    if (this.waiting.length > 0) {
      const resolve = this.waiting.shift()!
      resolve()
    } else {
      this.permits++
    }
  }
}