src/lib/services/message-handler.ts

import type { WebSocketMessage } from "./websocket";
import { chatStore } from "../stores/chat-store";
import { contentManager } from "./content-manager";
import type { MessageSegment } from "../models/chat-data";

export class MessageHandler {
  private currentMessageId: string | null = null;
  private currentSegments: Map<string, MessageSegment> = new Map();
  private completedSegments: MessageSegment[] = [];
  private latestTodoSegmentId: string | null = null;
  private updateTimer: ReturnType<typeof setTimeout> | null = null;

  handleMessage(message: WebSocketMessage): void {
    switch (message.type) {
      case "status":
        this.handleStatus(message);
        break;
      case "stream_start":
        this.handleStreamStart(message);
        break;
      case "stream_token":
        this.handleStreamToken(message);
        break;
      case "stream_end":
        this.handleStreamEnd();
        break;
      case "chat":
        this.handleChat(message);
        break;
      case "error":
        this.handleError(message);
        break;
      case "editor_update":
        this.handleEditorUpdate(message);
        break;
      case "segment_start":
        this.handleSegmentStart(message);
        break;
      case "segment_token":
        this.handleSegmentToken(message);
        break;
      case "segment_end":
        this.handleSegmentEnd(message);
        break;
      case "tool_start":
      case "tool_end":
        break;
    }
  }

  private handleStatus(message: WebSocketMessage): void {
    const { processing, connected } = message.payload;
    if (processing !== undefined) {
      chatStore.setProcessing(processing as boolean);
    }
    if (connected !== undefined) {
      chatStore.setConnected(connected as boolean);
    }
  }

  private handleStreamStart(message: WebSocketMessage): void {
    const messageId =
      (message.payload.messageId as string) || `assistant_${Date.now()}`;
    this.currentMessageId = messageId;
    this.currentSegments.clear();
    this.completedSegments = [];
    this.latestTodoSegmentId = null;

    chatStore.addMessage({
      id: messageId,
      role: "assistant",
      content: "",
      timestamp: Date.now(),
      segments: [],
    });
  }

  private handleStreamToken(message: WebSocketMessage): void {
    const token = (message.payload.token as string) || "";
    if (this.currentMessageId && token) {
      chatStore.appendToLastMessage(token);
    }
  }

  private handleStreamEnd(): void {
    if (this.currentMessageId) {
      this.updateMessageSegments();
    }
    this.currentMessageId = null;
    this.currentSegments.clear();
    this.completedSegments = [];
  }

  private handleChat(message: WebSocketMessage): void {
    const { content } = message.payload;
    if (content) {
      const messageId = `assistant_${Date.now()}`;
      chatStore.addMessage({
        id: messageId,
        role: "assistant",
        content: content as string,
        timestamp: Date.now(),
      });
    }
  }

  private handleError(message: WebSocketMessage): void {
    chatStore.setError((message.payload.error as string) || null);
  }

  private handleEditorUpdate(message: WebSocketMessage): void {
    const content = message.payload.content as string;
    if (content) {
      contentManager.updateFromAgent(content);
    }
  }

  private handleSegmentStart(message: WebSocketMessage): void {
    const { segmentId, segmentType, toolName, toolArgs } = message.payload;
    if (!segmentId || !this.currentMessageId) return;

    const segment: MessageSegment = {
      id: segmentId as string,
      type: segmentType as MessageSegment["type"],
      content: "",
      toolName: toolName as string | undefined,
      toolArgs: toolArgs as Record<string, unknown> | undefined,
      startTime: Date.now(),
      streaming: segmentType === "text",
    };

    this.currentSegments.set(segmentId as string, segment);
    this.updateMessageSegments();
  }

  private handleSegmentToken(message: WebSocketMessage): void {
    const { segmentId, token } = message.payload;
    if (!segmentId || !token) return;

    const segment = this.currentSegments.get(segmentId as string);
    if (segment) {
      const updatedSegment = {
        ...segment,
        content: segment.content + (token as string),
      };
      this.currentSegments.set(segmentId as string, updatedSegment);

      // Throttle updates during streaming to avoid excessive re-renders
      if (this.updateTimer) {
        clearTimeout(this.updateTimer);
      }
      this.updateTimer = setTimeout(() => {
        this.updateMessageSegments();
        this.updateTimer = null;
      }, 50);
    }
  }

  private handleSegmentEnd(message: WebSocketMessage): void {
    const {
      segmentId,
      content,
      toolStatus,
      toolOutput,
      toolError,
      consoleOutput,
    } = message.payload;
    if (!segmentId) return;

    // Clear any pending update timer for immediate final update
    if (this.updateTimer) {
      clearTimeout(this.updateTimer);
      this.updateTimer = null;
    }

    const segment = this.currentSegments.get(segmentId as string);
    if (segment) {
      const completedSegment: MessageSegment = {
        ...segment,
        content: content ? (content as string) : segment.content,
        toolStatus: toolStatus
          ? (toolStatus as MessageSegment["toolStatus"])
          : segment.toolStatus,
        toolOutput: toolOutput ? (toolOutput as string) : segment.toolOutput,
        toolError: toolError ? (toolError as string) : segment.toolError,
        consoleOutput: consoleOutput
          ? (consoleOutput as string[])
          : segment.consoleOutput,
        endTime: Date.now(),
        streaming: false,
      };

      // Special handling for todo tools - merge with existing
      if (this.isTodoTool(completedSegment)) {
        this.handleTodoSegment(completedSegment);
      } else {
        this.completedSegments.push(completedSegment);
      }

      this.currentSegments.delete(segmentId as string);
      this.updateMessageSegments();
    }
  }

  private isTodoTool(segment: MessageSegment): boolean {
    return !!segment.toolName?.includes("task");
  }

  private handleTodoSegment(segment: MessageSegment): void {
    // Remove previous todo segment and replace with new one
    if (this.latestTodoSegmentId) {
      this.completedSegments = this.completedSegments.filter(
        (s) => s.id !== this.latestTodoSegmentId,
      );
    }

    this.latestTodoSegmentId = segment.id;
    this.completedSegments.push(segment);
  }

  private updateMessageSegments(): void {
    if (!this.currentMessageId) return;

    const allSegments = [
      ...this.completedSegments,
      ...Array.from(this.currentSegments.values()),
    ];

    const content = this.buildContentFromSegments(allSegments);

    chatStore.setLastMessageContent(content);
    chatStore.setLastMessageSegments(allSegments);
  }

  private buildContentFromSegments(segments: MessageSegment[]): string {
    let content = "";
    let lastWasText = false;

    for (const segment of segments) {
      if (segment.type === "text" && segment.content) {
        // Add spacing between text segments if needed
        if (lastWasText && content && !content.endsWith("\n")) {
          content += " ";
        }
        content += segment.content;
        lastWasText = true;
      } else {
        // Tool segments are handled in UI, but reset text flag
        lastWasText = false;
      }
    }

    return content.trim();
  }
}

export const messageHandler = new MessageHandler();