import type { UiMessage } from "@/data/dto/chat";

export type ChatMessageKeyResolver = (message: UiMessage) => string;

export type ChatRenderItem =
  | { type: "date"; date: string; key: string }
  | { type: "msg"; message: UiMessage; key: string };

export function createChatMessageKeyResolver(): ChatMessageKeyResolver {
  const localMessageKeys = new WeakMap<UiMessage, string>();
  let nextLocalMessageKey = 0;

  return (message) => {
    if (message.id && message.id.length > 0) return `msg-${message.id}`;

    const existing = localMessageKeys.get(message);
    if (existing) return existing;

    nextLocalMessageKey += 1;
    const next = `local-msg-${nextLocalMessageKey}`;
    localMessageKeys.set(message, next);
    return next;
  };
}

export function buildChatRenderItems(
  messages: readonly UiMessage[],
  getMessageKey: ChatMessageKeyResolver,
): ChatRenderItem[] {
  const items: ChatRenderItem[] = [];

  messages.forEach((message, index) => {
    const messageKey = getMessageKey(message);
    if (index === 0 || message.date !== messages[index - 1].date) {
      items.push({
        type: "date",
        date: message.date,
        key: `date-${message.date}-${messageKey}`,
      });
    }
    items.push({ type: "msg", message, key: messageKey });
  });

  return items;
}