open-brainy/src/migration/MigrationRunner.ts

/**
 * MigrationRunner: Executes schema migrations on Brainy storage
 *
 * Handles paginated iteration, resume-safe batching, and dry-run previews.
 * Uses BaseStorage methods directly — no adapter-level changes needed.
 */

import type { BaseStorage } from '../storage/baseStorage.js'
import type { NounMetadata, VerbMetadata } from '../coreTypes.js'
import type { Migration, MigrationState, MigrationPreview, MigrationResult, MigrateOptions, MigrationError } from './types.js'
import { MIGRATIONS } from './migrations.js'
import {
  splitNounMetadataRecord,
  splitVerbMetadataRecord,
  buildNounMetadataRecord,
  buildVerbMetadataRecord,
  RESERVED_ENTITY_FIELDS,
  RESERVED_RELATION_FIELDS
} from '../types/reservedFields.js'

const RESERVED_NOUN_SET: ReadonlySet<string> = new Set(RESERVED_ENTITY_FIELDS)
const RESERVED_VERB_SET: ReadonlySet<string> = new Set(RESERVED_RELATION_FIELDS)

/**
 * Normalize a stored record (either era: legacy flat OR v2 nested-bag) into
 * THE transform view — the one shape every migration transform receives:
 * engine fields top-level, the user's metadata bag nested under `metadata`.
 * Transforms never see the storage era; a migration written today works on
 * a brain of any age.
 */
function toTransformView(
  record: Record<string, unknown>,
  kind: 'noun' | 'verb'
): Record<string, unknown> {
  const { reserved, custom } =
    kind === 'noun' ? splitNounMetadataRecord(record) : splitVerbMetadataRecord(record)
  return { ...reserved, metadata: { ...custom } }
}

/**
 * Convert a transform's returned view back into a stamped v2 stored record.
 * LOUD CONTRACT: user fields belong inside `.metadata` — a stray top-level
 * key that is not an engine field is a migration bug under the
 * field-addressing law (pre-law transforms wrote user fields flat), and it
 * refuses with the fix in the message rather than silently dropping or
 * silently storing it as an engine key.
 */
function fromTransformView(
  view: Record<string, unknown>,
  kind: 'noun' | 'verb'
): Record<string, unknown> {
  const reservedSet = kind === 'noun' ? RESERVED_NOUN_SET : RESERVED_VERB_SET
  const engine: Record<string, unknown> = {}
  for (const [key, value] of Object.entries(view)) {
    if (key === 'metadata') continue
    if (!reservedSet.has(key)) {
      throw new Error(
        `migration transform returned a top-level key '${key}' that is not an ` +
          `engine field — under the field-addressing law user fields live inside ` +
          `.metadata (return { ...view, metadata: { ...view.metadata, ${key}: … } }).`
      )
    }
    engine[key] = value
  }
  const bag =
    view.metadata && typeof view.metadata === 'object' && !Array.isArray(view.metadata)
      ? (view.metadata as Record<string, unknown>)
      : {}
  return kind === 'noun'
    ? buildNounMetadataRecord(engine, bag)
    : buildVerbMetadataRecord(engine, bag)
}

const MIGRATION_STATE_KEY = '__migration_state__'
const PREVIEW_SAMPLE_SIZE = 5
const DEFAULT_MAX_ERRORS = 100

export class MigrationRunner {
  private storage: BaseStorage
  private stateCache: MigrationState | null | undefined = undefined

  constructor(storage: BaseStorage) {
    this.storage = storage
    MigrationRunner.validateMigrations(MIGRATIONS)
  }

  /**
   * Validate migration definitions.
   * Called automatically in constructor for the global MIGRATIONS array.
   * Also available as a static method for validating custom migration arrays.
   */
  static validateMigrations(migrations: Migration[]): void {
    if (migrations.length === 0) return

    const seenIds = new Set<string>()
    const validApplies = new Set(['nouns', 'verbs', 'both'])

    for (const m of migrations) {
      if (!m.id || typeof m.id !== 'string') {
        throw new Error(`Migration has missing or invalid id`)
      }
      if (seenIds.has(m.id)) {
        throw new Error(`Duplicate migration id: "${m.id}"`)
      }
      seenIds.add(m.id)

      if (!m.version || typeof m.version !== 'string') {
        throw new Error(`Migration "${m.id}" has missing or invalid version`)
      }
      if (!m.description || typeof m.description !== 'string') {
        throw new Error(`Migration "${m.id}" has missing or invalid description`)
      }
      if (!validApplies.has(m.applies)) {
        throw new Error(`Migration "${m.id}" has invalid applies value: "${m.applies}" (must be "nouns", "verbs", or "both")`)
      }
      if (typeof m.transform !== 'function') {
        throw new Error(`Migration "${m.id}" has non-function transform`)
      }
    }
  }

  /**
   * Check if there are pending migrations to run.
   * Single getMetadata() call — ~0ms overhead when no migrations exist.
   */
  async hasPendingMigrations(): Promise<boolean> {
    if (MIGRATIONS.length === 0) return false
    const state = await this.getState()
    return this.getPendingMigrations(state).length > 0
  }

  /**
   * Get the version string for the next pending migration.
   */
  nextMigrationVersion(): string {
    const pending = this.getPendingMigrationsFromCache()
    return pending.length > 0 ? pending[pending.length - 1].version : 'unknown'
  }

  /**
   * Get count of pending migrations (for log messages).
   */
  async pendingCount(): Promise<number> {
    if (MIGRATIONS.length === 0) return 0
    const state = await this.getState()
    return this.getPendingMigrations(state).length
  }

  /**
   * Preview what a migration would do without writing anything.
   * Scans entities, applies transforms in memory, reports counts + samples.
   */
  async preview(): Promise<MigrationPreview> {
    const state = await this.getState()
    const pending = this.getPendingMigrations(state)

    if (pending.length === 0) {
      return {
        pendingMigrations: [],
        affectedEntities: 0,
        totalEntities: 0,
        sampleChanges: [],
        estimatedTime: '0ms'
      }
    }

    let totalEntities = 0
    let affectedEntities = 0
    const sampleChanges: MigrationPreview['sampleChanges'] = []
    const batchConfig = this.storage.getBatchConfig()
    const batchSize = batchConfig.maxBatchSize

    // Scan nouns if any pending migration applies to nouns
    const nounMigrations = pending.filter(m => m.applies === 'nouns' || m.applies === 'both')
    if (nounMigrations.length > 0) {
      let offset = 0
      let hasMore = true

      while (hasMore) {
        const batch = await this.storage.getNouns({ pagination: { offset, limit: batchSize } })
        const ids = batch.items.map(e => e.id)
        const metadataBatch = await this.storage.getNounMetadataBatch(ids)

        for (const entity of batch.items) {
          totalEntities++
          const entityMeta = metadataBatch.get(entity.id)
          if (!entityMeta) continue

          // Transforms see THE view (engine fields + nested user bag),
          // never the raw storage era.
          const view = toTransformView(entityMeta as Record<string, unknown>, 'noun')
          const result = this.applyTransforms(view, nounMigrations)
          if (result !== null) {
            affectedEntities++
            if (sampleChanges.length < PREVIEW_SAMPLE_SIZE) {
              sampleChanges.push({
                id: entity.id,
                before: view,
                after: result
              })
            }
          }
        }
        hasMore = batch.hasMore
        offset += batch.items.length
      }
    }

    // Scan verbs if any pending migration applies to verbs
    const verbMigrations = pending.filter(m => m.applies === 'verbs' || m.applies === 'both')
    if (verbMigrations.length > 0) {
      let offset = 0
      let hasMore = true

      while (hasMore) {
        const batch = await this.storage.getVerbs({ pagination: { offset, limit: batchSize } })

        for (const verb of batch.items) {
          totalEntities++
          const verbMeta = await this.storage.getVerbMetadata(verb.id)
          if (!verbMeta) continue

          const view = toTransformView(verbMeta as Record<string, unknown>, 'verb')
          const result = this.applyTransforms(view, verbMigrations)
          if (result !== null) {
            affectedEntities++
            if (sampleChanges.length < PREVIEW_SAMPLE_SIZE) {
              sampleChanges.push({
                id: verb.id,
                before: view,
                after: result
              })
            }
          }
        }
        hasMore = batch.hasMore
        offset += batch.items.length
      }
    }

    return {
      pendingMigrations: pending.map(m => ({ id: m.id, description: m.description })),
      affectedEntities,
      totalEntities,
      sampleChanges,
      estimatedTime: this.estimateTime(totalEntities)
    }
  }

  /**
   * Run all pending migrations.
   * Iterates entities in paginated batches, transforms metadata, saves changes.
   * Resume-safe: saves offset after each batch so interrupted migrations can continue.
   *
   * Entity-level errors are tracked (not thrown). If maxErrors is exceeded, migration
   * stops early and returns partial results with errors.
   */
  async run(options?: Pick<MigrateOptions, 'onProgress' | 'maxErrors'>): Promise<Omit<MigrationResult, 'backupPath'>> {
    const state = await this.getState()
    const pending = this.getPendingMigrations(state)

    if (pending.length === 0) {
      return { migrationsApplied: [], entitiesProcessed: 0, entitiesModified: 0, errors: [] }
    }

    let totalProcessed = 0
    let totalModified = 0
    const appliedMigrations: string[] = []
    const errors: MigrationError[] = []
    const maxErrors = options?.maxErrors ?? DEFAULT_MAX_ERRORS
    const batchConfig = this.storage.getBatchConfig()
    const batchSize = batchConfig.maxBatchSize
    const batchDelay = batchConfig.batchDelayMs

    for (const migration of pending) {
      if (errors.length >= maxErrors) break

      const resumeOffset = state?.resumeState?.migrationId === migration.id
        ? state.resumeState.lastProcessedOffset
        : 0

      let processed = 0
      let modified = 0

      // Process nouns
      if (migration.applies === 'nouns' || migration.applies === 'both') {
        const result = await this.migrateNouns(migration, resumeOffset, batchSize, batchDelay, errors, maxErrors, options?.onProgress)
        processed += result.processed
        modified += result.modified
      }

      // Process verbs
      if (migration.applies === 'verbs' || migration.applies === 'both') {
        if (errors.length < maxErrors) {
          const result = await this.migrateVerbs(migration, 0, batchSize, batchDelay, errors, maxErrors, options?.onProgress)
          processed += result.processed
          modified += result.modified
        }
      }

      totalProcessed += processed
      totalModified += modified
      appliedMigrations.push(migration.id)

      // Save completed migration state
      await this.saveState({
        completedVersion: migration.version,
        completedAt: Date.now(),
        completedMigrations: [...(state?.completedMigrations || []), migration.id],
        resumeState: undefined
      })
    }

    // Clear state cache so next check reads fresh
    this.stateCache = undefined

    return {
      migrationsApplied: appliedMigrations,
      entitiesProcessed: totalProcessed,
      entitiesModified: totalModified,
      errors
    }
  }

  // ─── Private helpers ───────────────────────────────────────────────

  private async migrateNouns(
    migration: Migration,
    startOffset: number,
    batchSize: number,
    batchDelay: number,
    errors: MigrationError[],
    maxErrors: number,
    onProgress?: MigrateOptions['onProgress']
  ): Promise<{ processed: number; modified: number }> {
    let offset = startOffset
    let hasMore = true
    let processed = 0
    let modified = 0

    while (hasMore) {
      const batch = await this.storage.getNouns({ pagination: { offset, limit: batchSize } })

      for (const entity of batch.items) {
        if (errors.length >= maxErrors) {
          return { processed, modified }
        }

        processed++
        const metadata = await this.storage.getNounMetadataBatch([entity.id])
        const entityMeta = metadata.get(entity.id)
        if (!entityMeta) continue

        try {
          const transformed = migration.transform(
            toTransformView(entityMeta as Record<string, unknown>, 'noun')
          )
          if (transformed !== null) {
            // Re-stamp as a v2 record (also upgrades legacy records touched
            // by a migration onto the nested-bag shape).
            await this.storage.saveNounMetadata(
              entity.id,
              fromTransformView(transformed, 'noun') as NounMetadata
            )
            modified++
          }
        } catch (err) {
          errors.push({
            entityId: entity.id,
            migrationId: migration.id,
            error: err instanceof Error ? err.message : String(err)
          })
        }
      }

      hasMore = batch.hasMore
      offset += batch.items.length

      // Save resume state after each batch
      if (hasMore) {
        await this.saveResumeState(migration.id, offset)
      }

      // Report progress
      if (onProgress) {
        onProgress({
          migrationId: migration.id,
          processed,
          modified,
          hasMore
        })
      }

      // Respect adapter rate limiting
      if (batchDelay > 0 && hasMore) {
        await new Promise(resolve => setTimeout(resolve, batchDelay))
      }
    }

    return { processed, modified }
  }

  private async migrateVerbs(
    migration: Migration,
    startOffset: number,
    batchSize: number,
    batchDelay: number,
    errors: MigrationError[],
    maxErrors: number,
    onProgress?: MigrateOptions['onProgress']
  ): Promise<{ processed: number; modified: number }> {
    let offset = startOffset
    let hasMore = true
    let processed = 0
    let modified = 0

    while (hasMore) {
      const batch = await this.storage.getVerbs({ pagination: { offset, limit: batchSize } })

      for (const verb of batch.items) {
        if (errors.length >= maxErrors) {
          return { processed, modified }
        }

        processed++
        const metadata = await this.storage.getVerbMetadata(verb.id)
        if (!metadata) continue

        try {
          const transformed = migration.transform(
            toTransformView(metadata as Record<string, unknown>, 'verb')
          )
          if (transformed !== null) {
            await this.storage.saveVerbMetadata(
              verb.id,
              fromTransformView(transformed, 'verb') as VerbMetadata
            )
            modified++
          }
        } catch (err) {
          errors.push({
            entityId: verb.id,
            migrationId: migration.id,
            error: err instanceof Error ? err.message : String(err)
          })
        }
      }

      hasMore = batch.hasMore
      offset += batch.items.length

      // Save resume state after each batch
      if (hasMore) {
        await this.saveResumeState(migration.id, offset)
      }

      // Report progress
      if (onProgress) {
        onProgress({
          migrationId: migration.id,
          processed,
          modified,
          hasMore
        })
      }

      // Respect adapter rate limiting
      if (batchDelay > 0 && hasMore) {
        await new Promise(resolve => setTimeout(resolve, batchDelay))
      }
    }

    return { processed, modified }
  }

  private applyTransforms(metadata: Record<string, unknown>, migrations: Migration[]): Record<string, unknown> | null {
    let current = metadata
    let anyChanged = false

    for (const migration of migrations) {
      const result = migration.transform(current)
      if (result !== null) {
        current = result
        anyChanged = true
      }
    }

    return anyChanged ? current : null
  }

  private getPendingMigrations(state: MigrationState | null): Migration[] {
    const completed = new Set(state?.completedMigrations || [])
    return MIGRATIONS.filter(m => !completed.has(m.id))
  }

  private getPendingMigrationsFromCache(): Migration[] {
    const state = this.stateCache === undefined ? null : this.stateCache
    return this.getPendingMigrations(state)
  }

  private async getState(): Promise<MigrationState | null> {
    if (this.stateCache !== undefined) return this.stateCache
    const state = await this.storage.getMetadata(MIGRATION_STATE_KEY) as unknown as MigrationState | null
    this.stateCache = state
    return state
  }

  private async saveState(state: MigrationState): Promise<void> {
    await this.storage.saveMetadata(MIGRATION_STATE_KEY, state as unknown as NounMetadata)
    this.stateCache = state
  }

  private async saveResumeState(migrationId: string, offset: number): Promise<void> {
    const state = await this.getState()
    await this.saveState({
      completedVersion: state?.completedVersion || '',
      completedAt: state?.completedAt || 0,
      completedMigrations: state?.completedMigrations || [],
      resumeState: { migrationId, lastProcessedOffset: offset }
    })
  }

  private estimateTime(entityCount: number): string {
    if (entityCount === 0) return '0ms'
    if (entityCount < 1000) return '<1s'
    if (entityCount < 10000) return '~1-5s'
    if (entityCount < 100000) return '~10s-1min'
    if (entityCount < 1000000) return '~1-5min'
    return '~5min+'
  }
}