brainy/src/distributed/configManager.ts

/**
 * Distributed Configuration Manager
 * Manages shared configuration in S3 for distributed Brainy instances
 */

import { v4 as uuidv4 } from 'uuid'
import { 
  DistributedConfig, 
  SharedConfig, 
  InstanceInfo,
  InstanceRole 
} from '../types/distributedTypes.js'
import { StorageAdapter } from '../coreTypes.js'

export class DistributedConfigManager {
  private config: SharedConfig | null = null
  private instanceId: string
  private role: InstanceRole | undefined
  private configPath: string
  private heartbeatInterval: number
  private configCheckInterval: number
  private instanceTimeout: number
  private storage: StorageAdapter
  private heartbeatTimer?: NodeJS.Timeout
  private configWatchTimer?: NodeJS.Timeout
  private lastConfigVersion: number = 0
  private onConfigUpdate?: (config: SharedConfig) => void
  
  constructor(
    storage: StorageAdapter,
    distributedConfig?: DistributedConfig,
    brainyMode?: { readOnly?: boolean; writeOnly?: boolean }
  ) {
    this.storage = storage
    this.instanceId = distributedConfig?.instanceId || `instance-${uuidv4()}`
    this.configPath = distributedConfig?.configPath || '_brainy/config.json'
    this.heartbeatInterval = distributedConfig?.heartbeatInterval || 30000
    this.configCheckInterval = distributedConfig?.configCheckInterval || 10000
    this.instanceTimeout = distributedConfig?.instanceTimeout || 60000
    
    // Set role from distributed config if provided
    if (distributedConfig?.role) {
      this.role = distributedConfig.role
    }
    // Infer role from Brainy's read/write mode if not explicitly set
    else if (brainyMode) {
      if (brainyMode.writeOnly) {
        this.role = 'writer'
      } else if (brainyMode.readOnly) {
        this.role = 'reader'
      }
      // If neither readOnly nor writeOnly, role must be explicitly set
    }
  }
  
  /**
   * Initialize the distributed configuration
   */
  async initialize(): Promise<SharedConfig> {
    // Load or create configuration
    this.config = await this.loadOrCreateConfig()
    
    // Determine role if not explicitly set
    if (!this.role) {
      this.role = await this.determineRole()
    }
    
    // Register this instance
    await this.registerInstance()
    
    // Start heartbeat and config watching
    this.startHeartbeat()
    this.startConfigWatch()
    
    return this.config
  }
  
  /**
   * Load existing config or create new one
   */
  private async loadOrCreateConfig(): Promise<SharedConfig> {
    try {
      // Use metadata storage with a special ID for config
      const configData = await this.storage.getMetadata('_distributed_config')
      if (configData) {
        this.lastConfigVersion = configData.version
        return configData as SharedConfig
      }
    } catch (error) {
      // Config doesn't exist yet
    }
    
    // Create default config
    const newConfig: SharedConfig = {
      version: 1,
      updated: new Date().toISOString(),
      settings: {
        partitionStrategy: 'hash',
        partitionCount: 100,
        embeddingModel: 'text-embedding-ada-002',
        dimensions: 1536,
        distanceMetric: 'cosine',
        hnswParams: {
          M: 16,
          efConstruction: 200
        }
      },
      instances: {}
    }
    
    await this.saveConfig(newConfig)
    return newConfig
  }
  
  /**
   * Determine role based on configuration
   * IMPORTANT: Role must be explicitly set - no automatic assignment based on order
   */
  private async determineRole(): Promise<InstanceRole> {
    // Check environment variable first
    if (process.env.BRAINY_ROLE) {
      const role = process.env.BRAINY_ROLE.toLowerCase()
      if (role === 'writer' || role === 'reader' || role === 'hybrid') {
        return role as InstanceRole
      }
      throw new Error(`Invalid BRAINY_ROLE: ${process.env.BRAINY_ROLE}. Must be 'writer', 'reader', or 'hybrid'`)
    }
    
    // Check if explicitly passed in distributed config
    if (this.role) {
      return this.role
    }
    
    // DO NOT auto-assign roles based on deployment order or existing instances
    // This is dangerous and can lead to data corruption or loss
    throw new Error(
      'Distributed mode requires explicit role configuration. ' +
      'Set BRAINY_ROLE environment variable or pass role in distributed config. ' +
      'Valid roles: "writer", "reader", "hybrid"'
    )
  }
  
  /**
   * Check if an instance is still alive
   */
  private isInstanceAlive(instance: InstanceInfo): boolean {
    const lastSeen = new Date(instance.lastHeartbeat).getTime()
    const now = Date.now()
    return (now - lastSeen) < this.instanceTimeout
  }
  
  /**
   * Register this instance in the shared config
   */
  private async registerInstance(): Promise<void> {
    if (!this.config) return
    
    // Role must be set by this point
    if (!this.role) {
      throw new Error('Cannot register instance without a role')
    }
    
    const instanceInfo: InstanceInfo = {
      role: this.role,
      status: 'active',
      lastHeartbeat: new Date().toISOString(),
      metrics: {
        memoryUsage: process.memoryUsage().heapUsed
      }
    }
    
    // Add endpoint if available
    if (process.env.SERVICE_ENDPOINT) {
      instanceInfo.endpoint = process.env.SERVICE_ENDPOINT
    }
    
    this.config.instances[this.instanceId] = instanceInfo
    await this.saveConfig(this.config)
  }
  
  /**
   * Save configuration with version increment
   */
  private async saveConfig(config: SharedConfig): Promise<void> {
    config.version++
    config.updated = new Date().toISOString()
    this.lastConfigVersion = config.version
    
    // Use metadata storage with a special ID for config
    await this.storage.saveMetadata('_distributed_config', config)
    
    this.config = config
  }
  
  /**
   * Start heartbeat to keep instance alive in config
   */
  private startHeartbeat(): void {
    this.heartbeatTimer = setInterval(async () => {
      await this.updateHeartbeat()
    }, this.heartbeatInterval)
  }
  
  /**
   * Update heartbeat and clean stale instances
   */
  private async updateHeartbeat(): Promise<void> {
    if (!this.config) return
    
    // Reload config to get latest state
    try {
      const latestConfig = await this.loadConfig()
      if (latestConfig) {
        this.config = latestConfig
      }
    } catch (error) {
      console.error('Failed to reload config:', error)
    }
    
    // Update our heartbeat
    if (this.config.instances[this.instanceId]) {
      this.config.instances[this.instanceId].lastHeartbeat = new Date().toISOString()
      this.config.instances[this.instanceId].status = 'active'
      
      // Update metrics if available
      this.config.instances[this.instanceId].metrics = {
        memoryUsage: process.memoryUsage().heapUsed
      }
    } else {
      // Re-register if we were removed
      await this.registerInstance()
      return
    }
    
    // Clean up stale instances
    const now = Date.now()
    let hasChanges = false
    
    for (const [id, instance] of Object.entries(this.config.instances)) {
      if (id === this.instanceId) continue
      
      const lastSeen = new Date(instance.lastHeartbeat).getTime()
      if (now - lastSeen > this.instanceTimeout) {
        delete this.config.instances[id]
        hasChanges = true
      }
    }
    
    // Save if there were changes
    if (hasChanges) {
      await this.saveConfig(this.config)
    } else {
      // Just update our heartbeat without version increment
      await this.storage.saveMetadata('_distributed_config', this.config)
    }
  }
  
  /**
   * Start watching for config changes
   */
  private startConfigWatch(): void {
    this.configWatchTimer = setInterval(async () => {
      await this.checkForConfigUpdates()
    }, this.configCheckInterval)
  }
  
  /**
   * Check for configuration updates
   */
  private async checkForConfigUpdates(): Promise<void> {
    try {
      const latestConfig = await this.loadConfig()
      if (!latestConfig) return
      
      if (latestConfig.version > this.lastConfigVersion) {
        this.config = latestConfig
        this.lastConfigVersion = latestConfig.version
        
        // Notify listeners of config update
        if (this.onConfigUpdate) {
          this.onConfigUpdate(latestConfig)
        }
      }
    } catch (error) {
      console.error('Failed to check config updates:', error)
    }
  }
  
  /**
   * Load configuration from storage
   */
  private async loadConfig(): Promise<SharedConfig | null> {
    try {
      const configData = await this.storage.getMetadata('_distributed_config')
      if (configData) {
        return configData as SharedConfig
      }
    } catch (error) {
      console.error('Failed to load config:', error)
    }
    return null
  }
  
  /**
   * Get current configuration
   */
  getConfig(): SharedConfig | null {
    return this.config
  }
  
  /**
   * Get instance role
   */
  getRole(): InstanceRole {
    if (!this.role) {
      throw new Error('Role not initialized')
    }
    return this.role
  }
  
  /**
   * Get instance ID
   */
  getInstanceId(): string {
    return this.instanceId
  }
  
  /**
   * Set config update callback
   */
  setOnConfigUpdate(callback: (config: SharedConfig) => void): void {
    this.onConfigUpdate = callback
  }
  
  /**
   * Get all active instances of a specific role
   */
  getInstancesByRole(role: InstanceRole): InstanceInfo[] {
    if (!this.config) return []
    
    return Object.entries(this.config.instances)
      .filter(([_, instance]) => 
        instance.role === role && 
        this.isInstanceAlive(instance)
      )
      .map(([_, instance]) => instance)
  }
  
  /**
   * Update instance metrics
   */
  async updateMetrics(metrics: Partial<InstanceInfo['metrics']>): Promise<void> {
    if (!this.config || !this.config.instances[this.instanceId]) return
    
    this.config.instances[this.instanceId].metrics = {
      ...this.config.instances[this.instanceId].metrics,
      ...metrics
    }
    
    // Don't increment version for metric updates
    await this.storage.saveMetadata('_distributed_config', this.config)
  }
  
  /**
   * Cleanup resources
   */
  async cleanup(): Promise<void> {
    // Stop timers
    if (this.heartbeatTimer) {
      clearInterval(this.heartbeatTimer)
    }
    if (this.configWatchTimer) {
      clearInterval(this.configWatchTimer)
    }
    
    // Mark instance as inactive
    if (this.config && this.config.instances[this.instanceId]) {
      this.config.instances[this.instanceId].status = 'inactive'
      await this.saveConfig(this.config)
    }
  }
}