src/lib/services/segment-formatter.ts

import type { MessageSegment } from "../models/chat-data";
import type { TodoListView } from "../models/segment-view";
import { parseTodoList } from "../models/segment-view";

export class SegmentFormatter {
  private static todoListCache: Map<string, TodoListView> = new Map();

  static formatSegmentContent(segment: MessageSegment): string {
    switch (segment.type) {
      case "text":
        return segment.content;

      case "reasoning":
        return segment.content;

      case "tool-invocation":
        return this.formatToolInvocation(segment);

      case "tool-result":
        return this.formatToolResult(segment);

      default:
        return segment.content;
    }
  }

  static formatToolInvocation(segment: MessageSegment): string {
    const args = segment.toolArgs
      ? JSON.stringify(segment.toolArgs, null, 2)
      : "No arguments";

    return `Tool: ${segment.toolName}\nArguments:\n${args}`;
  }

  static formatToolResult(segment: MessageSegment): string {
    if (segment.toolError) {
      return `❌ Error: ${segment.toolError}`;
    }

    if (segment.toolName?.includes("task")) {
      return this.formatTodoResult(segment);
    }

    if (segment.toolName === "observe_console") {
      return this.formatConsoleOutput(segment);
    }

    return segment.toolOutput || segment.content || "No output";
  }

  static formatTodoResult(segment: MessageSegment): string {
    const content = segment.toolOutput || segment.content;
    const todoList = parseTodoList(content);

    if (todoList) {
      this.todoListCache.set(segment.id, todoList);
      return this.renderTodoList(todoList);
    }

    return content;
  }

  static formatConsoleOutput(segment: MessageSegment): string {
    const output = segment.toolOutput || segment.content;
    const lines = output.split("\n");

    const formatted = lines
      .map((line) => {
        if (line.includes("[error]")) {
          return `🔴 ${line}`;
        } else if (line.includes("[warn]")) {
          return `🟡 ${line}`;
        } else if (line.includes("[info]")) {
          return `🔵 ${line}`;
        } else if (line.includes("[debug]")) {
          return `⚪ ${line}`;
        }
        return line;
      })
      .join("\n");

    return formatted;
  }

  static renderTodoList(todoList: TodoListView): string {
    const header = `📋 Tasks (${todoList.completedCount}/${todoList.totalCount} completed)\n`;
    const separator = "─".repeat(40) + "\n";

    const tasks = todoList.tasks
      .map((task) => `${task.emoji} [${task.id}] ${task.description}`)
      .join("\n");

    return header + separator + tasks;
  }

  static getLatestTodoList(): TodoListView | null {
    if (this.todoListCache.size === 0) {
      return null;
    }

    let latest: TodoListView | null = null;
    let latestTime = 0;

    for (const todoList of this.todoListCache.values()) {
      if (todoList.lastUpdated > latestTime) {
        latest = todoList;
        latestTime = todoList.lastUpdated;
      }
    }

    return latest;
  }

  static shouldCollapseByDefault(segment: MessageSegment): boolean {
    if (segment.type !== "tool-invocation" && segment.type !== "tool-result") {
      return false;
    }

    if (segment.toolError) {
      return false;
    }

    if (segment.toolName?.includes("task")) {
      return false;
    }

    const output = segment.toolOutput || segment.content || "";
    const lineCount = output.split("\n").length;

    return lineCount > 10;
  }

  static getSegmentIcon(segment: MessageSegment): string {
    const iconMap: Record<string, string> = {
      text: "💬",
      reasoning: "🤔",
      "tool-invocation": "🔧",
      "tool-result": "📊",
    };

    if (segment.toolName) {
      const toolIcons: Record<string, string> = {
        plan_tasks: "📋",
        update_task: "✏️",
        view_tasks: "👀",
        observe_console: "📺",
      };
      return toolIcons[segment.toolName] || iconMap[segment.type] || "📄";
    }

    return iconMap[segment.type] || "📄";
  }

  static formatDuration(ms: number): string {
    if (ms < 1000) {
      return `${ms}ms`;
    } else if (ms < 60000) {
      return `${(ms / 1000).toFixed(1)}s`;
    } else {
      const minutes = Math.floor(ms / 60000);
      const seconds = Math.floor((ms % 60000) / 1000);
      return `${minutes}m ${seconds}s`;
    }
  }

  static truncateContent(content: string, maxLength: number = 100): string {
    if (content.length <= maxLength) {
      return content;
    }

    const truncated = content.substring(0, maxLength);
    const lastSpace = truncated.lastIndexOf(" ");

    if (lastSpace > maxLength * 0.8) {
      return truncated.substring(0, lastSpace) + "...";
    }

    return truncated + "...";
  }
}

export const segmentFormatter = new SegmentFormatter();