Agent-System/web/src/components/chat/MessageStream.tsx

import React, { useRef, useLayoutEffect, useEffect, useCallback } from 'react';
import { ChatMessage, ChatDayBucket, MessageCategory } from '../../types/workbench';
import { MessageBubble } from './MessageBubble';
import { DayDivider } from '../history/DayDivider';
import { Loader2 } from 'lucide-react';
import { useVirtualMessages, flattenDayStack } from '../../hooks/useVirtualMessages';
import {
  registerVirtualScrollHandle,
  unregisterVirtualScrollHandle,
} from '../../engine/anchorNavigator';

export interface MessageStreamProps {
  messages?: ChatMessage[];
  dayStack?: ChatDayBucket[];
  isLoadingOlder?: boolean;
  hasMoreBefore?: boolean;
  onLoadOlder?: () => void;
  onLoadFullDay?: (dayKey: string) => void;
  onOpenCase?: (docId: string) => void;
  onCardAction?: (cardId: string, actionId: string, payload?: any) => void;
  onErrorAction?: (messageId: string, action: 'retry' | 'continue' | 'restart') => void;
  onQuote?: (message: ChatMessage) => void;
  onNavigateToAnchor?: (messageId: number) => void;
  selectedCategory?: MessageCategory | 'all';
  /** P3 / Step 5: 未读消息 ID 集合 */
  unreadMessageIds?: number[];
  onSaveToPrivate?: (message: ChatMessage) => void;
  onExportDocx?: (message: ChatMessage) => void;
  onExportPdf?: (message: ChatMessage) => void;
}

export const MessageStream: React.FC<MessageStreamProps> = ({
  messages,
  dayStack,
  isLoadingOlder = false,
  hasMoreBefore = false,
  onLoadOlder,
  onLoadFullDay,
  onOpenCase,
  onCardAction,
  onErrorAction,
  onQuote,
  onNavigateToAnchor,
  selectedCategory = 'all',
  unreadMessageIds = [],
  onSaveToPrivate,
  onExportDocx,
  onExportPdf,
}) => {
  const containerRef = useRef<HTMLDivElement>(null);
  const isInitialMountRef = useRef<boolean>(true);
  const prevMsgCountRef = useRef<number>(0);

  const unreadSet = React.useMemo(() => new Set(unreadMessageIds), [unreadMessageIds]);

  // 计算当前有效日分片集
  const activeBuckets: ChatDayBucket[] = React.useMemo(() => {
    if (dayStack && dayStack.length > 0) {
      return dayStack;
    }
    if (messages && messages.length > 0) {
      return [
        {
          dayKey: messages[0].dayKey || '',
          count: messages.length,
          loadedCount: messages.length,
          truncated: false,
          firstMessageId: messages[0].serverId || 0,
          lastMessageId: messages[messages.length - 1].serverId || 0,
          messages,
        },
      ];
    }
    return [];
  }, [dayStack, messages]);

  // 展平消息列表供虚拟化器使用
  const flatMessages = React.useMemo(
    () => flattenDayStack(activeBuckets),
    [activeBuckets]
  );

  // 按 serverId 快速映射 dayKey（供虚拟化行渲染时确定日期分割线）
  const dayKeyByMsgId = React.useMemo(() => {
    const map = new Map<number | string, string>();
    for (const bucket of activeBuckets) {
      for (const msg of bucket.messages) {
        if (msg.serverId) map.set(msg.serverId, bucket.dayKey);
        map.set(msg.id, bucket.dayKey);
      }
    }
    return map;
  }, [activeBuckets]);

  // 每个桶的第一条消息 id（用于判断是否需要在该消息前渲染日期分割线）
  const bucketFirstMsgId = React.useMemo(() => {
    const set = new Set<number | string>();
    for (const bucket of activeBuckets) {
      if (bucket.messages.length > 0) {
        const first = bucket.messages[0];
        if (first.serverId) set.add(first.serverId);
        set.add(first.id);
      }
    }
    return set;
  }, [activeBuckets]);

  // 每个桶的 dayKey -> bucket 映射（供日期分割线渲染）
  const bucketByDayKey = React.useMemo(() => {
    const map = new Map<string, ChatDayBucket>();
    for (const b of activeBuckets) map.set(b.dayKey, b);
    return map;
  }, [activeBuckets]);

  const totalMessagesCount = flatMessages.length;

  // 触顶回调
  const handleNearTop = useCallback(() => {
    if (!hasMoreBefore || isLoadingOlder) return;
    if (onLoadOlder) onLoadOlder();
  }, [hasMoreBefore, isLoadingOlder, onLoadOlder]);

  // 虚拟化 hook
  const { virtualizer, handle } = useVirtualMessages({
    flatMessages,
    containerRef,
    onNearTop: handleNearTop,
  });

  // 注册虚拟化句柄到全局 anchorNavigator
  useEffect(() => {
    registerVirtualScrollHandle(handle);
    return () => unregisterVirtualScrollHandle();
  }, [handle]);

  // 首屏挂载 / 切换会话时瞬时落底
  useLayoutEffect(() => {
    if (isInitialMountRef.current && totalMessagesCount > 0) {
      handle.scrollToBottom();
      isInitialMountRef.current = false;
      prevMsgCountRef.current = totalMessagesCount;
    }
  }, [totalMessagesCount, handle]);

  // 新消息到达时的滚动逻辑
  useEffect(() => {
    if (totalMessagesCount > prevMsgCountRef.current) {
      if (handle.isAtBottom()) {
        handle.scrollToBottom();
      }
      prevMsgCountRef.current = totalMessagesCount;
    }
  }, [totalMessagesCount, handle]);

  return (
    <div
      ref={containerRef}
      style={{
        flex: 1,
        overflowY: 'auto',
        position: 'relative',
      }}
    >
      {/* 触顶加载指示器 */}
      {isLoadingOlder && (
        <div
          style={{
            display: 'flex',
            alignItems: 'center',
            justifyContent: 'center',
            gap: '8px',
            padding: '8px 0',
            fontSize: '12px',
            color: 'var(--text-faint, #9ca3af)',
          }}
        >
          <Loader2 size={14} className="animate-spin" />
          <span>正在加载更早历史消息…</span>
        </div>
      )}

      {activeBuckets.length === 0 ? (
        <div
          style={{
            display: 'flex',
            flexDirection: 'column',
            alignItems: 'center',
            justifyContent: 'center',
            height: '100%',
            color: 'var(--text-faint, #9ca3af)',
            gap: '8px',
            fontSize: '13px',
            padding: '40px 0',
          }}
        >
          <span>暂无对话记录，请在下方发送立项需求</span>
        </div>
      ) : (
        /* 虚拟滚动轨道 */
        <div
          style={{
            height: typeof window === 'undefined' ? 'auto' : `${virtualizer.getTotalSize()}px`,
            width: '100%',
            position: 'relative',
            padding: '16px 0',
            boxSizing: 'border-box',
          }}
        >
          {(typeof window === 'undefined'
            ? flatMessages.map((msg, index) => ({
                index,
                key: msg.serverId ?? msg.id,
                start: 0,
              }))
            : virtualizer.getVirtualItems()
          ).map((vi) => {
            const msg = flatMessages[vi.index];
            const msgKey = msg.serverId ?? msg.id;
            const isFirst = bucketFirstMsgId.has(msgKey) || bucketFirstMsgId.has(msg.id);
            const dayKey = dayKeyByMsgId.get(msgKey) || dayKeyByMsgId.get(msg.id) || '';
            const bucket = bucketByDayKey.get(dayKey);

            return (
              <div
                key={vi.key}
                ref={typeof window === 'undefined' ? undefined : virtualizer.measureElement}
                data-index={vi.index}
                style={
                  typeof window === 'undefined'
                    ? { width: '100%' }
                    : {
                        position: 'absolute',
                        top: 0,
                        left: 0,
                        width: '100%',
                        transform: `translateY(${vi.start}px)`,
                      }
                }
              >
                {/* 日期分割线：该桶的第一条消息前展示 */}
                {isFirst && bucket && bucket.dayKey && (
                  <DayDivider bucket={bucket} onLoadFullDay={onLoadFullDay} />
                )}
                <MessageBubble
                  message={msg}
                  onOpenCase={onOpenCase}
                  onCardAction={onCardAction}
                  onErrorAction={onErrorAction}
                  onQuote={onQuote}
                  onNavigateToAnchor={onNavigateToAnchor}
                  isCategoryDimmed={
                    selectedCategory !== 'all' && msg.category !== selectedCategory
                  }
                  isUnread={msg.serverId ? unreadSet.has(msg.serverId) : false}
                  onSaveToPrivate={onSaveToPrivate}
                  onExportDocx={onExportDocx}
                  onExportPdf={onExportPdf}
                />
              </div>
            );
          })}
        </div>
      )}
    </div>
  );
};