AirCoding/packages/runtime/src/storage/repositories/MessageDraftRepository.ts

/**
 * MessageDraftRepository - CRUD + upsert, delete_for_message for message_drafts table (§5)
 *
 * Implements Repository<MessageDraftRecord, MessageDraftInsert, MessageDraftUpdate>
 * per contracts §6.
 *
 * @module packages/runtime/src/storage/repositories/MessageDraftRepository
 */

import type {
  Repository,
  TransactionHandle,
  SessionID,
  MessageID,
  ISOTimeString,
} from '@aircoding/contracts'

import { DatabaseHandle } from '../MigrationRunner.js'
import { assertEnumValues } from '../assertEnum.js'

// =============================================================================
// Types - per db-schema §5
// =============================================================================

export interface MessageDraftRecord {
  message_id: MessageID
  session_id: SessionID
  role: 'user' | 'assistant' | 'system' | 'tool'
  canonical_format: 'anthropic'
  partial_content_json: string
  status: 'streaming' | 'interrupted' | 'error'
  created_at: ISOTimeString
  updated_at: ISOTimeString
  metadata_json?: string
}

export type MessageDraftInsert = Omit<MessageDraftRecord, 'message_id'> & {
  message_id?: MessageID
}

export type MessageDraftUpdate = Partial<Omit<MessageDraftRecord, 'message_id' | 'session_id' | 'created_at'>>

// =============================================================================
// MessageDraftRepository
// =============================================================================

export class MessageDraftRepository implements Repository<MessageDraftRecord, MessageDraftInsert, MessageDraftUpdate> {
  private db: DatabaseHandle

  constructor(db: DatabaseHandle) {
    this.db = db
  }

  /**
   * Get a draft by message ID.
   */
  async get(message_id: MessageID, tx?: TransactionHandle): Promise<MessageDraftRecord | undefined> {
    const stmt = (tx?.db ?? this.db).prepare('SELECT * FROM message_drafts WHERE message_id = ?')
    const row = stmt.get(message_id) as MessageDraftRecord | undefined
    return row
  }

  /**
   * Insert a new draft.
   */
  async insert(record: MessageDraftInsert, tx?: TransactionHandle): Promise<void> {
    // Validate enum columns (only status is a closed enum for message_drafts)
    assertEnumValues('message_drafts', {
      status: record.status,
    })

    const stmt = (tx?.db ?? this.db).prepare(`
      INSERT INTO message_drafts (
        message_id, session_id, role, canonical_format,
        partial_content_json, status, created_at, updated_at, metadata_json
      ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
    `)

    stmt.run(
      record.message_id,
      record.session_id,
      record.role,
      record.canonical_format,
      record.partial_content_json,
      record.status,
      record.created_at,
      record.updated_at,
      record.metadata_json ?? null,
    )
  }

  /**
   * Update an existing draft.
   */
  async update(message_id: MessageID, patch: MessageDraftUpdate, tx?: TransactionHandle): Promise<void> {
    // Validate enum columns if present (only status is a closed enum for message_drafts)
    if (patch.status !== undefined) {
      assertEnumValues('message_drafts', { status: patch.status })
    }

    const fields: string[] = []
    const values: unknown[] = []

    if (patch.partial_content_json !== undefined) {
      fields.push('partial_content_json = ?')
      values.push(patch.partial_content_json)
    }
    if (patch.status !== undefined) {
      fields.push('status = ?')
      values.push(patch.status)
    }
    if (patch.updated_at !== undefined) {
      fields.push('updated_at = ?')
      values.push(patch.updated_at)
    }
    if (patch.metadata_json !== undefined) {
      fields.push('metadata_json = ?')
      values.push(patch.metadata_json)
    }

    if (fields.length === 0) {
      return // Nothing to update
    }

    values.push(message_id)
    const stmt = (tx?.db ?? this.db).prepare(`UPDATE message_drafts SET ${fields.join(', ')} WHERE message_id = ?`)
    stmt.run(...values)
  }

  // =============================================================================
  // Extra methods
  // =============================================================================

  /**
   * Upsert a draft - insert or replace existing.
   */
  async upsert(record: MessageDraftRecord, tx?: TransactionHandle): Promise<void> {
    // Validate enum columns (only status is a closed enum for message_drafts)
    assertEnumValues('message_drafts', {
      status: record.status,
    })

    const stmt = (tx?.db ?? this.db).prepare(`
      INSERT OR REPLACE INTO message_drafts (
        message_id, session_id, role, canonical_format,
        partial_content_json, status, created_at, updated_at, metadata_json
      ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
    `)

    stmt.run(
      record.message_id,
      record.session_id,
      record.role,
      record.canonical_format,
      record.partial_content_json,
      record.status,
      record.created_at,
      record.updated_at,
      record.metadata_json ?? null,
    )
  }

  /**
   * Delete draft for a specific message.
   */
  async delete_for_message(message_id: MessageID, tx?: TransactionHandle): Promise<void> {
    const stmt = (tx?.db ?? this.db).prepare('DELETE FROM message_drafts WHERE message_id = ?')
    stmt.run(message_id)
  }
}