AirCoding/packages/runtime/src/context/ContextAssembler.ts

/**
 * ContextAssembler - Assembles prompt layers into Anthropic-canonical context
 *
 * Implements contracts §16; DD §10.1 + §10.2 layer-assembly table.
 *
 * @module packages/runtime/src/context/ContextAssembler
 */

import type {
  AgentType, PromptLayer, BudgetFitResult,
  SessionID, ProjectID, AgentID, TaskID, ArtifactID, ISOTimeString
} from '@aircoding/contracts'

import { PromptLayerLoader, createPromptLayerLoader } from './PromptLayerLoader.js'
import { CompactionPolicy, createCompactionPolicy } from './CompactionPolicy.js'

export interface AssembledContext {
  messages: AssembledMessage[]
  metadata: AssemblyMetadata
}

export interface AssembledMessage {
  role: 'system' | 'user' | 'assistant'
  content: string
  layer?: string
}

export interface AssemblyMetadata {
  total_tokens: number
  fitted_layers: string[]
  omitted_layers: string[]
  compaction_requested: boolean
  layers_compacted: boolean
  omissions: string[]
  messages_artifact_id?: string
  assembled_at: ISOTimeString
}

export interface AssemblyContext {
  session_id: SessionID
  project_id: ProjectID
  project_root: string
  agent_id: AgentID
  agent_type: AgentType
  task_id?: TaskID
  token_budget?: number
  additional_layers?: PromptLayer[]
}

export class ContextAssembler {
  private loader: PromptLayerLoader
  private policy: CompactionPolicy
  private message_repo?: any
  private evidence_store?: any

  constructor(loader?: PromptLayerLoader, policy?: CompactionPolicy) {
    this.loader = loader || createPromptLayerLoader()
    this.policy = policy || createCompactionPolicy()
  }

  /**
   * Inject database-backed data sources for L6/L7/L8 real content.
   * Without these, layers use descriptive placeholder text.
   */
  set_data_sources(sources: { message_repo?: any; evidence_store?: any }): void {
    this.message_repo = sources.message_repo
    this.evidence_store = sources.evidence_store
  }

  /**
   * Assemble context from all layers.
   * Returns Anthropic-canonical AssembledContext.
   */
  assemble(context: AssemblyContext): AssembledContext {
    const token_budget = context.token_budget || 200000
    const warnings: string[] = []

    // Collect all layers
    const layers = this.collect_layers(context)

    // Fit into budget
    const fit_result = this.policy.fit_to_budget(layers, token_budget)
    if (fit_result.omissions.length > 0) {
      warnings.push(...fit_result.omissions)
    }

    // Build messages from fitted layers
    const messages = this.build_messages(fit_result.fitted, context)

    // Check if compaction is needed
    const compaction_check = this.policy.should_compact(layers, fit_result.total_tokens)

    return {
      messages,
      metadata: {
        total_tokens: fit_result.total_tokens,
        fitted_layers: fit_result.fitted.map(l => l.level),
        omitted_layers: fit_result.omitted.map(l => l.level),
        compaction_requested: compaction_check.should_compact,
        layers_compacted: false,
        omissions: fit_result.omissions,
        assembled_at: new Date().toISOString() as ISOTimeString
      }
    }
  }

  /**
   * Collect all layers in order L0-L9.
   */
  private collect_layers(context: AssemblyContext): PromptLayer[] {
    const layers: PromptLayer[] = []

    // L0: Runtime invariants (ALWAYS FIRST, never omitted)
    const l0 = this.loader.load_runtime_invariant()
    layers.push(l0)

    // L1: Role (worker AgentType)
    const l1 = this.loader.load_role(context.agent_type)
    layers.push(l1)

    // L2: Safety - built-in
    layers.push({
      level: 'safety' as any,
      priority: 2,
      content: '# Safety Rules\n- Never execute destructive commands\n- Always validate inputs\n- Report errors immediately',
      token_estimate: 50
    })

    // L3: Project rules
    const project_rules = this.loader.load_project_rules({
      project_id: context.project_id,
      project_root: context.project_root
    })
    layers.push(...project_rules)

    // L4: Architecture (if available)
    if (context.additional_layers) {
      const arch_layers = context.additional_layers.filter(l => l.level === 'architecture')
      layers.push(...arch_layers)
    }

    // L5: Task spec (if task_id provided)
    if (context.task_id) {
      const task_layers = this.loader.load_task_context(
        {
          id: context.task_id,
          type: 'execute',
          title: 'Current Task',
          description: 'Task from session context',
          acceptance_criteria: ['Task completed successfully']
        },
        {}
      )
      layers.push(...task_layers)
    }

    // L6: Evidence — try EvidenceStore if available, else descriptive
    const evidence_layers = context.additional_layers?.filter(l => l.level === 'evidence') || []
    if (evidence_layers.length > 0) {
      layers.push(...evidence_layers)
    } else {
      let evidence_content = ''
      if (this.evidence_store && context.task_id) {
        try {
          const records = this.evidence_store.list_for_entity?.(context.task_id) || []
          if (records.length > 0) {
            evidence_content = records.map((r: any) =>
              `- [${r.type || 'evidence'}] ${r.summary || r.id}`).join('\n')
          }
        } catch { /* fall through to descriptive */ }
      }
      layers.push({
        level: 'evidence' as any,
        priority: 6,
        content: evidence_content || [
          '# Evidence Context (L6)',
          `Session: ${context.session_id}`,
          context.task_id ? `Task: ${context.task_id}` : '',
          'No evidence records available for this task.',
        ].filter(Boolean).join('\n'),
        token_estimate: evidence_content ? evidence_content.length / 4 : 80,
        source_ref: `session:${context.session_id}:evidence`
      })
    }

    // L7: Conversation history — try MessageRepository if available
    const conv_layers = context.additional_layers?.filter(l => l.level === 'conversation') || []
    if (conv_layers.length > 0) {
      layers.push(...conv_layers)
    } else {
      let conv_content = ''
      if (this.message_repo) {
        try {
          const messages = this.message_repo.list_by_session?.(context.session_id) || []
          conv_content = messages.slice(-20).map((m: any) =>
            `[${m.role}]: ${String(m.content_json || m.content || '').slice(0, 200)}`).join('\n')
        } catch { /* fall through */ }
      }
      layers.push({
        level: 'conversation' as any,
        priority: 7,
        content: conv_content || [
          '# Conversation History (L7)',
          `Session: ${context.session_id}`,
          'No message history available.',
        ].join('\n'),
        token_estimate: conv_content ? conv_content.length / 4 : 60,
        source_ref: `session:${context.session_id}:messages`
      })
    }

    // L8: Recent tool outputs — try DB if available
    const tool_layers = context.additional_layers?.filter(l => l.level === 'tool_output') || []
    if (tool_layers.length > 0) {
      layers.push(...tool_layers)
    } else {
      let tool_content = ''
      if (this.message_repo) {
        try {
          const msgs = this.message_repo.list_by_session?.(context.session_id) || []
          const tool_msgs = msgs.filter((m: any) => m.role === 'tool_result' || m.role === 'tool_use').slice(-10)
          tool_content = tool_msgs.map((m: any) =>
            `[${m.role}]: ${String(m.content_json || '').slice(0, 300)}`).join('\n')
        } catch { /* fall through */ }
      }
      layers.push({
        level: 'tool_output' as any,
        priority: 8,
        content: tool_content || [
          '# Recent Tool Outputs (L8)',
          'No tool output history available.',
        ].join('\n'),
        token_estimate: tool_content ? tool_content.length / 4 : 50,
        source_ref: `session:${context.session_id}:tool_outputs`
      })
    }

    // L9: User override from additional_layers
    const user_layers = context.additional_layers?.filter(l => l.level === 'user_override') || []
    if (user_layers.length > 0) {
      layers.push(...user_layers)
    } else {
      layers.push({
        level: 'user_override',
        priority: 9,
        content: '# User Overrides (L9)\n// No user overrides active',
        token_estimate: 15
      })
    }

    // Add any additional layers
    if (context.additional_layers) {
      const others = context.additional_layers.filter(l => l.level !== 'architecture')
      layers.push(...others)
    }

    return layers
  }

  /**
   * Build assembled messages from fitted layers.
   */
  private build_messages(layers: PromptLayer[], context: AssemblyContext): AssembledMessage[] {
    const messages: AssembledMessage[] = []

    // System message: L0 + L1 + L2 + L3
    const system_content = layers
      .filter(l => ['runtime_invariant', 'role', 'safety', 'project_rules'].includes(l.level))
      .map(l => l.content)
      .join('\n\n---\n\n')

    if (system_content) {
      messages.push({ role: 'system', content: system_content, layer: 'system' })
    }

    // Architecture context
    const arch_layers = layers.filter(l => l.level === 'architecture')
    for (const l of arch_layers) {
      messages.push({ role: 'user', content: String(l.content), layer: 'architecture' })
    }

    // Task spec
    const task_layers = layers.filter(l => l.level === 'task_spec')
    for (const l of task_layers) {
      messages.push({ role: 'user', content: String(l.content), layer: 'task_spec' })
    }

    // Evidence
    const evidence_layers = layers.filter(l => l.level === 'evidence')
    for (const l of evidence_layers) {
      messages.push({ role: 'user', content: String(l.content), layer: 'evidence' })
    }

    // Conversation (L7)
    const conv_layers = layers.filter(l => l.level === 'conversation')
    for (const l of conv_layers) {
      messages.push({ role: 'assistant', content: String(l.content), layer: 'conversation' })
    }

    // Tool output (L8)
    const tool_layers = layers.filter(l => l.level === 'tool_output')
    for (const l of tool_layers) {
      messages.push({ role: 'user', content: `[Tool Output]\n${l.content}`, layer: 'tool_output' })
    }

    return messages
  }
}

export function createContextAssembler(loader?: PromptLayerLoader, policy?: CompactionPolicy): ContextAssembler {
  return new ContextAssembler(loader, policy)
}