logancyang_obsidian-copilot/src/core/MessageRepository.ts

import { PromptContextEnvelope } from "@/context/PromptContextTypes";
import { formatDateTime } from "@/utils";
import { ChatMessage, MessageContext, NewChatMessage, StoredMessage } from "@/types/message";
import { logInfo } from "@/logger";

/**
 * MessageRepository - Single source of truth for all messages
 *
 * This implements a minimal clean architecture where:
 * - Each message is stored once with both display and processed text
 * - Display messages are computed views for UI
 * - LLM messages are computed views for AI communication
 * - No complex dual message systems or ID matching
 */
export class MessageRepository {
  private messages: StoredMessage[] = [];

  /**
   * Generate a unique message ID
   */
  private generateId(): string {
    return `msg-${Date.now()}-${Math.random().toString(36).slice(2, 11)}`;
  }

  /**
   * Add a message from a ChatMessage object
   */
  addMessage(message: NewChatMessage): string;
  /**
   * Add a message with separate display and processed text
   */
  addMessage(
    displayText: string,
    processedText: string,
    sender: string,
    context?: MessageContext,
    content?: any[]
  ): string;
  addMessage(
    messageOrDisplayText: NewChatMessage | string,
    processedText?: string,
    sender?: string,
    context?: MessageContext,
    content?: any[]
  ): string {
    // If first parameter is a ChatMessage object
    if (typeof messageOrDisplayText === "object") {
      const message = messageOrDisplayText;
      const id = message.id || this.generateId();
      const timestamp = message.timestamp || formatDateTime(new Date());

      const storedMessage: StoredMessage = {
        id,
        displayText: message.message,
        processedText: message.originalMessage || message.message,
        sender: message.sender,
        timestamp,
        context: message.context,
        contextEnvelope: message.contextEnvelope,
        isVisible: message.isVisible !== false,
        isErrorMessage: message.isErrorMessage,
        sources: message.sources,
        content: message.content,
        responseMetadata: message.responseMetadata,
      };

      this.messages.push(storedMessage);
      logInfo(`[MessageRepository] Added message with ID: ${id}`);
      return id;
    }

    // Otherwise, use string parameters
    if (processedText === undefined || sender === undefined) {
      throw new Error("processedText and sender are required when using string-based addMessage");
    }

    const displayText = messageOrDisplayText;
    const id = this.generateId();
    const timestamp = formatDateTime(new Date());

    const message: StoredMessage = {
      id,
      displayText,
      processedText,
      sender,
      timestamp,
      context,
      contextEnvelope: undefined,
      isVisible: true,
      isErrorMessage: false,
      content,
    };

    this.messages.push(message);
    logInfo(`[MessageRepository] Added message with ID: ${id}`);

    return id;
  }

  /**
   * Edit a message's display text
   */
  editMessage(id: string, newDisplayText: string): boolean {
    const message = this.messages.find((msg) => msg.id === id);
    if (!message) {
      logInfo(`[MessageRepository] Message not found for edit: ${id}`);
      return false;
    }

    if (message.displayText === newDisplayText) {
      logInfo(`[MessageRepository] No changes needed for message: ${id}`);
      return true;
    }

    // Update display text
    message.displayText = newDisplayText;

    // For user messages, mark that processed text needs updating
    if (message.sender === "user" || message.sender === "USER") {
      // ProcessedText will be updated by ContextManager
      logInfo(`[MessageRepository] Edited user message ${id}, needs context reprocessing`);
    } else {
      // For AI messages, display and processed are the same
      message.processedText = newDisplayText;
      logInfo(`[MessageRepository] Edited AI message ${id}`);
    }

    return true;
  }

  /**
   * Update the processed text for a message (after context processing)
   *
   * TRANSITIONAL METHOD - Updates both processedText (legacy) and contextEnvelope (new)
   * during Phase 1 migration. After ChainRunner migration (Phase 2), this can be
   * simplified to only update contextEnvelope.
   */
  updateProcessedText(
    id: string,
    processedText: string,
    contextEnvelope?: PromptContextEnvelope
  ): boolean {
    const message = this.messages.find((msg) => msg.id === id);
    if (!message) {
      logInfo(`[MessageRepository] Message not found for processed text update: ${id}`);
      return false;
    }

    // TRANSITIONAL: Update both for backward compatibility
    message.processedText = processedText;
    message.contextEnvelope = contextEnvelope;
    logInfo(`[MessageRepository] Updated processed text for message ${id}`);
    return true;
  }

  /**
   * Delete a message
   */
  deleteMessage(id: string): boolean {
    const index = this.messages.findIndex((msg) => msg.id === id);
    if (index === -1) {
      logInfo(`[MessageRepository] Message not found for deletion: ${id}`);
      return false;
    }

    this.messages.splice(index, 1);
    logInfo(`[MessageRepository] Deleted message ${id}`);
    return true;
  }

  /**
   * Clear all messages
   */
  clear(): void {
    this.messages = [];
    logInfo(`[MessageRepository] Cleared all messages`);
  }

  /**
   * Truncate messages after a specific index
   */
  truncateAfter(index: number): void {
    this.messages = this.messages.slice(0, index + 1);
    logInfo(`[MessageRepository] Truncated messages after index ${index}`);
  }

  /**
   * Truncate messages after a specific message ID
   */
  truncateAfterMessageId(messageId: string): void {
    const index = this.messages.findIndex((msg) => msg.id === messageId);
    if (index !== -1) {
      this.messages = this.messages.slice(0, index + 1);
      logInfo(`[MessageRepository] Truncated messages after message ${messageId}`);
    }
  }

  /**
   * Get display messages (computed view for UI)
   * Shows displayText for all visible messages
   */
  getDisplayMessages(): ChatMessage[] {
    return this.messages
      .filter((msg) => msg.isVisible)
      .map((msg) => ({
        id: msg.id,
        message: msg.displayText,
        originalMessage: msg.displayText,
        sender: msg.sender,
        timestamp: msg.timestamp,
        isVisible: true,
        context: msg.context,
        contextEnvelope: msg.contextEnvelope,
        isErrorMessage: msg.isErrorMessage,
        sources: msg.sources,
        content: msg.content,
        responseMetadata: msg.responseMetadata,
      }));
  }

  /**
   * Get a specific message for LLM processing with full context
   *
   * TRANSITIONAL METHOD - Returns processedText (concatenated context) for
   * legacy ChainRunners that haven't migrated to envelope-based prompts.
   *
   * MIGRATION NOTE:
   * - Phase 1: Use this for current turn processing in legacy runners
   * - Phase 2+: Prefer contextEnvelope with LayerToMessagesConverter
   *
   * Returns processedText (with context) for the message.
   */
  getLLMMessage(id: string): ChatMessage | undefined {
    const msg = this.messages.find((m) => m.id === id);
    if (!msg) return undefined;

    return {
      id: msg.id,
      message: msg.processedText, // TRANSITIONAL: Full context (legacy format)
      originalMessage: msg.displayText,
      sender: msg.sender,
      timestamp: msg.timestamp,
      isVisible: false, // LLM messages are not for display
      context: msg.context,
      contextEnvelope: msg.contextEnvelope, // NEW: Use this for envelope-based runners
      isErrorMessage: msg.isErrorMessage,
      sources: msg.sources,
      content: msg.content,
      responseMetadata: msg.responseMetadata,
    };
  }

  /**
   * Get all messages for LLM conversation history
   * IMPORTANT: Returns displayText only (raw messages without context)
   * to prevent context duplication in chat memory.
   *
   * Context should be added per-turn via the envelope (L3 layer),
   * not baked into the chat history.
   */
  getLLMMessages(): ChatMessage[] {
    return this.messages.map((msg) => ({
      id: msg.id,
      message: msg.displayText, // Changed from processedText to prevent context duplication
      originalMessage: msg.displayText,
      sender: msg.sender,
      timestamp: msg.timestamp,
      isVisible: false,
      context: msg.context,
      contextEnvelope: msg.contextEnvelope,
      isErrorMessage: msg.isErrorMessage,
      sources: msg.sources,
      content: msg.content,
    }));
  }

  /**
   * Get a message by ID (returns display version)
   */
  getMessage(id: string): ChatMessage | undefined {
    const msg = this.messages.find((m) => m.id === id);
    if (!msg) return undefined;

    return {
      id: msg.id,
      message: msg.displayText,
      originalMessage: msg.displayText,
      sender: msg.sender,
      timestamp: msg.timestamp,
      isVisible: msg.isVisible,
      context: msg.context,
      contextEnvelope: msg.contextEnvelope,
      isErrorMessage: msg.isErrorMessage,
      sources: msg.sources,
      content: msg.content,
    };
  }

  /**
   * Load messages from persistence
   */
  loadMessages(messages: ChatMessage[]): void {
    this.clear();
    messages.forEach((msg) => {
      this.messages.push({
        id: msg.id || this.generateId(),
        displayText: msg.message,
        processedText: msg.originalMessage || msg.message,
        sender: msg.sender,
        timestamp: msg.timestamp || formatDateTime(new Date()),
        context: msg.context,
        contextEnvelope: msg.contextEnvelope,
        isVisible: msg.isVisible !== false,
        isErrorMessage: msg.isErrorMessage,
        sources: msg.sources,
        content: msg.content,
      });
    });
    logInfo(`[MessageRepository] Loaded ${messages.length} messages`);
  }

  /**
   * Get debug information
   */
  getDebugInfo() {
    return {
      totalMessages: this.messages.length,
      visibleMessages: this.messages.filter((m) => m.isVisible).length,
      userMessages: this.messages.filter((m) => m.sender === "user" || m.sender === "USER").length,
      aiMessages: this.messages.filter((m) => m.sender === "AI" || m.sender === "assistant").length,
    };
  }
}