import type { WorkNotice } from "@code/primitives/work";

import {
  getMessageText,
  getRawMessageText,
  getReasoningText,
} from "@/lib/chat/transforms";
import type { ConversationMessage } from "@/lib/chat/types";

import type { WorkspaceTimelineItem } from "./types";

export const isVisibleConversationMessage = (
  message: ConversationMessage
): boolean =>
  message.role === "user" ||
  getMessageText(message).length > 0 ||
  getReasoningText(message).length > 0;

const targetIndexForNotice = (
  messages: readonly ConversationMessage[],
  notice: WorkNotice
): number => {
  const sourceIndexes = notice.sourceTexts.flatMap((sourceText) => {
    const index = messages.findIndex(
      (message) =>
        message.role === "user" && getRawMessageText(message) === sourceText
    );
    return index === -1 ? [] : [index];
  });
  if (sourceIndexes.length === 0) {
    return messages.length - 1;
  }
  const sourceIndex = Math.max(...sourceIndexes);
  const responseOffset = messages
    .slice(sourceIndex + 1)
    .findIndex((message) => message.role === "assistant");
  return responseOffset === -1 ? sourceIndex : sourceIndex + responseOffset + 1;
};

export const buildWorkspaceTimeline = (
  allMessages: readonly ConversationMessage[],
  notices: readonly WorkNotice[]
): readonly WorkspaceTimelineItem[] => {
  const messages = allMessages.filter(isVisibleConversationMessage);
  const noticesByMessageIndex = new Map<number, WorkNotice[]>();
  for (const notice of notices) {
    const targetIndex = targetIndexForNotice(messages, notice);
    const atTarget = noticesByMessageIndex.get(targetIndex) ?? [];
    atTarget.push(notice);
    noticesByMessageIndex.set(targetIndex, atTarget);
  }

  const timeline: WorkspaceTimelineItem[] = [];
  for (const [index, message] of messages.entries()) {
    timeline.push({ kind: "message", message });
    for (const notice of noticesByMessageIndex.get(index) ?? []) {
      timeline.push({ kind: "work", notice });
    }
  }
  if (messages.length === 0) {
    for (const notice of noticesByMessageIndex.get(-1) ?? []) {
      timeline.push({ kind: "work", notice });
    }
  }
  return timeline;
};

export const findSourceMessageTarget = (
  messages: readonly ConversationMessage[],
  sourceText: string
): string | undefined =>
  messages.find(
    (message) =>
      message.role === "user" && getRawMessageText(message) === sourceText
  )?.id;
