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'

export interface ToolExecutor {
  (call: ToolCall, context: ToolExecutionContext): Promise<ToolResultEnvelope>
}

export interface ToolExecutionContext {
  session_id: string
  project_id: string
  project_root: string
  agent_id: string
  agent_type: AgentType
}

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

/**
 * Branching behavior per DD §9.3
 */
const ACTION_BRANCHES: Record<PermissionAction, (decision: PermissionDecision, call: ToolCall, ctx: ToolExecutionContext) => Promise<ToolResultEnvelope>> = {
  allow: async (_decision, call, ctx) => {
    // Execute directly
    const definition = global_tool_registry?.get(call.name)
    if (!definition) {
      return create_error_result(call.id, 'tool_not_found', 'Tool not found')
    }
    const executor = global_tool_registry?.executors.get(call.name)
    if (!executor) {
      return create_error_result(call.id, 'executor_not_found', 'Executor not registered')
    }
    return executor(call, ctx)
  },

  deny: async (decision) => {
    return create_error_result('', 'permission_denied', decision.reason)
  },

  prompt: async (_decision, _call, _ctx) => {
    // TODO: Integrate with UI for user prompt
    // For now, deny with prompt message
    return create_error_result('', 'user_prompt_required', 'User confirmation required')
  },

  read_only: async (_decision, call, ctx) => {
    // Downgrade write operations to read-only
    const modified_call = this.downgrade_to_readonly(call)
    const definition = global_tool_registry?.get(call.name)
    const executor = global_tool_registry?.executors.get(call.name)
    if (!executor) {
      return create_error_result(call.id, 'executor_not_found', 'Executor not registered')
    }
    return executor(modified_call as ToolCall, ctx)
  },

  sandbox: async (decision, call, ctx) => {
    // Execute in sandboxed mode with restricted environment
    const sandboxed_call = {
      ...call,
      arguments: this.apply_sandbox_restrictions(call.arguments, decision.flags)
    }
    const executor = global_tool_registry?.executors.get(call.name)
    if (!executor) {
      return create_error_result(call.id, 'executor_not_found', 'Executor not registered')
    }
    return executor(sandboxed_call, ctx)
  },

  audit_log: async (_decision, call, ctx) => {
    // Execute and log for audit
    const definition = global_tool_registry?.get(call.name)
    const executor = global_tool_registry?.executors.get(call.name)
    if (!executor) {
      return create_error_result(call.id, 'executor_not_found', 'Executor not registered')
    }
    const result = await executor(call, ctx)
    // Add audit flag to result
    return {
      ...result,
      metadata: { ...result.metadata, audit_logged: true }
    }
  }
}

/**
 * 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

  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.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.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)

    // Step 5: Branch on permission action (DD §9.3)
    const branch = ACTION_BRANCHES[decision.action]
    if (!branch) {
      return create_error_result(call.id, 'invalid_decision', 'Invalid permission decision')
    }

    // Step 6: Execute branch
    try {
      const result = await branch(decision, call, context)

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

      return result
    } catch (error) {
      return create_error_result(call.id, 'execution_error', error instanceof Error ? error.message : String(error))
    }
  }

  /**
   * 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.id, 'executor_not_found', 'Executor not registered')
      return
    }

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

    if (decision.action !== 'allow') {
      yield create_error_result(call.id, 'permission_denied', decision.reason)
      return
    }

    // Execute with streaming support
    // The executor yields intermediate results, final result comes at end
    let final_result: ToolResultEnvelope | undefined

    for await (const chunk of this.execute_streaming(call, context, executor)) {
      if (chunk.type === 'final') {
        final_result = chunk
      } else {
        yield chunk
      }
    }

    // Yield final result exactly once
    if (final_result) {
      yield final_result
    } else {
      yield create_error_result(call.id, 'no_final_result', 'Streaming tool did not produce final result')
    }
  }

  /**
   * 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: undefined, // Would be loaded from task context
      permission_profile: undefined // Would be loaded from agent config
    }
  }

  /**
   * 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 streaming tool.
   */
  private async *execute_streaming(
    call: ToolCall,
    context: ToolExecutionContext,
    executor: ToolExecutor
  ): AsyncGenerator<ToolResultEnvelope> {
    // This is a placeholder - actual implementation would depend on the tool
    // For now, just execute normally
    const result = await executor(call, context)
    yield result
  }
}

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 {
    call_id,
    tool_name: '',
    type: 'error',
    content: { error_type, message },
    metadata: { timestamp: new Date().toISOString() as ISOTimeString }
  }
}