Agent-System/web/src/hooks/useChatWindow.ts

import { useState, useEffect, useCallback, useRef } from 'react';
import { ChatMessage, ChatDayBucket, MessageCategory } from '../types/workbench';
import {
  fetchChatWindow,
  fetchDayMessages,
} from '../api/chatWindow';
import {
  loadSessionWindowCache,
  saveSessionWindowCache,
  clearSessionWindowCache,
  loadSessionDraft,
  saveSessionDraft,
} from '../data/storage';
import { enrichMessagesWithCategories } from '../engine/categoryHelper';
import {
  filterPendingLiveMessages,
  isLiveMessageInHistory,
  mergeLiveOverHistory,
} from '../engine/messageIdentity';

export type WindowPhase =
  | 'loading_window'
  | 'ready'
  | 'loading_older'
  | 'authorized_empty'
  | 'forbidden'
  | 'error';

export interface UseChatWindowOptions {
  sessionId: string;
  userId?: number;
  days?: number;
  initialFallbackMessages?: ChatMessage[];
  liveMessages?: ChatMessage[];
}

export interface UseChatWindowReturn {
  phase: WindowPhase;
  dayStack: ChatDayBucket[];
  displayDayStack: ChatDayBucket[];
  allMessages: ChatMessage[];
  hasMoreBefore: boolean;
  loadOlder: () => Promise<void>;
  loadFullDay: (dayKey: string) => Promise<void>;
  jumpToDate: (dayKey: string) => Promise<void>;
  loadDayForAnchor: (messageId: number) => Promise<boolean>;
  selectedCategory: MessageCategory | 'all';
  setSelectedCategory: (cat: MessageCategory | 'all') => void;
  appendLiveMessage: (msg: ChatMessage) => void;
  updateLiveMessage: (id: string, patch: Partial<ChatMessage>) => void;
  replaceMessageServerId: (tempId: string, serverId: number) => void;
  draft: string;
  setDraft: (draft: string) => void;
  refresh: () => Promise<void>;
}

const MAX_WINDOW_DAYS_STACK = 31;
const MAX_WINDOW_MSGS_STACK = 3000;

function getTodayKey(): string {
  const now = new Date();
  return `${now.getFullYear()}-${String(now.getMonth() + 1).padStart(2, '0')}-${String(now.getDate()).padStart(2, '0')}`;
}

function bucketHasMessages(days: ChatDayBucket[]): boolean {
  return days.some((d) => (d.messages?.length ?? 0) > 0);
}

function isForbiddenResponse(err: unknown): boolean {
  if (!(err instanceof Error)) return false;
  const msg = err.message.toLowerCase();
  return msg.includes('403') || msg.includes('forbidden');
}

export function useChatWindow({
  sessionId,
  userId,
  days = 7,
  initialFallbackMessages,
  liveMessages,
}: UseChatWindowOptions): UseChatWindowReturn {
  const [phase, setPhase] = useState<WindowPhase>('loading_window');
  const [dayStack, setDayStack] = useState<ChatDayBucket[]>(() => {
    const cached = loadSessionWindowCache(sessionId, userId);
    if (cached && cached.length > 0 && bucketHasMessages(cached)) return cached;
    return [];
  });
  const [hasMoreBefore, setHasMoreBefore] = useState<boolean>(true);
  const [oldestLoadedDay, setOldestLoadedDay] = useState<string>('');

  const [selectedCategory, setSelectedCategory] = useState<MessageCategory | 'all'>('all');
  const [draft, setDraftState] = useState<string>(() => loadSessionDraft(sessionId));
  const generationRef = useRef(0);
  const abortRef = useRef<AbortController | null>(null);

  // session 切换时立即恢复该 session 的窗口缓存骨架，避免展示上一 session 内容
  useEffect(() => {
    abortRef.current?.abort();
    generationRef.current += 1;
    const cached = loadSessionWindowCache(sessionId, userId);
    const hasCachedMessages = cached && cached.length > 0 && bucketHasMessages(cached);
    setDayStack(hasCachedMessages ? cached! : []);
    setHasMoreBefore(true);
    setOldestLoadedDay('');
    setPhase(hasCachedMessages ? 'ready' : 'loading_window');
    setDraftState(loadSessionDraft(sessionId));
  }, [sessionId, userId]);

  const setDraft = useCallback(
    (newDraft: string) => {
      setDraftState(newDraft);
      saveSessionDraft(sessionId, newDraft);
    },
    [sessionId]
  );


  // 初次拉取窗口
  const loadInitialWindow = useCallback(async () => {
    if (!sessionId) return;
    const generation = ++generationRef.current;
    abortRef.current?.abort();
    const controller = new AbortController();
    abortRef.current = controller;

    setPhase('loading_window');

    try {
      const res = await fetchChatWindow(sessionId, days, undefined, controller.signal);
      if (controller.signal.aborted || generation !== generationRef.current) return;
      if (res.sessionId !== sessionId) return;
      if (bucketHasMessages(res.days)) {
        setDayStack(res.days);
        setHasMoreBefore(res.hasMoreBefore);
        setOldestLoadedDay(res.oldestLoadedDay);
        saveSessionWindowCache(sessionId, res.days, userId);
        setPhase('ready');
      } else if (res.days.length > 0) {
        clearSessionWindowCache(sessionId, userId);
        setDayStack([]);
        setHasMoreBefore(res.hasMoreBefore);
        setOldestLoadedDay(res.oldestLoadedDay);
        setPhase('authorized_empty');
      } else if (initialFallbackMessages && initialFallbackMessages.length > 0) {
        // 若服务端尚无历史数据（如纯本地 Demo 会话），将 fallback 消息包装为今日 bucket (Spec §3.5 A5)
        const enrichedFallback = enrichMessagesWithCategories(initialFallbackMessages);
        const todayStr = getTodayKey();
        const fallbackBucket: ChatDayBucket = {
          dayKey: todayStr,
          count: enrichedFallback.length,
          loadedCount: enrichedFallback.length,
          truncated: false,
          firstMessageId: 0,
          lastMessageId: 0,
          messages: enrichedFallback,
        };
        setDayStack([fallbackBucket]);
        setHasMoreBefore(false);
        setPhase('ready');
      } else {
        setDayStack([]);
        setHasMoreBefore(false);
        setPhase('authorized_empty');
      }
    } catch (err) {
      if (controller.signal.aborted || generation !== generationRef.current) return;
      if (isForbiddenResponse(err)) {
        clearSessionWindowCache(sessionId, userId);
        setDayStack([]);
        setPhase('forbidden');
        return;
      }
      console.warn('[useChatWindow] fetch window failed, falling back to local cache/fallback', err);
      const cached = loadSessionWindowCache(sessionId, userId);
      if (cached && bucketHasMessages(cached)) {
        setDayStack(cached);
        setPhase('ready');
        return;
      }
      if (!cached && initialFallbackMessages && initialFallbackMessages.length > 0) {
        const enrichedFallback = enrichMessagesWithCategories(initialFallbackMessages);
        const todayStr = getTodayKey();
        setDayStack([
          {
            dayKey: todayStr,
            count: enrichedFallback.length,
            loadedCount: enrichedFallback.length,
            truncated: false,
            firstMessageId: 0,
            lastMessageId: 0,
            messages: enrichedFallback,
          },
        ]);
      }
      setPhase('error');
    }
  }, [sessionId, userId, days, initialFallbackMessages]);

  useEffect(() => {
    loadInitialWindow();
    return () => {
      abortRef.current?.abort();
    };
  }, [loadInitialWindow]);

  // 触顶加载更早日 (P3: 滚动位置由虚拟化器 scrollToIndex 恢复，无需 DOM 补偿)
  const loadOlder = useCallback(
    async () => {
      if (!hasMoreBefore || phase === 'loading_older' || !oldestLoadedDay || !sessionId) {
        return;
      }

      setPhase('loading_older');

      try {
        const res = await fetchChatWindow(sessionId, 7, oldestLoadedDay);
        if (res.days.length > 0) {
          setDayStack((prev) => {
            // 合并并去重
            const existingKeys = new Set(prev.map((b) => b.dayKey));
            const newBuckets = res.days.filter((b) => !existingKeys.has(b.dayKey));
            const combined = [...newBuckets, ...prev].sort((a, b) => (a.dayKey < b.dayKey ? -1 : 1));

            // 容量保护：若超出上限，从最早日出栈
            let protectedStack = combined;
            if (protectedStack.length > MAX_WINDOW_DAYS_STACK) {
              protectedStack = protectedStack.slice(protectedStack.length - MAX_WINDOW_DAYS_STACK);
            }
            let totalMsgs = protectedStack.reduce((acc, b) => acc + b.loadedCount, 0);
            while (totalMsgs > MAX_WINDOW_MSGS_STACK && protectedStack.length > 1) {
              const removed = protectedStack.shift();
              if (removed) totalMsgs -= removed.loadedCount;
            }

            saveSessionWindowCache(sessionId, protectedStack, userId);
            return protectedStack;
          });

          setHasMoreBefore(res.hasMoreBefore);
          setOldestLoadedDay(res.oldestLoadedDay);
        } else {
          setHasMoreBefore(false);
        }
      } catch (err) {
        console.error('[useChatWindow] loadOlder failed', err);
      } finally {
        setPhase('ready');
      }
    },
    [hasMoreBefore, phase, oldestLoadedDay, sessionId, userId]
  );

  // 加载当日全量消息
  const loadFullDay = useCallback(
    async (dayKey: string) => {
      if (!sessionId || !dayKey) return;
      try {
        const res = await fetchDayMessages(sessionId, dayKey, { limit: 300 });
        setDayStack((prev) => {
          const updated = prev.map((b) => {
            if (b.dayKey === dayKey) {
              return {
                ...b,
                messages: res.messages,
                loadedCount: res.messages.length,
                truncated: false,
              };
            }
            return b;
          });
          saveSessionWindowCache(sessionId, updated, userId);
          return updated;
        });
      } catch (err) {
        console.error('[useChatWindow] loadFullDay failed', err);
      }
    },
    [sessionId, userId]
  );

  // 日期跳转: 聚焦指定日期
  const jumpToDate = useCallback(
    async (dayKey: string) => {
      const exists = dayStack.some((b) => b.dayKey === dayKey);
      if (exists) {
        if (typeof document !== 'undefined') {
          const el = document.getElementById(`chat-day-divider-${dayKey}`);
          if (el) {
            el.scrollIntoView({ block: 'start', behavior: 'auto' });
          }
        }
        return;
      }

      try {
        const res = await fetchDayMessages(sessionId, dayKey, { limit: 100 });
        if (res.messages.length > 0) {
          const bucket: ChatDayBucket = {
            dayKey: res.dayKey,
            count: res.totalCount,
            loadedCount: res.loadedCount,
            truncated: res.totalCount > res.loadedCount,
            firstMessageId: res.messages[0]?.serverId || 0,
            lastMessageId: res.messages[res.messages.length - 1]?.serverId || 0,
            messages: res.messages,
          };

          setDayStack((prev) => {
            const next = prev.filter((b) => b.dayKey !== dayKey);
            next.push(bucket);
            next.sort((a, b) => (a.dayKey < b.dayKey ? -1 : 1));
            if (next.length > MAX_WINDOW_DAYS_STACK) {
              next.splice(0, next.length - MAX_WINDOW_DAYS_STACK);
            }
            saveSessionWindowCache(sessionId, next, userId);
            return next;
          });

          setTimeout(() => {
            if (typeof document !== 'undefined') {
              const el = document.getElementById(`chat-day-divider-${dayKey}`);
              if (el) {
                el.scrollIntoView({ block: 'start', behavior: 'auto' });
              }
            }
          }, 60);
        }
      } catch (err) {
        console.warn(`[useChatWindow] jumpToDate failed for ${dayKey}:`, err);
      }
    },
    [sessionId, userId, dayStack]
  );

  // 跨日锚点拉取: 根据 messageId 拉取所属日分片并合并进窗口
  const loadDayForAnchor = useCallback(
    async (messageId: number): Promise<boolean> => {
      if (!sessionId) return false;
      for (const b of dayStack) {
        if (b.messages.some((m) => m.serverId === messageId)) {
          return true;
        }
      }

      try {
        const res = await fetchDayMessages(sessionId, '', { anchorId: messageId, limit: 100 });
        if (res.messages.length > 0) {
          const bucket: ChatDayBucket = {
            dayKey: res.dayKey,
            count: res.totalCount,
            loadedCount: res.loadedCount,
            truncated: res.totalCount > res.loadedCount,
            firstMessageId: res.messages[0]?.serverId || 0,
            lastMessageId: res.messages[res.messages.length - 1]?.serverId || 0,
            messages: res.messages,
          };

          setDayStack((prev) => {
            const next = prev.filter((b) => b.dayKey !== res.dayKey);
            next.push(bucket);
            next.sort((a, b) => (a.dayKey < b.dayKey ? -1 : 1));
            if (next.length > MAX_WINDOW_DAYS_STACK) {
              next.splice(0, next.length - MAX_WINDOW_DAYS_STACK);
            }
            saveSessionWindowCache(sessionId, next, userId);
            return next;
          });
          return true;
        }
      } catch (err) {
        console.warn(`[useChatWindow] loadDayForAnchor failed for message ${messageId}:`, err);
      }
      return false;
    },
    [sessionId, userId, dayStack]
  );

  // 追加新消息 (SSE 或用户发送)
  const appendLiveMessage = useCallback(
    (msg: ChatMessage) => {
      const todayStr = msg.dayKey || getTodayKey();
      const enrichedMsg: ChatMessage = {
        ...msg,
        dayKey: todayStr,
      };

      setDayStack((prev) => {
        let found = false;
        const updated = prev.map((bucket) => {
          if (bucket.dayKey === todayStr) {
            found = true;
            return {
              ...bucket,
              count: bucket.count + 1,
              loadedCount: bucket.loadedCount + 1,
              messages: [...bucket.messages, enrichedMsg],
            };
          }
          return bucket;
        });

        if (!found) {
          const newBucket: ChatDayBucket = {
            dayKey: todayStr,
            count: 1,
            loadedCount: 1,
            truncated: false,
            firstMessageId: msg.serverId || 0,
            lastMessageId: msg.serverId || 0,
            messages: [enrichedMsg],
          };
          updated.push(newBucket);
          updated.sort((a, b) => (a.dayKey < b.dayKey ? -1 : 1));
        }

        saveSessionWindowCache(sessionId, updated, userId);
        return updated;
      });
    },
    [sessionId, userId]
  );

  // 原地更新指定消息
  const updateLiveMessage = useCallback(
    (id: string, patch: Partial<ChatMessage>) => {
      setDayStack((prev) => {
        let changed = false;
        const updated = prev.map((bucket) => {
          const mIdx = bucket.messages.findIndex((m) => m.id === id);
          if (mIdx >= 0) {
            changed = true;
            const copyMsgs = [...bucket.messages];
            copyMsgs[mIdx] = { ...copyMsgs[mIdx], ...patch };
            return { ...bucket, messages: copyMsgs };
          }
          return bucket;
        });
        if (changed) {
          saveSessionWindowCache(sessionId, updated, userId);
          return updated;
        }
        return prev;
      });
    },
    [sessionId, userId]
  );

  // 将流式临时 ID 替换为服务端落库真实 serverId
  const replaceMessageServerId = useCallback(
    (tempId: string, serverId: number) => {
      setDayStack((prev) => {
        let changed = false;
        const updated = prev.map((bucket) => {
          const mIdx = bucket.messages.findIndex((m) => m.id === tempId);
          if (mIdx >= 0) {
            changed = true;
            const copyMsgs = [...bucket.messages];
            copyMsgs[mIdx] = {
              ...copyMsgs[mIdx],
              serverId,
            };
            return {
              ...bucket,
              lastMessageId: serverId,
              messages: copyMsgs,
            };
          }
          return bucket;
        });
        if (changed) {
          saveSessionWindowCache(sessionId, updated, userId);
          return updated;
        }
        return prev;
      });
    },
    [sessionId, userId]
  );

  // 将 liveMessages (来自组件层实时状态，如刚发送的 user 消息或 streaming 中的 assistant 气泡)
  // 无缝合并进对应的日 bucket 中供视图渲染
  const displayDayStack = useCallback(() => {
    if (!liveMessages || liveMessages.length === 0) {
      return dayStack;
    }

    const todayStr = getTodayKey();
    const historyMessages = dayStack.flatMap((b) => b.messages);
    const pendingLive = filterPendingLiveMessages(liveMessages, historyMessages);

    if (dayStack.length === 0) {
      return [
        {
          dayKey: todayStr,
          count: liveMessages.length,
          loadedCount: liveMessages.length,
          truncated: false,
          firstMessageId: liveMessages[0].serverId || 0,
          lastMessageId: liveMessages[liveMessages.length - 1].serverId || 0,
          messages: liveMessages,
        },
      ];
    }

    if (pendingLive.length === 0) {
      let hasUpdate = false;
      const updatedBuckets = dayStack.map((b) => {
        let bucketChanged = false;
        const newMsgs = b.messages.map((m) => {
          const liveM = liveMessages.find((live) => isLiveMessageInHistory(live, [m]));
          if (
            liveM &&
            (liveM.content !== m.content ||
              liveM.status !== m.status ||
              liveM.livePhaseLabel !== m.livePhaseLabel ||
              liveM.errorInfo !== m.errorInfo ||
              liveM.serverId !== m.serverId)
          ) {
            bucketChanged = true;
            hasUpdate = true;
            return mergeLiveOverHistory(m, liveM);
          }
          return m;
        });
        if (bucketChanged) {
          return { ...b, messages: newMsgs };
        }
        return b;
      });
      return hasUpdate ? updatedBuckets : dayStack;
    }

    const result = dayStack.map((b) => ({ ...b, messages: [...b.messages] }));
    let todayBucket = result.find((b) => b.dayKey === todayStr);
    if (!todayBucket) {
      todayBucket = {
        dayKey: todayStr,
        count: 0,
        loadedCount: 0,
        truncated: false,
        firstMessageId: 0,
        lastMessageId: 0,
        messages: [],
      };
      result.push(todayBucket);
      result.sort((a, b) => (a.dayKey < b.dayKey ? -1 : 1));
    }

    // 建立 liveMessages 的原始索引映射，确保合并时绝对保序
    const liveIndexMap = new Map<string, number>();
    liveMessages.forEach((m, idx) => {
      liveIndexMap.set(m.id, idx);
      if (m.serverId) liveIndexMap.set(`msg-${m.serverId}`, idx);
    });

    for (const msg of pendingLive) {
      todayBucket.messages.push(msg);
      todayBucket.count++;
      todayBucket.loadedCount++;
    }

    // 严谨保序：按照真实时间与发生时序排列，保证新追加的对话永远沉淀在最底部
    const getMessageSortTime = (m: ChatMessage): number => {
      if (m.createdAt) {
        const t = new Date(m.createdAt).getTime();
        if (!isNaN(t)) return t;
      }
      if (m.timestamp && m.timestamp.includes(':')) {
        const [hh, mm] = m.timestamp.split(':').map(Number);
        if (!isNaN(hh) && !isNaN(mm)) {
          const d = new Date();
          d.setHours(hh, mm, 0, 0);
          return d.getTime();
        }
      }
      if (m.id && (m.id.startsWith('msg-') || m.id.startsWith('ast-'))) {
        const numPart = m.id.replace(/^[a-z]+-/, '');
        const parsed = parseInt(numPart, 10);
        if (!isNaN(parsed) && parsed > 1000000000000) {
          return parsed;
        }
      }
      return 0;
    };

    todayBucket.messages.sort((a, b) => {
      const tA = getMessageSortTime(a);
      const tB = getMessageSortTime(b);
      if (tA !== 0 && tB !== 0 && tA !== tB) {
        return tA - tB;
      }
      const idxA = liveIndexMap.get(a.id) ?? (a.serverId ? liveIndexMap.get(`msg-${a.serverId}`) : undefined);
      const idxB = liveIndexMap.get(b.id) ?? (b.serverId ? liveIndexMap.get(`msg-${b.serverId}`) : undefined);
      if (idxA !== undefined && idxB !== undefined) {
        return idxA - idxB;
      }
      if (idxA !== undefined) return 1;
      if (idxB !== undefined) return -1;
      return (a.serverId || 0) - (b.serverId || 0);
    });

    return result;
  }, [dayStack, liveMessages])();

  // 扁平展开供现有下游（卡片处理、DST等）快速访问
  const allMessages = displayDayStack.flatMap((b) => b.messages);

  return {
    phase,
    dayStack,
    displayDayStack,
    allMessages,
    hasMoreBefore,
    loadOlder,
    loadFullDay,
    jumpToDate,
    loadDayForAnchor,
    selectedCategory,
    setSelectedCategory,
    appendLiveMessage,
    updateLiveMessage,
    replaceMessageServerId,
    draft,
    setDraft,
    refresh: loadInitialWindow,
  };
}