brainy/src/storage/adapters/memoryStorage.ts

/**
 * Memory Storage Adapter
 * In-memory storage adapter for environments where persistent storage is not available or needed
 */

import * as fs from 'node:fs'
import * as nodePath from 'node:path'
import * as zlib from 'node:zlib'
import {
  GraphVerb,
  HNSWNoun,
  HNSWVerb,
  NounMetadata,
  VerbMetadata,
  HNSWNounWithMetadata,
  HNSWVerbWithMetadata,
  StatisticsData,
  NounType
} from '../../coreTypes.js'
import { BaseStorage, StorageBatchConfig, STATISTICS_KEY } from '../baseStorage.js'
import { PaginatedResult } from '../../types/paginationTypes.js'

// No type aliases needed - using the original types directly

/**
 * In-memory storage adapter
 * Uses Maps to store data in memory
 */
export class MemoryStorage extends BaseStorage {
  // Removed redundant Maps (nouns, verbs) - objectStore handles all storage via type-first paths
  private statistics: StatisticsData | null = null

  // Unified object store for primitive operations (replaces metadata, nounMetadata, verbMetadata)
  private objectStore: Map<string, any> = new Map()

  // Raw binary-blob store, keyed by blob key. Holds opaque byte payloads
  // (column-store segments, batch vectors) verbatim — no JSON envelope. In-memory
  // storage has no local file, so getBinaryBlobPath() returns null.
  private blobStore: Map<string, Buffer> = new Map()

  // Transaction log (8.0 MVCC) — the in-memory equivalent of the filesystem
  // adapter's `_system/tx-log.jsonl`. One JSON document per committed
  // transact(); serialized to a real JSONL file by snapshotToDirectory().
  private txLogLines: string[] = []

  // Backward compatibility aliases
  private get metadata(): Map<string, any> {
    return this.objectStore
  }
  private get nounMetadata(): Map<string, any> {
    return this.objectStore
  }
  private get verbMetadata(): Map<string, any> {
    return this.objectStore
  }

  constructor() {
    super()
  }

  /**
   * Get Memory-optimized batch configuration
   *
   * Memory storage has no rate limits and can handle very high throughput:
   * - Large batch sizes (1000 items)
   * - No delays needed (0ms)
   * - High concurrency (1000 operations)
   * - Parallel processing maximizes throughput
   *
   * @returns Memory-optimized batch configuration
   */
  public override getBatchConfig(): StorageBatchConfig {
    return {
      maxBatchSize: 1000,
      batchDelayMs: 0,
      maxConcurrent: 1000,
      supportsParallelWrites: true,  // Memory loves parallel operations
      rateLimit: {
        operationsPerSecond: 100000,  // Virtually unlimited
        burstCapacity: 100000
      }
    }
  }

  /**
   * Initialize the storage adapter
   * Calls super.init() to initialize GraphAdjacencyIndex and type statistics
   */
  public override async init(): Promise<void> {
    await super.init()
  }

  // Removed saveNoun_internal and getNoun_internal - using BaseStorage's type-first implementation

  /**
   * Get nouns with pagination and filtering
   * Returns HNSWNounWithMetadata[] (includes metadata field)
   * @param options Pagination and filtering options
   * @returns Promise that resolves to a paginated result of nouns with metadata
   */
  // Removed public method overrides (getNouns, getNounsWithPagination, getVerbs) - using BaseStorage's type-first implementation

  // Removed getNounsByNounType_internal and deleteNoun_internal - using BaseStorage's type-first implementation

  // Removed saveVerb_internal and getVerb_internal - using BaseStorage's type-first implementation

  // Removed verb *_internal method overrides - using BaseStorage's type-first implementation

  /**
   * Primitive operation: Write object to path
   * All metadata operations use this internally via base class routing
   */
  protected async writeObjectToPath(path: string, data: any): Promise<void> {
    // Store in unified object store using path as key
    this.objectStore.set(path, JSON.parse(JSON.stringify(data)))
  }

  /**
   * Primitive operation: Read object from path
   * All metadata operations use this internally via base class routing
   */
  protected async readObjectFromPath(path: string): Promise<any | null> {
    const data = this.objectStore.get(path)
    if (!data) {
      return null
    }
    return JSON.parse(JSON.stringify(data))
  }

  /**
   * Primitive operation: Delete object from path
   * All metadata operations use this internally via base class routing
   */
  protected async deleteObjectFromPath(path: string): Promise<void> {
    this.objectStore.delete(path)
  }

  /**
   * Primitive operation: List objects under path prefix
   * All metadata operations use this internally via base class routing
   */
  protected async listObjectsUnderPath(prefix: string): Promise<string[]> {
    const paths: string[] = []
    for (const key of this.objectStore.keys()) {
      if (key.startsWith(prefix)) {
        paths.push(key)
      }
    }
    return paths.sort()
  }

  // ===========================================================================
  // Raw binary-blob primitive
  // ===========================================================================

  /**
   * Persist a raw binary blob in memory under `key`. A defensive copy of the
   * bytes is stored so later mutations to the caller's buffer don't corrupt the
   * stored blob. Overwrites any existing blob at the same key.
   *
   * @param key - The blob key.
   * @param data - The exact bytes to store.
   */
  public async saveBinaryBlob(key: string, data: Buffer): Promise<void> {
    this.blobStore.set(key, Buffer.from(data))
  }

  /**
   * Load a copy of the bytes stored under `key`, or `null` if absent. A copy is
   * returned so callers cannot mutate the stored blob in place.
   *
   * @param key - The blob key.
   * @returns The blob bytes, or `null` if absent.
   */
  public async loadBinaryBlob(key: string): Promise<Buffer | null> {
    const data = this.blobStore.get(key)
    return data ? Buffer.from(data) : null
  }

  /**
   * Delete the blob stored under `key`. Missing blobs are ignored.
   *
   * @param key - The blob key.
   */
  public async deleteBinaryBlob(key: string): Promise<void> {
    // Registered-blob contract — parity with the filesystem adapter:
    // a declared family member is undeletable (throws ProtectedArtifactError).
    await this.assertBlobKeyDeletable(key)
    this.blobStore.delete(key)
  }

  /**
   * In-memory storage has no local filesystem path to mmap, so this always
   * returns `null`. Callers must use {@link loadBinaryBlob} instead.
   *
   * @param _key - The blob key (unused).
   * @returns Always `null`.
   */
  public getBinaryBlobPath(_key: string): string | null {
    return null
  }

  // ===========================================================================
  // Generational record layer (8.0 MVCC) — tx-log + snapshot primitives
  // ===========================================================================

  /**
   * Append one line to the in-memory transaction log (mirror of the
   * filesystem adapter's `_system/tx-log.jsonl` append).
   *
   * @param line - One complete JSON document, without trailing newline.
   */
  public async appendTxLogLine(line: string): Promise<void> {
    this.txLogLines.push(line)
  }

  /**
   * Read all transaction-log lines, oldest first (a copy — callers cannot
   * mutate the log).
   */
  public async readTxLogLines(): Promise<string[]> {
    return [...this.txLogLines]
  }

  // ===========================================================================
  // Binary raw-byte primitives — in-memory mirror of the filesystem adapter's
  // append-only substrate (the generation fact log's segments), so memory
  // brains dual-write facts too and the compat suite runs on both adapters.
  // ===========================================================================

  /** Raw binary files, keyed by verbatim path. */
  private rawBytesStore: Map<string, Uint8Array> = new Map()

  /** Append bytes to a raw binary file, creating it when absent. */
  public async appendRawBytes(rawPath: string, bytes: Uint8Array): Promise<void> {
    const existing = this.rawBytesStore.get(rawPath)
    if (!existing) {
      this.rawBytesStore.set(rawPath, bytes.slice())
      return
    }
    const merged = new Uint8Array(existing.length + bytes.length)
    merged.set(existing, 0)
    merged.set(bytes, existing.length)
    this.rawBytesStore.set(rawPath, merged)
  }

  /** Read a raw binary file whole (a copy); absent → null. */
  public async readRawBytes(rawPath: string): Promise<Uint8Array | null> {
    const bytes = this.rawBytesStore.get(rawPath)
    return bytes ? bytes.slice() : null
  }

  /** Replace a raw binary file (atomic by construction in memory). */
  public async writeRawBytes(rawPath: string, bytes: Uint8Array): Promise<void> {
    this.rawBytesStore.set(rawPath, bytes.slice())
  }

  /** Byte size of a raw binary file, or null when absent. */
  public async rawByteSize(rawPath: string): Promise<number | null> {
    const bytes = this.rawBytesStore.get(rawPath)
    return bytes ? bytes.length : null
  }

  /**
   * Serialize the entire in-memory store to a directory in the exact layout
   * the filesystem adapter uses (uncompressed JSON objects, `_blobs/*.bin`
   * payloads, `_system/tx-log.jsonl`). The resulting directory is a
   * self-contained store openable via `Brainy.load(path)` — persisting an
   * in-memory brain produces a real, durable snapshot.
   *
   * @param targetPath - Absolute directory for the snapshot. Created if
   *   missing; must be empty or absent (refuses to overwrite).
   */
  public async snapshotToDirectory(targetPath: string): Promise<void> {
    try {
      const existing = await fs.promises.readdir(targetPath)
      if (existing.length > 0) {
        throw new Error(
          `snapshotToDirectory: target ${targetPath} already exists and is not empty. ` +
          `Choose a fresh directory per snapshot.`
        )
      }
    } catch (error: any) {
      if (error.code !== 'ENOENT') throw error
    }
    await fs.promises.mkdir(targetPath, { recursive: true })

    for (const [key, value] of this.objectStore.entries()) {
      const filePath = nodePath.join(targetPath, ...key.split('/'))
      await fs.promises.mkdir(nodePath.dirname(filePath), { recursive: true })
      await fs.promises.writeFile(filePath, JSON.stringify(value, null, 2), 'utf-8')
    }

    for (const [key, data] of this.blobStore.entries()) {
      const blobPath = nodePath.join(targetPath, '_blobs', ...key.split('/')) + '.bin'
      await fs.promises.mkdir(nodePath.dirname(blobPath), { recursive: true })
      await fs.promises.writeFile(blobPath, data)
    }

    if (this.txLogLines.length > 0) {
      const logPath = nodePath.join(targetPath, '_system', 'tx-log.jsonl')
      await fs.promises.mkdir(nodePath.dirname(logPath), { recursive: true })
      await fs.promises.writeFile(logPath, this.txLogLines.map((l) => `${l}\n`).join(''), 'utf-8')
    }
  }

  /**
   * Replace the in-memory store's contents from a snapshot directory —
   * accepts snapshots produced by either adapter (handles the filesystem
   * adapter's gzip-compressed `.json.gz` objects transparently), then reloads
   * all derived state.
   *
   * @param sourcePath - Absolute path of a snapshot directory.
   */
  public async restoreFromDirectory(sourcePath: string): Promise<void> {
    const sourceStat = await fs.promises.stat(sourcePath).catch(() => null)
    if (!sourceStat || !sourceStat.isDirectory()) {
      throw new Error(`restoreFromDirectory: ${sourcePath} is not a directory`)
    }

    this.objectStore.clear()
    this.blobStore.clear()
    this.txLogLines = []
    this.statistics = null

    await this.loadDirectoryIntoStore(sourcePath, '')
    await this.reloadDerivedState()
  }

  /**
   * Recursively load a snapshot directory into the object/blob/tx-log stores.
   * Lock directories from filesystem-adapter snapshots are skipped (process-
   * local state, meaningless in memory).
   */
  private async loadDirectoryIntoStore(dirAbs: string, relPrefix: string): Promise<void> {
    const entries = await fs.promises.readdir(dirAbs, { withFileTypes: true })

    for (const entry of entries) {
      const rel = relPrefix ? `${relPrefix}/${entry.name}` : entry.name
      const abs = nodePath.join(dirAbs, entry.name)

      if (entry.isDirectory()) {
        if (relPrefix === '' && entry.name === 'locks') continue
        await this.loadDirectoryIntoStore(abs, rel)
        continue
      }
      if (!entry.isFile()) continue

      if (rel === '_system/tx-log.jsonl') {
        const content = await fs.promises.readFile(abs, 'utf-8')
        this.txLogLines = content.split('\n').filter((l) => l.length > 0)
      } else if (relPrefix.startsWith('_blobs') && rel.endsWith('.bin')) {
        const key = rel.slice('_blobs/'.length, -'.bin'.length)
        this.blobStore.set(key, await fs.promises.readFile(abs))
      } else if (rel.endsWith('.json.gz') || rel.endsWith('.gz')) {
        const raw = await fs.promises.readFile(abs)
        const key = rel.slice(0, -'.gz'.length)
        this.objectStore.set(key, JSON.parse(zlib.gunzipSync(raw).toString('utf-8')))
      } else if (entry.name.includes('.tmp.')) {
        // In-flight atomic-write remnant from a crashed filesystem store — skip.
      } else {
        const content = await fs.promises.readFile(abs, 'utf-8')
        this.objectStore.set(rel, JSON.parse(content))
      }
    }
  }

  /**
   * Get multiple metadata objects in batches (CRITICAL: Prevents socket exhaustion)
   * Memory storage implementation is simple since all data is already in memory
   */
  public async getMetadataBatch(ids: string[]): Promise<Map<string, any>> {
    const results = new Map<string, any>()

    // Memory storage can handle all IDs at once since it's in-memory
    for (const id of ids) {
      // CRITICAL: Use getNounMetadata() instead of deprecated getMetadata()
      // This ensures we fetch from the correct noun metadata store (2-file system)
      const metadata = await this.getNounMetadata(id)
      if (metadata) {
        results.set(id, metadata)
      }
    }

    return results
  }

  /**
   * Clear all data from storage
   * Clears objectStore (type-first paths)
   * Also clears writeCache to prevent stale data after clear
   */
  public async clear(): Promise<void> {
    this.objectStore.clear()
    this.blobStore.clear()
    this.rawBytesStore.clear()
    this.txLogLines = []
    this.statistics = null

    // Reset ALL shared derived state via the same path restore uses: the
    // write-through cache, id→type/subtype caches, per-type/subtype count
    // rollups, statistics cache, graph-index singleton, and the BlobStorage
    // instance (its LRU cache holds deleted blobs), then recount from the
    // now-empty object store. Without the rollup reset, stats() would keep
    // reporting per-type counts for deleted entities.
    await this.reloadDerivedState()
  }

  /**
   * Get information about storage usage and capacity
   * Uses BaseStorage counts
   */
  public async getStorageStatus(): Promise<{
    type: string
    used: number
    quota: number | null
    details?: Record<string, any>
  }> {
    return {
      type: 'memory',
      used: 0, // In-memory storage doesn't have a meaningful size
      quota: null, // In-memory storage doesn't have a quota
      details: {
        nodeCount: this.totalNounCount,
        edgeCount: this.totalVerbCount,
        objectStoreSize: this.objectStore.size
      }
    }
  }

  /**
   * Save statistics data to storage
   * @param statistics The statistics data to save
   */
  protected async saveStatisticsData(statistics: StatisticsData): Promise<void> {
    // For memory storage, we just need to store the statistics in memory
    // Create a deep copy to avoid reference issues
    this.statistics = {
      nounCount: {...statistics.nounCount},
      verbCount: {...statistics.verbCount},
      metadataCount: {...statistics.metadataCount},
      hnswIndexSize: statistics.hnswIndexSize,
      lastUpdated: statistics.lastUpdated,
      // Include serviceActivity if present
      ...(statistics.serviceActivity && {
        serviceActivity: Object.fromEntries(
          Object.entries(statistics.serviceActivity).map(([k, v]) => [k, {...v}])
        )
      }),
      // Include services if present
      ...(statistics.services && {
        services: statistics.services.map(s => ({...s}))
      })
    }

    // Since this is in-memory, there's no need for time-based partitioning
    // or legacy file handling
  }

  /**
   * Get statistics data from storage
   * @returns Promise that resolves to the statistics data or null if not found
   */
  protected async getStatisticsData(): Promise<StatisticsData | null> {
    if (!this.statistics) {
      // CRITICAL FIX: Statistics don't exist yet (first init)
      // Return minimal stats with counts instead of null
      // This prevents HNSW from seeing entityCount=0 during index rebuild
      return {
        nounCount: {},
        verbCount: {},
        metadataCount: {},
        hnswIndexSize: 0,
        totalNodes: this.totalNounCount,
        totalEdges: this.totalVerbCount,
        totalMetadata: 0,
        lastUpdated: new Date().toISOString()
      }
    }

    // Return a deep copy to avoid reference issues
    return {
      nounCount: {...this.statistics.nounCount},
      verbCount: {...this.statistics.verbCount},
      metadataCount: {...this.statistics.metadataCount},
      hnswIndexSize: this.statistics.hnswIndexSize,
      // CRITICAL FIX: Populate totalNodes and totalEdges from in-memory counts
      // HNSW rebuild depends on these fields to determine entity count
      totalNodes: this.totalNounCount,
      totalEdges: this.totalVerbCount,
      lastUpdated: this.statistics.lastUpdated,
      // Include serviceActivity if present
      ...(this.statistics.serviceActivity && {
        serviceActivity: Object.fromEntries(
          Object.entries(this.statistics.serviceActivity).map(([k, v]) => [k, {...v}])
        )
      }),
      // Include services if present
      ...(this.statistics.services && {
        services: this.statistics.services.map(s => ({...s}))
      })
    }

    // Since this is in-memory, there's no need for fallback mechanisms
    // to check multiple storage locations
  }

  /**
   * Initialize counts from in-memory storage - O(1) operation
   */
  protected async initializeCounts(): Promise<void> {
    // Scan objectStore paths (ID-first structure) to count entities
    this.entityCounts.clear()
    this.verbCounts.clear()

    let totalNouns = 0
    let totalVerbs = 0

    // Scan all paths in objectStore
    for (const path of this.objectStore.keys()) {
      // Count nouns (entities/nouns/{shard}/{id}/vectors.json)
      const nounMatch = path.match(/^entities\/nouns\/[0-9a-f]{2}\/[^/]+\/vectors\.json$/)
      if (nounMatch) {
        // Type is in metadata, not path - just count total
        totalNouns++
      }

      // Count verbs (entities/verbs/{shard}/{id}/vectors.json)
      const verbMatch = path.match(/^entities\/verbs\/[0-9a-f]{2}\/[^/]+\/vectors\.json$/)
      if (verbMatch) {
        // Type is in metadata, not path - just count total
        totalVerbs++
      }
    }

    this.totalNounCount = totalNouns
    this.totalVerbCount = totalVerbs
  }

  /**
   * Persist counts to storage - no-op for memory storage
   */
  protected async persistCounts(): Promise<void> {
    // No persistence needed for in-memory storage
    // Counts are always accurate from the live data structures
  }

  // =============================================
  // HNSW Index Persistence
  // =============================================

  /**
   * Get vector for a noun
   * Uses BaseStorage's type-first implementation
   */
  public async getNounVector(id: string): Promise<number[] | null> {
    const noun = await this.getNoun(id)
    return noun ? [...noun.vector] : null
  }

  // CRITICAL FIX: Mutex locks for HNSW concurrency control
  // Even in-memory operations need serialization to prevent async race conditions
  private hnswLocks = new Map<string, Promise<void>>()

  /**
   * Save HNSW graph data for a noun
   *
   * CRITICAL FIX: Mutex locking to prevent race conditions during concurrent HNSW updates
   * Even in-memory operations can race due to async/await interleaving
   * Prevents data corruption when multiple entities connect to same neighbor simultaneously
   */
  public async saveVectorIndexData(nounId: string, hnswData: {
    level: number
    connections: Record<string, string[]>
  }): Promise<void> {
    const path = `hnsw/${nounId}.json`

    // MUTEX LOCK: Wait for any pending operations on this entity
    while (this.hnswLocks.has(path)) {
      await this.hnswLocks.get(path)
    }

    // Acquire lock by creating a promise that we'll resolve when done
    let releaseLock!: () => void
    const lockPromise = new Promise<void>(resolve => { releaseLock = resolve })
    this.hnswLocks.set(path, lockPromise)

    try {
      // Read existing data (if exists)
      let existingNode: any = {}
      const existing = this.objectStore.get(path)
      if (existing) {
        existingNode = existing
      }

      // Preserve id and vector, update only HNSW graph metadata
      const updatedNode = {
        ...existingNode,  // Preserve all existing fields
        level: hnswData.level,
        connections: hnswData.connections
      }

      // Write atomically (in-memory, but now serialized by mutex)
      this.objectStore.set(path, JSON.parse(JSON.stringify(updatedNode)))
    } finally {
      // Release lock
      this.hnswLocks.delete(path)
      releaseLock()
    }
  }

  /**
   * Get HNSW graph data for a noun
   */
  public async getVectorIndexData(nounId: string): Promise<{
    level: number
    connections: Record<string, string[]>
  } | null> {
    const path = `hnsw/${nounId}.json`
    const data = await this.readObjectFromPath(path)
    return data || null
  }

  /**
   * Save HNSW system data (entry point, max level)
   *
   * CRITICAL FIX: Mutex locking to prevent race conditions
   */
  public async saveHNSWSystem(systemData: {
    entryPointId: string | null
    maxLevel: number
  }): Promise<void> {
    const path = 'system/hnsw-system.json'

    // MUTEX LOCK: Wait for any pending operations
    while (this.hnswLocks.has(path)) {
      await this.hnswLocks.get(path)
    }

    // Acquire lock
    let releaseLock!: () => void
    const lockPromise = new Promise<void>(resolve => { releaseLock = resolve })
    this.hnswLocks.set(path, lockPromise)

    try {
      // Write atomically (serialized by mutex)
      this.objectStore.set(path, JSON.parse(JSON.stringify(systemData)))
    } finally {
      // Release lock
      this.hnswLocks.delete(path)
      releaseLock()
    }
  }

  /**
   * Get HNSW system data
   */
  public async getHNSWSystem(): Promise<{
    entryPointId: string | null
    maxLevel: number
  } | null> {
    const path = 'system/hnsw-system.json'
    const data = await this.readObjectFromPath(path)
    return data || null
  }
}