src/lib/message-segments.ts

import type { ChatMessage, MessageSegment } from './contracts';

function upsertContentSegment(
  segments: MessageSegment[],
  id: string,
  kind: 'text' | 'reasoning',
  content: string,
): MessageSegment[] {
  const index = segments.findIndex((segment) => segment.id === id);
  if (!content.trim()) {
    return index < 0 ? segments : segments.filter((_, candidate) => candidate !== index);
  }
  const next: MessageSegment = { id, kind, content };
  if (index < 0) return [...segments, next];
  return segments.map((segment, candidate) => candidate === index ? next : segment);
}

export function updateRoundSegments(
  segments: MessageSegment[],
  roundId: string,
  reasoning: string,
  text: string,
): MessageSegment[] {
  let next = upsertContentSegment(segments, `${roundId}:reasoning`, 'reasoning', reasoning);
  next = upsertContentSegment(next, `${roundId}:text`, 'text', text);
  return next;
}

export function appendToolSegment(
  segments: MessageSegment[],
  toolCallId: string,
  roundId?: string,
): MessageSegment[] {
  if (segments.some((segment) => segment.kind === 'tool' && segment.toolCallId === toolCallId)) {
    return segments;
  }
  return [...segments, {
    id: `tool:${toolCallId}`,
    kind: 'tool',
    toolCallId,
    ...(roundId ? { roundId } : {}),
  }];
}

export function segmentsForMessage(message: ChatMessage): MessageSegment[] {
  if (message.segments) return message.segments;
  const segments: MessageSegment[] = [];
  if (message.reasoning?.trim()) {
    segments.push({ id: `${message.id}:legacy-reasoning`, kind: 'reasoning', content: message.reasoning });
  }
  if (message.content.trim()) {
    segments.push({ id: `${message.id}:legacy-text`, kind: 'text', content: message.content });
  }
  const referencedArtifacts = new Set<string>();
  for (const tool of message.tools ?? []) {
    segments.push({ id: `tool:${tool.id}`, kind: 'tool', toolCallId: tool.id });
    if (tool.artifactId) referencedArtifacts.add(tool.artifactId);
  }
  for (const artifact of message.artifacts ?? []) {
    if (!referencedArtifacts.has(artifact.id)) {
      segments.push({ id: `artifact:${artifact.id}`, kind: 'artifact', artifactId: artifact.id });
    }
  }
  return segments;
}

export function normalizeMessageSegments(message: ChatMessage): ChatMessage {
  if (message.role !== 'assistant') return message;
  const artifactIds = new Set((message.artifacts ?? []).map((artifact) => artifact.id));
  const tools = message.tools?.map((tool) => !tool.artifactId && artifactIds.has(tool.id)
    ? { ...tool, artifactId: tool.id }
    : tool);
  const withTools = tools ? { ...message, tools } : message;
  return { ...withTools, segments: segmentsForMessage(withTools) };
}