AirCoding/packages/runtime/src/tools/ToolRegistry.ts

/**
 * ToolRegistry - tool lookup, validation, permission, execution
 *
 * Implements contracts §12; DD §9.1 + §9.3 branching table.
 * INV-3: all tool execution must go through this registry.
 *
 * @module packages/runtime/src/tools/ToolRegistry
 */

import type { ToolDefinition, ToolCall, ToolResultEnvelope, ISOTimeString } from '@aircoding/contracts'

import { PermissionEngine, createPermissionEngine, type PermissionContext, type PermissionDecision, type PermissionAction } from '../security/PermissionEngine.js'
import type { AgentType } from '@aircoding/contracts'
import { eventIngestor } from '../events/EventIngestor.js'
import { eventBus, type Subscription } from '../events/EventBus.js'

export type ToolExecutionReturn =
  | ToolResultEnvelope
  | Promise<ToolResultEnvelope>
  | AsyncIterable<ToolResultEnvelope>

export interface ToolExecutor {
  (call: ToolCall, context: ToolExecutionContext): ToolExecutionReturn
}

export interface ToolExecutionContext {
  session_id: string
  project_id: string
  project_root: string
  agent_id: string
  agent_type: AgentType
  task_id?: string
  task_scope?: PermissionContext['task_scope']
  permission_profile?: PermissionContext['permission_profile']
}

export interface ToolCallContext {
  tool_definition: ToolDefinition
  executor: ToolExecutor
  permission_context: PermissionContext
}

/**
 * Global tool registry (singleton)
 */
let global_tool_registry: ToolRegistry | undefined

export class ToolRegistry {
  private tools: Map<string, ToolDefinition> = new Map()
  private executors: Map<string, ToolExecutor> = new Map()
  private permission_engine: PermissionEngine
  private project_root: string
  private readonly permission_timeout_ms = 5 * 60 * 1000

  constructor(project_root: string) {
    this.project_root = project_root
    this.permission_engine = createPermissionEngine(project_root)
    global_tool_registry = this
  }

  /**
   * Register a tool with its definition and executor.
   */
  register(name: string, definition: ToolDefinition, executor: ToolExecutor): void {
    this.tools.set(name, definition)
    this.executors.set(name, executor)
  }

  /**
   * Unregister a tool.
   */
  unregister(name: string): void {
    this.tools.delete(name)
    this.executors.delete(name)
  }

  /**
   * Get tool definition by name.
   */
  get(name: string): ToolDefinition | undefined {
    return this.tools.get(name)
  }

  /**
   * List all registered tools.
   */
  list(): ToolDefinition[] {
    return Array.from(this.tools.values())
  }

  /**
   * Call a tool with permission evaluation and branching.
   * Implements DD §9.1 algorithm.
   */
  async call(call: ToolCall, context: ToolExecutionContext): Promise<ToolResultEnvelope> {
    // Step 1: Lookup tool definition
    const definition = this.tools.get(call.name)
    if (!definition) {
      return create_error_result(call.call_id, 'tool_not_found', `Tool ${call.name} not found`)
    }

    // Step 2: Validate input schema
    const validation = this.validate_input(call, definition)
    if (!validation.valid) {
      return create_error_result(call.call_id, 'invalid_input', validation.error || 'Invalid input')
    }

    // Step 3: Build permission context
    const permission_context = this.build_permission_context(call, context)

    // Step 4: Evaluate permissions (layered per DD §9.2)
    const decision = await this.permission_engine.evaluate(call, permission_context, definition)

    // B.2: Emit tool.started event before execution
    const startedAt = new Date().toISOString() as ISOTimeString
    const eventPayload = {
      tool_run_id: call.call_id,
      call_id: call.call_id,
      tool_name: call.name,
      input_json: call.arguments || {},
      started_at: startedAt,
      agent_id: context.agent_id,
      task_id: context.task_id,
      session_id: context.session_id,
    }
    try {
      await eventIngestor.ingest({
        id: `tool_start_${call.call_id}_${Date.now()}`,
        type: 'tool.started',
        version: 1,
        timestamp: startedAt,
        session_id: context.session_id,
        project_id: context.project_id,
        source: { kind: 'tool' as const, id: context.agent_id },
        route: [],
        payload: eventPayload,
      })
    } catch (e) {
      // Log but don't fail tool execution if event emission fails
      console.warn('Failed to emit tool.started event:', e)
    }

    // Step 5: Branch on permission action (DD §9.3)
    // Step 6: Execute branch
    let result: ToolResultEnvelope
    try {
      result = await this.execute_branch(decision, call, context)
    } catch (error) {
      // B.2: Emit tool.failed event on exception
      const failedAt = new Date().toISOString() as ISOTimeString
      const errorPayload = {
        call_id: call.call_id,
        tool_name: call.name,
        completed_at: failedAt,
        exit_code: 'error',
        output_kind: 'error',
        error: error instanceof Error ? error.message : String(error),
      }
      try {
        await eventIngestor.ingest({
          id: `tool_fail_${call.call_id}_${Date.now()}`,
          type: 'tool.failed',
          version: 1,
          timestamp: failedAt,
          session_id: context.session_id,
          project_id: context.project_id,
          source: { kind: 'tool' as const, id: context.agent_id },
          route: [],
          payload: errorPayload,
        })
      } catch (e) {
        console.warn('Failed to emit tool.failed event:', e)
      }
      return create_error_result(call.call_id, 'execution_error', error instanceof Error ? error.message : String(error))
    }

    // B.2: Emit tool.completed event on success
    const completedAt = new Date().toISOString() as ISOTimeString
    const completedPayload = {
      call_id: call.call_id,
      tool_name: call.name,
      completed_at: completedAt,
      exit_code: 'ok',
      output_kind: 'text',
    }
    try {
      await eventIngestor.ingest({
        id: `tool_end_${call.call_id}_${Date.now()}`,
        type: 'tool.completed',
        version: 1,
        timestamp: completedAt,
        session_id: context.session_id,
        project_id: context.project_id,
        source: { kind: 'tool' as const, id: context.agent_id },
        route: [],
        payload: completedPayload,
      })
    } catch (e) {
      console.warn('Failed to emit tool.completed event:', e)
    }

    // Step 7: Record decision (if enabled)
    await this.permission_engine.record(decision, context)

    return result
  }

  /**
   * Streaming call - returns chunks for real-time output.
   * Ends with exactly one final ToolResultEnvelope.
   */
  async *call_streaming(call: ToolCall, context: ToolExecutionContext): AsyncGenerator<ToolResultEnvelope> {
    const definition = this.tools.get(call.name)
    if (!definition?.streaming) {
      // Non-streaming tool, call normally and yield single result
      const result = await this.call(call, context)
      yield result
      return
    }

    // For streaming tools, we need to get the executor
    const executor = this.executors.get(call.name)
    if (!executor) {
      yield create_error_result(call.call_id, 'executor_not_found', 'Executor not registered')
      return
    }

    // Permission check first (same branching as call)
    const permission_context = this.build_permission_context(call, context)
    const decision = await this.permission_engine.evaluate(call, permission_context, definition)

    // Emit tool.started event before execution
    const startedAt = new Date().toISOString() as ISOTimeString
    try {
      await eventIngestor.ingest({
        id: `tool_start_${call.call_id}_${Date.now()}`,
        type: 'tool.started',
        version: 1,
        timestamp: startedAt,
        session_id: context.session_id,
        project_id: context.project_id,
        source: { kind: 'tool' as const, id: context.agent_id },
        route: [],
        payload: {
          call_id: call.call_id,
          tool_name: call.name,
          started_at: startedAt,
          agent_id: context.agent_id,
          task_id: context.task_id,
          session_id: context.session_id,
        },
      })
    } catch (e) {
      console.warn('Failed to emit tool.started event (streaming):', e)
    }

    // Full permission branching (same as execute_branch)
    switch (decision.action) {
      case 'deny':
        yield create_error_result(call.call_id, 'permission_denied', decision.reason)
        return
      case 'block':
        yield create_error_result(call.call_id, 'blocked', `Action blocked: ${decision.reason}`)
        return
      case 'refuse':
        yield create_error_result(call.call_id, 'policy_error', `Refused: ${decision.reason}`)
        return
      case 'ask_user': {
        const prompt_id = `perm_${crypto.randomUUID()}`
        try {
          await eventIngestor.ingest({
            id: `evt_${prompt_id}`,
            type: 'permission.prompt.requested',
            version: 1,
            session_id: context.session_id,
            project_id: context.project_id,
            timestamp: new Date().toISOString(),
            source: { kind: 'tool', id: call.name },
            route: ['tool_registry', 'permission'],
            payload: {
              prompt_id,
              subject: call.name,
              risk_level: decision.risk_level,
              reason: decision.reason,
              options: ['allow_once', 'deny'],
              default_option: 'deny',
              request_ref: { call_id: call.call_id, tool_name: call.name, agent_id: context.agent_id },
            },
          })
        } catch (e) {
          console.warn('Failed to emit permission.prompt.requested:', e)
        }
        const selected = await this.wait_for_permission(prompt_id, context)
        if (selected !== 'allow_once' && selected !== 'allow') {
          yield create_error_result(call.call_id, 'permission_denied', `User selected ${selected}`)
          return
        }
        break
      }
      case 'announce_then_run':
        // Fall through to execution with announced metadata
        break
      case 'allow':
        break
      default:
        yield create_error_result(call.call_id, 'invalid_decision', `Unknown action: ${decision.action}`)
        return
    }

    let saw_final = false
    let lastError: Error | undefined

    try {
      for await (const chunk of this.execute_streaming(call, context, executor)) {
        if (chunk.metadata && (chunk.metadata as any).is_final === true) saw_final = true
        if (decision.action === 'announce_then_run' && chunk.metadata) {
          (chunk.metadata as any).announced = true
        }
        yield chunk
      }
    } catch (error) {
      lastError = error instanceof Error ? error : new Error(String(error))
      yield create_error_result(call.call_id, 'execution_error', lastError.message)
    }

    // Emit tool.completed or tool.failed event
    const endAt = new Date().toISOString() as ISOTimeString
    const endEventType = lastError ? 'tool.failed' : 'tool.completed'
    try {
      await eventIngestor.ingest({
        id: `tool_${lastError ? 'fail' : 'end'}_${call.call_id}_${Date.now()}`,
        type: endEventType,
        version: 1,
        timestamp: endAt,
        session_id: context.session_id,
        project_id: context.project_id,
        source: { kind: 'tool' as const, id: context.agent_id },
        route: [],
        payload: {
          call_id: call.call_id,
          tool_name: call.name,
          completed_at: endAt,
          exit_code: lastError ? 'error' : 'ok',
          output_kind: lastError ? 'error' : 'text',
          ...(lastError ? { error: lastError.message } : {}),
        },
      })
    } catch (e) {
      console.warn(`Failed to emit ${endEventType} event (streaming):`, e)
    }

    if (!saw_final && !lastError) {
      yield create_error_result(call.call_id, 'no_final_result', 'Streaming tool did not produce final result')
    }

    await this.permission_engine.record(decision, context)
  }

  /**
   * Validate tool input against schema.
   */
  private validate_input(call: ToolCall, definition: ToolDefinition): { valid: boolean; error?: string } {
    // Basic validation - in production, use JSON Schema validation
    if (!call.arguments || typeof call.arguments !== 'object') {
      return { valid: false, error: 'Arguments must be an object' }
    }

    // Check required fields if specified in definition
    // This is a simplified check - full implementation would use JSON Schema
    return { valid: true }
  }

  /**
   * Build permission context from tool call and execution context.
   */
  private build_permission_context(call: ToolCall, context: ToolExecutionContext): PermissionContext {
    return {
      session_id: context.session_id,
      project_id: context.project_id,
      project_root: context.project_root,
      agent_type: context.agent_type,
      agent_id: context.agent_id,
      task_scope: context.task_scope,
      permission_profile: context.permission_profile,
    }
  }

  /**
   * Downgrade write operations to read-only.
   */
  private downgrade_to_readonly(call: ToolCall): ToolCall {
    // Modify arguments to make operation read-only
    const modified = { ...call.arguments }

    // For filesystem operations, remove write-related flags
    if ('mode' in modified && typeof modified.mode === 'string') {
      if (modified.mode.includes('w') || modified.mode.includes('a')) {
        modified.mode = modified.mode.replace(/[wa]/g, 'r')
      }
    }

    // Remove force flags
    delete modified.force
    delete modified.overwrite

    return { ...call, arguments: modified }
  }

  /**
   * Apply sandbox restrictions to arguments.
   */
  private apply_sandbox_restrictions(args: Record<string, unknown>, flags: string[]): Record<string, unknown> {
    const restricted = { ...args }

    // Add sandbox restrictions based on flags
    if (flags.includes('no_network')) {
      delete restricted.url
      delete restricted.endpoint
    }

    if (flags.includes('read_only')) {
      // Already handled in read_only branch
    }

    // Add sandbox metadata
    return restricted
  }

  /**
   * Execute branching behavior per DD §9.3.
   * Replaces the module-level ACTION_BRANCHES to fix `this` binding.
   */
  private async execute_branch(
    decision: PermissionDecision,
    call: ToolCall,
    ctx: ToolExecutionContext,
  ): Promise<ToolResultEnvelope> {
    switch (decision.action) {
      case 'allow': {
        const executor = this.executors.get(call.name)
        if (!executor) {
          return create_error_result(call.call_id, 'executor_not_found', 'Executor not registered')
        }

        // FR-011: Handle backup requirement for out-of-project writes
        if ((decision as any).backup_required && (decision as any).backup_path) {
          await this.perform_backup(call, (decision as any).backup_path)
        }

        return this.execute_executor_final(executor, call, ctx)
      }

      case 'announce_then_run': {
        // Emit visible notice, then execute unless interrupted
        const executor = this.executors.get(call.name)
        if (!executor) {
          return create_error_result(call.call_id, 'executor_not_found', 'Executor not registered')
        }

        // FR-011: Handle backup requirement for out-of-project writes
        if ((decision as any).backup_required && (decision as any).backup_path) {
          await this.perform_backup(call, (decision as any).backup_path)
        }

        const result = await this.execute_executor_final(executor, call, ctx)
        return {
          ...result,
          metadata: { ...result.metadata, announced: true },
        }
      }

      case 'ask_user': {
        const prompt_id = `perm_${crypto.randomUUID()}`
        await eventIngestor.ingest({
          id: `evt_${prompt_id}`,
          type: 'permission.prompt.requested',
          version: 1,
          session_id: ctx.session_id,
          project_id: ctx.project_id,
          timestamp: new Date().toISOString(),
          source: { kind: 'tool', id: call.name },
          route: ['tool_registry', 'permission'],
          payload: {
            prompt_id,
            subject: call.name,
            risk_level: decision.risk_level,
            reason: decision.reason,
            options: ['allow_once', 'deny'],
            default_option: 'deny',
            request_ref: { call_id: call.call_id, tool_name: call.name, agent_id: ctx.agent_id },
          },
        })

        const selected = await this.wait_for_permission(prompt_id, ctx)
        if (selected !== 'allow_once' && selected !== 'allow') {
          return create_error_result(call.call_id, 'permission_denied', `User selected ${selected}`)
        }

        const executor = this.executors.get(call.name)
        if (!executor) {
          return create_error_result(call.call_id, 'executor_not_found', 'Executor not registered')
        }
        return this.execute_executor_final(executor, call, ctx)
      }

      case 'deny':
        return create_error_result(call.call_id, 'permission_denied', decision.reason)

      case 'block': {
        // Return blocked outcome → task.blocked upstream
        return create_error_result(call.call_id, 'blocked', `Action blocked: ${decision.reason}`)
      }

      case 'refuse': {
        // Return policy error; no execution
        return create_error_result(call.call_id, 'policy_error', `Refused: ${decision.reason}`)
      }

      default:
        return create_error_result(call.call_id, 'invalid_decision', `Unknown action: ${decision.action}`)
    }
  }

  /**
   * Execute a tool and return the final envelope.
   * Streaming executors are consumed until their final result.
   */
  private async execute_executor_final(
    executor: ToolExecutor,
    call: ToolCall,
    context: ToolExecutionContext,
  ): Promise<ToolResultEnvelope> {
    const result = executor(call, context)
    if (this.is_async_iterable(result)) {
      let final_result: ToolResultEnvelope | undefined
      let last_chunk: ToolResultEnvelope | undefined
      for await (const chunk of result) {
        last_chunk = chunk
        if (chunk.metadata && (chunk.metadata as any).is_final === true) {
          final_result = chunk
        }
      }
      return final_result || last_chunk || create_error_result(call.call_id, 'no_result', 'Tool produced no result')
    }
    return await result
  }

  private is_async_iterable(value: unknown): value is AsyncIterable<ToolResultEnvelope> {
    return Boolean(value && typeof (value as any)[Symbol.asyncIterator] === 'function')
  }

  private wait_for_permission(prompt_id: string, ctx: ToolExecutionContext): Promise<string> {
    return new Promise((resolve) => {
      let settled = false
      let subscription: Subscription | undefined
      const finish = (selected: string) => {
        if (settled) return
        settled = true
        clearTimeout(timeout)
        if (subscription) eventBus.unsubscribe(subscription)
        resolve(selected)
      }

      const timeout = setTimeout(() => {
        void eventIngestor.ingest({
          id: `evt_${prompt_id}_timeout`,
          type: 'permission.prompt.resolved',
          version: 1,
          session_id: ctx.session_id,
          project_id: ctx.project_id,
          timestamp: new Date().toISOString(),
          source: { kind: 'tool', id: 'permission_timeout' },
          route: ['tool_registry', 'permission'],
          payload: {
            prompt_id,
            selected_option: 'deny',
            decision_id: `decision_${crypto.randomUUID()}`,
            resolved_by: 'timeout',
          },
        }).catch(() => finish('deny'))
      }, this.permission_timeout_ms)

      subscription = eventBus.subscribe({ session_id: ctx.session_id, types: ['permission.prompt.resolved'] }, (event) => {
        const payload = event.payload as Record<string, unknown>
        if (payload.prompt_id !== prompt_id) return
        finish(String(payload.selected_option || 'deny'))
      })
    })
  }

  /**
   * Execute streaming tool.
   */
  private async *execute_streaming(
    call: ToolCall,
    context: ToolExecutionContext,
    executor: ToolExecutor
  ): AsyncGenerator<ToolResultEnvelope> {
    const result = executor(call, context)
    if (this.is_async_iterable(result)) {
      for await (const chunk of result) yield chunk
      return
    }
    yield await result
  }

  /**
   * FR-011: Perform backup before out-of-project write.
   * Creates backup in .air/local/backups/
   */
  private async perform_backup(call: ToolCall, backup_path: string): Promise<void> {
    try {
      const { mkdirSync, existsSync, cpSync } = await import('fs')
      const { dirname } = await import('path')

      // Extract source path from tool call
      const source_path = this.extract_source_path(call)
      if (!source_path || !existsSync(source_path)) {
        console.warn('Backup skipped: source file does not exist:', source_path)
        return
      }

      // Ensure backup directory exists
      const backup_dir = dirname(backup_path)
      if (!existsSync(backup_dir)) {
        mkdirSync(backup_dir, { recursive: true })
      }

      // Copy source to backup location
      cpSync(source_path, backup_path)

      // Emit backup event
      await eventIngestor.ingest({
        id: `backup_${Date.now()}`,
        type: 'file.backup.created',
        version: 1,
        timestamp: new Date().toISOString(),
        session_id: call.call_id, // Use call_id as session_id placeholder
        project_id: this.project_root,
        source: { kind: 'tool', id: 'tool_registry' },
        route: ['tool_registry', 'backup'],
        payload: {
          original_path: source_path,
          backup_path,
          tool_name: call.name,
          call_id: call.call_id
        }
      })
    } catch (e) {
      console.error('Backup failed:', e)
      // Continue with operation even if backup fails - log but don't block
    }
  }

  /**
   * Extract source file path from tool call for backup.
   */
  private extract_source_path(call: ToolCall): string | null {
    const args = call.arguments as Record<string, unknown>
    const path_keys = ['path', 'file', 'file_path', 'source', 'target']

    for (const key of path_keys) {
      if (typeof args[key] === 'string') {
        return args[key] as string
      }
    }
    return null
  }
}

export function createToolRegistry(project_root: string): ToolRegistry {
  return new ToolRegistry(project_root)
}

// ============================================================================
// Result helpers
// ============================================================================

function create_error_result(call_id: string, error_type: string, message: string): ToolResultEnvelope {
  return {
    status: 'error',
    error: {
      error_id: call_id,
      kind: error_type === 'not_found' ? 'unknown_error' : 'tool_error',
      severity: 'error',
      message,
      retryability: 'not_retryable',
      semantic_signature: error_type,
    },
    metadata: { timestamp: new Date().toISOString() as ISOTimeString, call_id }
  }
}