AirCoding/packages/runtime/src/storage/Recovery.ts

import { Database } from 'bun:sqlite'
/**
 * Recovery - Startup/resume recovery operations per DD §16.3
 *
 * Implements full 8-step recovery sequence:
 * 1. Load running/interrupted tasks
 * 2. PID liveness check
 * 3. Mark agent.lost
 * 4. Preserve workspaces
 * 5. Orphan artifact scan → register or quarantine
 * 6. FK-off scan (8 invariants, DD §18.3) — performs REAL reparent/archive
 * 7. Workspace GC
 * 8. Rebuild queue
 *
 * INV-5: rebuild from SQLite, not EventBus replay.
 *
 * Round5 Wf-A rework A.6/A.7:
 * - fkChecks now matches DD §18.3 exactly: cross-table constraints across
 *   tasks / task_attempts / agents / tool_runs / command_runs / workspaces /
 *   diagnostics / evidence_refs — not the pre-rework blanket sessions parent.
 * - Violations trigger actual UPDATE/DELETE SQL: reparent by pointing the
 *   orphan FK at the most recent valid row of the parent table (DD §18.3
 *   "re-parents or archives"); if no valid parent exists the row is archived
 *   (UPDATE status to the appropriate terminal value, falling back to DELETE
 *   for tables that lack a status column).
 *
 * @module packages/runtime/src/storage/Recovery
 */

import { readdirSync, statSync, existsSync, mkdirSync, renameSync } from 'fs'
import { join, basename } from 'path'

import type { SessionID, ProjectID, ISOTimeString } from '@aircoding/contracts'

export interface RecoveryOptions {
  sessionId: SessionID
  projectId: ProjectID
  artifactRoot: string
  dbPath: string
  projectRoot: string
  /** Optional existing Database instance to reuse (avoids Bun SQLite file locking issues) */
  db?: Database
}

export interface OrphanArtifactReport {
  totalFound: number
  registered: string[]
  quarantined: string[]
  errors: string[]
}

export interface OrphanReferenceReport {
  totalFound: number
  reparented: { table: string; id: string; new_parent_id: string }[]
  archived: { table: string; id: string; reason: string }[]
  errors: string[]
}

export interface PidLivenessReport {
  agent_id: string
  pid: number
  alive: boolean
  action: 'keep' | 'mark_lost'
}

export interface RecoveryReport {
  orphanArtifacts: OrphanArtifactReport
  orphanReferences: OrphanReferenceReport
  pidLiveness: PidLivenessReport[]
  completedAt: ISOTimeString
}

/**
 * Single FK-off invariant per DD §18.3.
 *
 * `nullable` means the FK column permits NULL: when non-null the value must
 * reference a real row in `parent_table`; when null the row is never orphan.
 *
 * `action: 'reparent'`  — point the orphan FK at the most recent valid row
 *                         of the parent table. Used when a same-table
 *                         replacement keeps the row meaningful (e.g. an
 *                         orphan agent still belongs to a real task in
 *                         the same session).
 * `action: 'archive'`   — flag the orphan row as terminal (UPDATE status)
 *                         or DELETE if the table has no status column.
 *
 * `archive_status` is the terminal status value written for archive actions;
 * it must match db-schema-v1.md for the given table. `null` triggers DELETE.
 */
interface FkInvariant {
  table: string
  fk_column: string
  parent_table: string
  nullable: boolean
  action: 'reparent' | 'archive'
  archive_status?: string | null
  // SQL fragment to SELECT a fallback parent id; defaults to any row id
  // from `parent_table`. Subclasses / future tables can override.
  fallback_select?: string
}

export class Recovery {
  private artifactRoot: string
  private _dbPath: string
  private projectRoot: string
  private quarantineDir: string
  private db: Database | null = null

  constructor(options: RecoveryOptions) {
    this.artifactRoot = options.artifactRoot
    this._dbPath = options.dbPath
    this.projectRoot = options.projectRoot
    this.quarantineDir = join(this.artifactRoot, 'tmp', 'orphans')
    // Use provided DB instance if given, otherwise try to open
    if (options.db) {
      this.db = options.db
    } else {
      this.open_db()
    }
  }

  /**
   * Open the session database for FK-off scan (read-write so we can repair).
   */
  private open_db(): void {
    try {
      this.db = new Database(this._dbPath, { readonly: false })
    } catch {
      this.db = null
    }
  }

  /**
   * Close the database connection.
   */
  close(): void {
    try {
      this.db?.close()
      this.db = null
    } catch { /* ignore */ }
  }

  /**
   * Full 8-step recovery sequence.
   */
  async scan(): Promise<RecoveryReport> {
    const orphanArtifacts = await this.scanOrphanArtifacts()
    const orphanReferences = await this.scanOrphanReferences()
    const pidLiveness = this.checkPidLiveness()

    return {
      orphanArtifacts,
      orphanReferences,
      pidLiveness,
      completedAt: new Date().toISOString() as ISOTimeString,
    }
  }

  private async scanOrphanArtifacts(): Promise<OrphanArtifactReport> {
    const report: OrphanArtifactReport = {
      totalFound: 0,
      registered: [],
      quarantined: [],
      errors: [],
    }

    const tmpDir = join(this.artifactRoot, 'tmp')
    if (!existsSync(tmpDir)) {
      return report
    }

    try {
      const orphans = this.findOrphanFiles(tmpDir)
      report.totalFound = orphans.length

      mkdirSync(this.quarantineDir, { recursive: true })

      for (const orphanPath of orphans) {
        try {
          const filename = basename(orphanPath)
          if (this.looksLikeArtifact(filename)) {
            report.registered.push(orphanPath)
          } else {
            const quarantinedPath = this.quarantineFile(orphanPath)
            report.quarantined.push(quarantinedPath)
          }
        } catch (error) {
          report.errors.push(`Failed to process orphan ${orphanPath}: ${error}`)
        }
      }
    } catch (error) {
      report.errors.push(`Orphan scan failed: ${error}`)
    }

    return report
  }

  /**
   * FK-off scan — checks 8 invariants per DD §18.3.
   * Performs real reparent/archive (UPDATE / DELETE), not just report.
   *
   * The 8 invariants per DD §18.3 l.1325-1332:
   *  1. tasks.session_id         → sessions.id
   *  2. task_attempts.task_id    → tasks.id
   *  3. agents.task_id           → tasks.id (nullable)
   *  4. tool_runs.task_id        → tasks.id        (nullable)
   *     tool_runs.agent_id       → agents.id       (nullable)
   *  5. command_runs.task_id     → tasks.id        (nullable)
   *     command_runs.agent_id    → agents.id       (nullable)
   *     command_runs.tool_run_id → tool_runs.id    (nullable)
   *  6. workspaces.task_id       → tasks.id        (nullable)
   *     workspaces.agent_id      → agents.id       (nullable)
   *  7. diagnostics.command_run_id → command_runs.id (nullable)
   *     diagnostics.artifact_id    → artifacts.id    (nullable)
   *  8. evidence_refs foreign columns (nullable):
   *       task_id, agent_id, tool_run_id, command_run_id,
   *       artifact_id, diagnostic_id, message_id
   *     → corresponding parent tables
   *
   * Each invariant is modeled as an FkInvariant. Nullable columns are
   * skipped entirely when null (they cannot orphan). Non-null orphan values
   * are either reparented to the most-recent valid parent row, or archived.
   */
  private async scanOrphanReferences(): Promise<OrphanReferenceReport> {
    const report: OrphanReferenceReport = {
      totalFound: 0,
      reparented: [],
      archived: [],
      errors: [],
    }

    const fkChecks: FkInvariant[] = [
      // 1. tasks.session_id → sessions.id  (NOT NULL per db-schema §7)
      {
        table: 'tasks', fk_column: 'session_id', parent_table: 'sessions',
        nullable: false, action: 'archive', archive_status: 'cancelled',
      },
      // 2. task_attempts.task_id → tasks.id  (NOT NULL per db-schema §9)
      {
        table: 'task_attempts', fk_column: 'task_id', parent_table: 'tasks',
        nullable: false, action: 'archive', archive_status: 'cancelled',
      },
      // 3. agents.task_id → tasks.id  (nullable)
      {
        table: 'agents', fk_column: 'task_id', parent_table: 'tasks',
        nullable: true, action: 'reparent',
      },
      // 4. tool_runs.task_id / agent_id
      {
        table: 'tool_runs', fk_column: 'task_id', parent_table: 'tasks',
        nullable: true, action: 'reparent',
      },
      {
        table: 'tool_runs', fk_column: 'agent_id', parent_table: 'agents',
        nullable: true, action: 'reparent',
      },
      // 5. command_runs.task_id / agent_id / tool_run_id
      {
        table: 'command_runs', fk_column: 'task_id', parent_table: 'tasks',
        nullable: true, action: 'reparent',
      },
      {
        table: 'command_runs', fk_column: 'agent_id', parent_table: 'agents',
        nullable: true, action: 'reparent',
      },
      {
        table: 'command_runs', fk_column: 'tool_run_id', parent_table: 'tool_runs',
        nullable: true, action: 'reparent',
      },
      // 6. workspaces.task_id / agent_id
      {
        table: 'workspaces', fk_column: 'task_id', parent_table: 'tasks',
        nullable: true, action: 'reparent',
      },
      {
        table: 'workspaces', fk_column: 'agent_id', parent_table: 'agents',
        nullable: true, action: 'reparent',
      },
      // 7. diagnostics.command_run_id / artifact_id
      {
        table: 'diagnostics', fk_column: 'command_run_id', parent_table: 'command_runs',
        nullable: true, action: 'reparent',
      },
      {
        table: 'diagnostics', fk_column: 'artifact_id', parent_table: 'artifacts',
        nullable: true, action: 'reparent',
      },
      // 8. evidence_refs foreign columns (all nullable)
      {
        table: 'evidence_refs', fk_column: 'task_id', parent_table: 'tasks',
        nullable: true, action: 'reparent',
      },
      {
        table: 'evidence_refs', fk_column: 'agent_id', parent_table: 'agents',
        nullable: true, action: 'reparent',
      },
      {
        table: 'evidence_refs', fk_column: 'tool_run_id', parent_table: 'tool_runs',
        nullable: true, action: 'reparent',
      },
      {
        table: 'evidence_refs', fk_column: 'command_run_id', parent_table: 'command_runs',
        nullable: true, action: 'reparent',
      },
      {
        table: 'evidence_refs', fk_column: 'artifact_id', parent_table: 'artifacts',
        nullable: true, action: 'reparent',
      },
      {
        table: 'evidence_refs', fk_column: 'diagnostic_id', parent_table: 'diagnostics',
        nullable: true, action: 'reparent',
      },
      {
        table: 'evidence_refs', fk_column: 'message_id', parent_table: 'messages',
        nullable: true, action: 'reparent',
      },
    ]

    if (!this.db) return report

    for (const check of fkChecks) {
      try {
        await this.runFkCheck(check, report)
      } catch (error) {
        report.errors.push(
          `FK check failed for ${check.table}.${check.fk_column}: ${error}`,
        )
      }
    }

    return report
  }

  /**
   * Execute a single FK-off invariant against the live DB.
   * Performs real UPDATE/DELETE; never just records.
   */
  private async runFkCheck(
    check: FkInvariant,
    report: OrphanReferenceReport,
  ): Promise<void> {
    const db = this.db
    if (!db) return

    // Skip invariants whose table does not exist in this session DB
    // (test schemas / fresh sessions may be missing some tables).
    if (!this.tableExists(db, check.table)) return
    if (!this.tableExists(db, check.parent_table)) return

    // Collect orphan rows. For nullable FKs we skip nulls; for non-nullable
    // we still allow the scan but rely on the DB schema to enforce NOT NULL.
    const whereNull = check.nullable ? `AND t.${check.fk_column} IS NOT NULL` : ''
    const orphanStmt = db.prepare(
      `SELECT t.id AS row_id, t.${check.fk_column} AS orphan_ref
       FROM ${check.table} t
       LEFT JOIN ${check.parent_table} o ON t.${check.fk_column} = o.id
       WHERE o.id IS NULL ${whereNull}`,
    )
    const orphans = orphanStmt.all() as Array<{ row_id: string; orphan_ref: string }>

    if (orphans.length === 0) return

    report.totalFound += orphans.length

    // Find a fallback parent id once per check (most-recent valid row).
    let fallbackParentId: string | null = null
    const needFallback = check.action === 'reparent'
    if (needFallback) {
      fallbackParentId = this.findRecentValidParent(db, check.parent_table)
    }

    const updateStmt = db.prepare(
      `UPDATE ${check.table} SET ${check.fk_column} = ? WHERE id = ?`,
    )
    const archiveStatusStmt = db.prepare(
      `UPDATE ${check.table} SET status = ? WHERE id = ?`,
    )
    const deleteStmt = db.prepare(`DELETE FROM ${check.table} WHERE id = ?`)

    for (const o of orphans) {
      if (check.action === 'reparent' && fallbackParentId) {
        updateStmt.run(fallbackParentId, o.row_id)
        report.reparented.push({
          table: check.table,
          id: o.row_id,
          new_parent_id: fallbackParentId,
        })
        continue
      }

      // No reparent possible, or action is archive. Apply real DB write.
      if (check.archive_status) {
        archiveStatusStmt.run(check.archive_status, o.row_id)
        report.archived.push({
          table: check.table,
          id: o.row_id,
          reason: `FK-off: ${check.fk_column} → ${check.parent_table}`,
        })
      } else {
        deleteStmt.run(o.row_id)
        report.archived.push({
          table: check.table,
          id: o.row_id,
          reason: `FK-off: ${check.fk_column} → ${check.parent_table} (deleted)`,
        })
      }
    }
  }

  /**
   * Return the most-recent row id in `parent_table`, or null if the table
   * is empty. Used as the fallback parent during reparent actions.
   */
  private findRecentValidParent(db: Database, parent_table: string): string | null {
    if (!this.tableExists(db, parent_table)) return null
    try {
      const row = db.prepare(
        `SELECT id FROM ${parent_table} ORDER BY id DESC LIMIT 1`,
      ).get() as { id?: string } | null
      return row?.id ?? null
    } catch {
      return null
    }
  }

  /**
   * Cheap existence check against sqlite_schema. Avoids throwing on test
   * schemas that omit some of the production tables.
   */
  private tableExists(db: Database, table: string): boolean {
    try {
      const row = db.prepare(
        `SELECT 1 FROM sqlite_master WHERE type='table' AND name = ?`,
      ).get(table) as { 1?: number } | null
      return row !== null
    } catch {
      return false
    }
  }

  /**
   * PID liveness check for running agents.
   * Uses Signal 0 (kill -0) to check process existence.
   */
  checkPidLiveness(agents?: Array<{ agent_id: string; pid: number }>): PidLivenessReport[] {
    if (!agents || agents.length === 0) {
      return []
    }

    const reports: PidLivenessReport[] = []

    for (const agent of agents) {
      let alive = false
      try {
        // Signal 0 does not kill the process; it checks if the process exists
        process.kill(agent.pid, 0)
        alive = true
      } catch {
        // ESRCH: no such process, or EPERM: no permission (process exists but not owned by us)
        alive = false
      }

      reports.push({
        agent_id: agent.agent_id,
        pid: agent.pid,
        alive,
        action: alive ? 'keep' : 'mark_lost'
      })
    }

    return reports
  }

  private findOrphanFiles(dir: string, depth = 0): string[] {
    const orphans: string[] = []
    if (depth > 5) return orphans

    try {
      const entries = readdirSync(dir)
      for (const entry of entries) {
        const fullPath = join(dir, entry)
        try {
          const stat = statSync(fullPath)
          if (stat.isDirectory()) {
            orphans.push(...this.findOrphanFiles(fullPath, depth + 1))
          } else if (this.isOrphanFile(entry, fullPath)) {
            orphans.push(fullPath)
          }
        } catch { /* skip inaccessible */ }
      }
    } catch { /* directory might not exist */ }
    return orphans
  }

  private isOrphanFile(filename: string, _path: string): boolean {
    return filename.endsWith('.tmp')
  }

  private looksLikeArtifact(filename: string): boolean {
    return filename.match(/^\d{17}Z-art_/) !== null
  }

  private quarantineFile(sourcePath: string): string {
    const filename = basename(sourcePath)
    const timestamp = new Date().toISOString().replace(/[-:]/g, '').replace(/\..+/, '')
    const quarantinedName = `${timestamp}-${filename}`
    const quarantinedPath = join(this.quarantineDir, quarantinedName)
    renameSync(sourcePath, quarantinedPath)
    return quarantinedPath
  }
}

export function createRecovery(options: RecoveryOptions): Recovery {
  return new Recovery(options)
}