obsidian_ollama/src/chat-view.ts

import { ItemView, WorkspaceLeaf, Notice } from 'obsidian';
/// <reference lib="dom" />
// Use global types from JSDOM setup
type KeyboardEvent = globalThis.KeyboardEvent;
type HTMLTextAreaElement = globalThis.HTMLTextAreaElement;
type HTMLButtonElement = globalThis.HTMLButtonElement;

const DEFAULT_VAULT_SEARCH_LIMIT = 3;
const MAX_MESSAGE_HISTORY = 50;
const MAX_STREAM_CHUNKS = 1000;
import {
  PluginSettings,
  OllamaMessage,
  ChatMessage,
  OllamaTool,
  ToolCall,
  ToolResult,
} from './types';
import { OllamaClient } from './ollama-client';
import { VaultIndexer } from './vault-indexer';
import { ToolExecutor } from './tool-executor';
import { ErrorHandler } from './error-handler';

export class ChatView extends ItemView {
  private settings: PluginSettings;
  private messages: ChatMessage[] = [];
  private ollamaClient: OllamaClient;
  private vaultIndexer: VaultIndexer;
  private toolExecutor: ToolExecutor;
  private lastMessageEl: HTMLElement | null = null;
  private newChatButton: HTMLButtonElement | null = null;
  private sendButton: HTMLButtonElement | null = null;
  private inputEl: HTMLTextAreaElement | null = null;
  private chatContainer: HTMLElement | null = null;
  private sendButtonClickHandler: (() => Promise<void>) | null = null;
  private inputKeyDownHandler: ((e: KeyboardEvent) => Promise<void>) | null = null;
  private newChatButtonClickHandler: (() => void) | null = null;
  private sendButtonClickWrapper: (() => void) | null = null;
  private inputKeyDownWrapper: ((e: KeyboardEvent) => void) | null = null;
  private newChatButtonClickWrapper: (() => void) | null = null;
  private listenersAttached = false;

  // Getters for testing
  public getSendButtonClickHandler(): (() => Promise<void>) | null {
    return this.sendButtonClickHandler;
  }

  public getInputKeyDownHandler(): ((e: KeyboardEvent) => Promise<void>) | null {
    return this.inputKeyDownHandler;
  }

  public getNewChatButtonClickHandler(): (() => void) | null {
    return this.newChatButtonClickHandler;
  }

  constructor(leaf: WorkspaceLeaf, settings: PluginSettings) {
    super(leaf);
    this.settings = settings;
    this.ollamaClient = new OllamaClient(settings.ollamaUrl, settings.model);
    this.vaultIndexer = new VaultIndexer(this.app.vault);
    this.toolExecutor = new ToolExecutor(this.app.vault, this.app);
  }

  public updateSettings(newSettings: PluginSettings): void {
    this.settings = newSettings;
    this.ollamaClient = new OllamaClient(newSettings.ollamaUrl, newSettings.model);
  }

  getViewType(): string {
    return 'ollama-chat-view';
  }

  getDisplayText(): string {
    return 'Ollama Chat';
  }

  onOpen(): Promise<void> {
    this.render();
    this.removeEventListeners(); // Clean up any existing listeners before reattaching
    this.setupEventListeners();
    return Promise.resolve();
  }

  public onSettingsChange(newSettings: PluginSettings): void {
    this.updateSettings(newSettings);
  }

  onClose(): Promise<void> {
    this.ollamaClient.cancelStream();
    this.removeEventListeners();
    this.cleanupStreamingResources();
    this.lastMessageEl = null;
    this.sendButton = null;
    this.inputEl = null;
    this.chatContainer = null;
    return Promise.resolve();
  }

  private cleanupStreamingResources(): void {
    // Only cleanup if there's still an active streaming message
    const streamingMessage = this.messages.find((msg) => msg.isStreaming);
    if (streamingMessage && this.lastMessageEl && this.lastMessageEl.parentElement) {
      this.lastMessageEl.parentElement.removeChild(this.lastMessageEl);
      this.lastMessageEl = null;
    }
  }

  render() {
    const container =
      this.chatContainer || this.contentEl.createEl('div', { cls: 'ollama-chat-container' });
    this.chatContainer = container;
    const inputContainer =
      this.contentEl.querySelector('.ollama-input-container') ||
      this.contentEl.createEl('div', { cls: 'ollama-input-container' });

    if (!this.inputEl) {
      this.inputEl = inputContainer.createEl('textarea', { cls: 'ollama-input' });
    }
    if (!this.sendButton) {
      this.sendButton = inputContainer.createEl('button', {
        cls: 'ollama-send-button',
      });
      this.sendButton.textContent = 'Send';
    }

    if (!this.newChatButton) {
      const newChatContainer =
        this.contentEl.querySelector('.ollama-new-chat') ||
        this.contentEl.createEl('div', { cls: 'ollama-new-chat' });
      this.newChatButton = newChatContainer.createEl('button', {
        cls: 'ollama-new-chat-button',
      });
      this.newChatButton.textContent = '🔄 New Chat';
      this.newChatButton.title = 'Start a new conversation';
    }

    // Create immutable snapshot for rendering
    const messagesSnapshot = [...this.messages];

    // Only render messages that are not currently streaming
    const nonStreamingMessages = messagesSnapshot.filter((msg) => !msg.isStreaming);

    // Differential update: only update messages that have changed
    const existingMessages = container.querySelectorAll('.ollama-message');

    for (const msg of nonStreamingMessages) {
      const existingEl = container.querySelector(`.ollama-message[data-msg-id="${msg.id}"]`);
      if (existingEl) {
        existingEl.textContent = msg.content;
      } else {
        const messageEl = container.createEl('div', {
          cls: `ollama-message ${msg.role}`,
        }) as HTMLElement;
        messageEl.setAttribute('data-msg-id', msg.id);
        messageEl.textContent = msg.content;
      }
    }

    // Remove messages that are no longer in the array
    for (const el of Array.from(existingMessages)) {
      const id = el.getAttribute('data-msg-id');
      if (!id || !nonStreamingMessages.some((m) => m.id === id)) {
        el.remove();
      }
    }

    // Re-attach streaming message if it exists
    const streamingMessage = messagesSnapshot.find((msg) => msg.isStreaming);
    if (streamingMessage && this.lastMessageEl) {
      const existingStreamingEl = container.querySelector(
        `.ollama-message[data-msg-id="${streamingMessage.id}"]`
      );
      if (!existingStreamingEl) {
        container.appendChild(this.lastMessageEl);
      }
    }
  }

  private setupEventListeners(): void {
    if (!this.sendButton || !this.inputEl || this.listenersAttached) return;

    // Create handlers if they don't exist
    if (!this.sendButtonClickHandler) {
      this.sendButtonClickHandler = async () => {
        if (!this.inputEl) return;
        await this.handleUserInput(this.inputEl.value);
        this.inputEl.value = '';
      };
    }

    if (!this.inputKeyDownHandler) {
      this.inputKeyDownHandler = async (e: KeyboardEvent) => {
        if (!this.inputEl || e.key !== 'Enter' || e.shiftKey) return;
        e.preventDefault();
        await this.handleUserInput(this.inputEl.value);
        this.inputEl.value = '';
      };
    }

    // Create wrapper functions for event listeners
    this.sendButtonClickWrapper = () => {
      void this.sendButtonClickHandler?.();
    };
    this.inputKeyDownWrapper = (e: KeyboardEvent) => {
      void this.inputKeyDownHandler?.(e);
    };
    this.newChatButtonClickWrapper = () => {
      void this.newChatButtonClickHandler?.();
    };

    // Add event listeners using wrappers
    this.sendButton.addEventListener('click', this.sendButtonClickWrapper);
    this.inputEl.addEventListener('keydown', this.inputKeyDownWrapper);
    if (this.newChatButton) {
      if (!this.newChatButtonClickHandler) {
        this.newChatButtonClickHandler = () => this.clearConversation();
      }
      this.newChatButton.addEventListener('click', this.newChatButtonClickWrapper);
    }
    this.listenersAttached = true;
  }

  private removeEventListeners(): void {
    if (this.sendButton && this.sendButtonClickWrapper) {
      this.sendButton.removeEventListener('click', this.sendButtonClickWrapper);
    }
    if (this.inputEl && this.inputKeyDownWrapper) {
      this.inputEl.removeEventListener('keydown', this.inputKeyDownWrapper);
    }
    if (this.newChatButton && this.newChatButtonClickWrapper) {
      this.newChatButton.removeEventListener('click', this.newChatButtonClickWrapper);
    }
    this.sendButtonClickWrapper = null;
    this.inputKeyDownWrapper = null;
    this.newChatButtonClickWrapper = null;
    this.listenersAttached = false;
  }

  private clearConversation(): void {
    // Create new array to ensure immutability
    this.messages = [];
    this.lastMessageEl = null;
    this.render();
    new Notice('Conversation cleared');
  }

  private updateMessageById(id: string, partial: Partial<ChatMessage>): boolean {
    const index = this.messages.findIndex((m) => m.id === id);
    if (index < 0) return false;
    this.messages = [
      ...this.messages.slice(0, index),
      { ...this.messages[index], ...partial },
      ...this.messages.slice(index + 1),
    ];
    return true;
  }

  private updateLastMessage(content: string) {
    const streamingMessage = this.messages.find((msg) => msg.isStreaming);
    if (streamingMessage && !this.lastMessageEl) {
      this.lastMessageEl = this.contentEl.createEl('div', {
        cls: `ollama-message assistant`,
      }) as HTMLElement;
      this.lastMessageEl.setAttribute('data-msg-id', streamingMessage.id);
    }
    if (this.lastMessageEl) {
      this.lastMessageEl.textContent = content;
    }
  }

  private async handleUserInput(content: string) {
    if (!this.sendButton || !this.inputEl) return;
    this.sendButton.disabled = true;

    try {
      // Guard against empty messages
      const userMessage = content.trim();
      if (!userMessage) return;

      // Search vault using user message as query
      const entries = await this.vaultIndexer.searchVault(userMessage, DEFAULT_VAULT_SEARCH_LIMIT);
      let context = entries.map((entry) => `### ${entry.title}\n${entry.content}`).join('\n\n');

      // Cap context size to prevent prompt bloat with large vaults
      const MAX_CONTEXT_LENGTH = 4000;
      if (context.length > MAX_CONTEXT_LENGTH) {
        context = context.substring(0, MAX_CONTEXT_LENGTH) + '\n\n... (truncated)';
      }

      const systemContent = context
        ? `You are a helpful assistant.\n\nRelevant vault context:\n${context}`
        : 'You are a helpful assistant.';
      const systemMessage: OllamaMessage = {
        role: 'system',
        content: systemContent,
      };
      const userMessageWithContext: OllamaMessage = {
        role: 'user',
        content: userMessage,
      };

      const messages: OllamaMessage[] = [
        systemMessage,
        ...this.messages.map(
          (m) =>
            ({
              role: m.role,
              content: m.content,
              tool_calls: m.tool_calls,
            }) as OllamaMessage
        ),
        userMessageWithContext,
      ];

      const tools: OllamaTool[] = [
        {
          type: 'function',
          function: {
            name: 'create_file',
            description: 'Create a new file in the vault',
            parameters: {
              type: 'object' as const,
              properties: {
                path: {
                  type: 'string' as const,
                  description: "Relative path within the vault, e.g. 'Notes/todo.md'",
                },
                content: { type: 'string' as const, description: 'Content of the file to create' },
              },
              required: ['path', 'content'],
            },
          },
        },
      ];

      const messageId = `${Date.now()}-${Math.random().toString(36).slice(2, 8)}`;

      const userMessageId = messageId;
      const assistantMessageId = `${messageId}-assistant`;

      // Store user message in conversation history
      const userChatMessage: ChatMessage = {
        id: userMessageId,
        role: 'user' as const,
        content: userMessage,
        timestamp: Date.now(),
      };

      const assistantMessage: ChatMessage = {
        id: assistantMessageId,
        role: 'assistant' as const,
        content: '',
        timestamp: Date.now(),
        isStreaming: true,
      };

      // Update messages immutably
      this.messages = [...this.messages, userChatMessage, assistantMessage];

      try {
        this.render();

        const stream = this.ollamaClient.streamChat(messages, tools);
        let fullResponse = '';
        let toolCalls: ToolCall[] = [];
        let chunkCount = 0;
        for await (const chunk of stream) {
          chunkCount++;
          if (chunkCount > MAX_STREAM_CHUNKS) {
            throw new Error('Response too long, stopped streaming');
          }

          if (chunk.content) {
            fullResponse += chunk.content;
          }

          if (chunk.tool_calls) {
            toolCalls = toolCalls.concat(chunk.tool_calls);
          }

          this.updateLastMessage(fullResponse);
        }

        // Update the assistant message with the full response immutably
        this.updateMessageById(assistantMessageId, {
          content: fullResponse,
          tool_calls: toolCalls,
        });

        // Process tool calls with proper follow-up context
        if (toolCalls.length > 0) {
          // Validate tool calls before processing
          const MAX_TOOL_CALLS = 10;
          if (toolCalls.length > MAX_TOOL_CALLS) {
            throw new Error(`Too many tool calls (max ${MAX_TOOL_CALLS})`);
          }

          // Collect all tool results using allSettled to support partial results
          const settledResults = await Promise.allSettled(
            toolCalls.map((call) => this.toolExecutor.handleToolCall(call))
          );

          const toolResults: ToolResult[] = [];
          for (const result of settledResults) {
            if (result.status === 'fulfilled') {
              toolResults.push(result.value);
            } else {
              // Use centralized error handler for tool errors
              ErrorHandler.handleError(result.reason, 'ChatView.handleUserInput');
            }
          }

          // Only create follow-up when we have tool results
          if (toolResults.length > 0) {
            // Create follow-up messages including the assistant's tool calls and results
            const followUpMessages: OllamaMessage[] = [
              ...messages,
              { role: 'assistant' as const, content: fullResponse, tool_calls: toolCalls },
              ...toolResults.map((result) => ({
                role: 'tool' as const,
                content: JSON.stringify(result),
              })),
            ];

            const followUp = await this.ollamaClient.chat(followUpMessages, tools);
            fullResponse += followUp.content;
            this.updateLastMessage(fullResponse);

            // Update the assistant message with the final response immutably
            this.updateMessageById(assistantMessageId, {
              content: fullResponse,
              isStreaming: false,
            });
          } else {
            // Even if no tool results were successful, mark streaming as complete
            // to prevent the assistant message from disappearing
            this.updateMessageById(assistantMessageId, {
              content: fullResponse,
              isStreaming: false,
            });
          }
        }

        // Update assistant message immutably — only if no tool calls were processed
        if (toolCalls.length === 0) {
          this.updateMessageById(assistantMessageId, {
            isStreaming: false,
          });
        }

        // Limit conversation history to prevent memory issues
        if (this.messages.length > MAX_MESSAGE_HISTORY) {
          this.messages = this.messages.slice(-MAX_MESSAGE_HISTORY);
        }
        this.render();
      } finally {
        // Clean up streaming resources regardless of outcome
        this.cleanupStreamingResources();
      }
    } catch (error) {
      // Use centralized error handler
      ErrorHandler.handleError(error, 'ChatView.handleUserInput');
      // Update any streaming messages to non-streaming state to prevent stale messages
      // This ensures that if an error occurs during streaming, the assistant message
      // is still visible (with any partial content received) but won't cause issues
      // in subsequent requests due to stale isStreaming: true flag
      this.messages = this.messages.map((msg) =>
        msg.isStreaming ? { ...msg, isStreaming: false } : msg
      );
      this.cleanupStreamingResources();
    } finally {
      if (this.sendButton) {
        this.sendButton.disabled = false;
      }
    }
  }
}