course-agent-od/web/src/store/chat.ts

import { create } from "zustand";
import { ChatSession, ChatMessage, SourceInfo, ThinkingStep } from "@/types";
import { chatAPI } from "@/lib/api";

interface ChatState {
  sessions: ChatSession[];
  currentSession: ChatSession | null;
  messages: ChatMessage[];
  isLoading: boolean;
  isStreaming: boolean;
  error: string | null;
  abortController: AbortController | null;
}

interface ChatActions {
  // 会话管理
  loadSessions: () => Promise<void>;
  createSession: (title: string) => Promise<ChatSession | null>;
  selectSession: (sessionId: number) => Promise<void>;
  deleteSession: (sessionId: number) => Promise<void>;
  renameSession: (sessionId: number, newTitle: string) => Promise<void>;
  exportSession: (sessionId: number, format?: string) => Promise<void>;
  
  // 消息管理
  loadMessages: (sessionId: number) => Promise<void>;
  sendMessage: (message: string) => Promise<void>;
  streamMessage: (message: string, mode?: string, knowledgeBaseIds?: string[], model?: string) => Promise<void>;
  addMessage: (message: ChatMessage) => void;
  updateMessage: (messageId: number, content: string) => void;
  editMessage: (messageId: number, newContent: string) => Promise<void>;
  regenerateMessage: (messageId: number, model?: string) => Promise<void>;
  feedbackMessage: (messageId: number, feedback: "like" | "dislike") => Promise<void>;
  stopGeneration: () => void;
  
  // 状态管理
  setLoading: (loading: boolean) => void;
  setStreaming: (streaming: boolean) => void;
  setError: (error: string | null) => void;
  clearError: () => void;
  clearMessages: () => void;
}

type ChatStore = ChatState & ChatActions;

export const useChatStore = create<ChatStore>((set, get) => ({
  // 初始状态
  sessions: [],
  currentSession: null,
  messages: [],
  isLoading: false,
  isStreaming: false,
  error: null,
  abortController: null,

  // 加载会话列表
  loadSessions: async () => {
    try {
      const sessions = await chatAPI.getSessions();
      set({ sessions });
    } catch (error) {
      const errorMessage = error instanceof Error ? error.message : "加载会话失败";
      set({ error: errorMessage });
    }
  },

  // 创建新会话
  createSession: async (title: string) => {
    try {
      // 直接调用创建会话API
      const response = await chatAPI.createSession(title);
      const newSession: ChatSession = {
        id: response.id,
        title: response.title,
        created_at: response.created_at,
        updated_at: response.updated_at || response.created_at,
        message_count: response.message_count || 0,
      };
      
      set((state) => ({
        sessions: [newSession, ...state.sessions],
        currentSession: newSession,
        messages: [],
      }));
      
      return newSession;
    } catch (error) {
      console.error("创建会话失败:", error);
      return null;
    }
  },

  // 选择会话
  selectSession: async (sessionId: number) => {
    set({ isLoading: true, error: null });
    
    try {
      const session = get().sessions.find(s => s.id === sessionId);
      if (!session) {
        throw new Error("会话不存在");
      }
      
      set({ currentSession: session });
      await get().loadMessages(sessionId);
    } catch (error) {
      const errorMessage = error instanceof Error ? error.message : "选择会话失败";
      set({ error: errorMessage, isLoading: false });
    }
  },

  // 删除会话
  deleteSession: async (sessionId: number) => {
    try {
      await chatAPI.deleteSession(sessionId);
      set((state) => ({
        sessions: state.sessions.filter(s => s.id !== sessionId),
        currentSession: state.currentSession?.id === sessionId ? null : state.currentSession,
        messages: state.currentSession?.id === sessionId ? [] : state.messages,
      }));
    } catch (error) {
      console.error("删除会话失败:", error);
      const errorMessage = error instanceof Error ? error.message : "删除会话失败";
      set({ error: errorMessage });
    }
  },

  // 重命名会话
  renameSession: async (sessionId: number, newTitle: string) => {
    try {
      await chatAPI.renameSession(sessionId, newTitle);
      set((state) => ({
        sessions: state.sessions.map(s =>
          s.id === sessionId ? { ...s, title: newTitle } : s
        ),
        currentSession: state.currentSession?.id === sessionId
          ? { ...state.currentSession, title: newTitle }
          : state.currentSession,
      }));
    } catch (error) {
      console.error("重命名会话失败:", error);
      const errorMessage = error instanceof Error ? error.message : "重命名会话失败";
      set({ error: errorMessage });
    }
  },

  // 导出会话
  exportSession: async (sessionId: number, format: string = "json") => {
    try {
      const exportData = await chatAPI.exportSession(sessionId, format);
      
      // 创建下载链接
      const blob = new Blob([JSON.stringify(exportData, null, 2)], {
        type: "application/json",
      });
      const url = URL.createObjectURL(blob);
      const a = document.createElement("a");
      a.href = url;
      a.download = `${exportData.title || `session-${sessionId}`}.${format}`;
      document.body.appendChild(a);
      a.click();
      document.body.removeChild(a);
      URL.revokeObjectURL(url);
    } catch (error) {
      console.error("导出会话失败:", error);
      const errorMessage = error instanceof Error ? error.message : "导出会话失败";
      set({ error: errorMessage });
    }
  },

  // 加载消息
  loadMessages: async (sessionId: number) => {
    try {
      const messages = await chatAPI.getSessionMessages(sessionId);
      set({ messages, isLoading: false });
    } catch (error) {
      const errorMessage = error instanceof Error ? error.message : "加载消息失败";
      set({ error: errorMessage, isLoading: false });
    }
  },

  // 发送消息（非流式）
  sendMessage: async (message: string) => {
    const { currentSession } = get();
    if (!currentSession) return;

    set({ isLoading: true, error: null });

    // 添加用户消息
    const userMessage: ChatMessage = {
      id: Date.now(),
      role: "user",
      content: message,
      created_at: new Date().toISOString(),
    };
    
    set((state) => ({
      messages: [...state.messages, userMessage],
    }));

    try {
      const response = await chatAPI.sendMessage(message, currentSession.id);
      
      // 添加助手回复
      const assistantMessage: ChatMessage = {
        id: response.message_id,
        role: "assistant",
        content: response.answer,
        created_at: new Date().toISOString(),
        metadata: { sources: response.sources },
      };
      
      set((state) => ({
        messages: [...state.messages, assistantMessage],
        isLoading: false,
      }));
    } catch (error) {
      const errorMessage = error instanceof Error ? error.message : "发送消息失败";
      set({ error: errorMessage, isLoading: false });
    }
  },

  // 流式发送消息
  streamMessage: async (message: string, mode: string = "normal", knowledgeBaseIds?: string[], model?: string) => {
    const { currentSession } = get();
    if (!currentSession) return;

    // 创建新的 AbortController
    const abortController = new AbortController();

    set({ isStreaming: true, error: null, abortController });

    // 添加用户消息
    const userMessage: ChatMessage = {
      id: Date.now(),
      role: "user",
      content: message,
      created_at: new Date().toISOString(),
    };
    
    set((state) => ({
      messages: [...state.messages, userMessage],
    }));

    // 创建助手消息占位符
    const assistantMessage: ChatMessage = {
      id: Date.now() + 1,
      role: "assistant",
      content: "",
      created_at: new Date().toISOString(),
    };
    
    set((state) => ({
      messages: [...state.messages, assistantMessage],
    }));

    let chunkCount = 0;
    let totalChars = 0;
    let thinkingSteps: ThinkingStep[] = [];

    try {
      await chatAPI.streamMessage(
        message,
        currentSession.id,
        mode, // 传递实际选择的模式
        knowledgeBaseIds,
        abortController.signal, // 新增参数
        (chunk: string) => {
          // 解析chunk
          try {
            const data = JSON.parse(chunk);
            
            if (data.type === 'thinking') {
              // 处理思考过程
              thinkingSteps.push({
                stage: data.stage,
                message: data.message,
                doc_count: data.doc_count,
                time: data.time
              });
              
              // 立即更新思考过程到UI
              requestAnimationFrame(() => {
                set((state) => ({
                  messages: state.messages.map(msg =>
                    msg.id === assistantMessage.id
                      ? { ...msg, thinking: [...thinkingSteps] }
                      : msg
                  ),
                }));
              });
            } else if (data.type === 'chunk') {
              // 处理内容chunk
              chunkCount++;
              totalChars += data.content.length;

              requestAnimationFrame(() => {
                set((state) => ({
                  messages: state.messages.map(msg =>
                    msg.id === assistantMessage.id
                      ? { ...msg, content: msg.content + data.content }
                      : msg
                  ),
                }));
              });
            }
          } catch (e) {
            // 向后兼容：纯文本chunk
            chunkCount++;
            totalChars += chunk.length;

            requestAnimationFrame(() => {
              set((state) => ({
                messages: state.messages.map(msg =>
                  msg.id === assistantMessage.id
                    ? { ...msg, content: msg.content + chunk }
                    : msg
                ),
              }));
            });
          }
        },
        (_sessionId: number, messageId?: number, userMessageId?: number) => {
          if (messageId) {
            set((state) => ({
              messages: state.messages.map(msg => {
                if (msg.id === assistantMessage.id) return { ...msg, id: messageId };
                if (userMessageId && msg.id === userMessage.id) return { ...msg, id: userMessageId };
                return msg;
              }),
              isStreaming: false,
              abortController: null,
            }));
          } else {
            // 前端已通过流式获得完整内容，不需要从服务端重新加载
            set({ isStreaming: false, abortController: null });
          }
        },
        (error: string) => {
          // 流式错误
          set({ error, isStreaming: false, abortController: null });
        },
        undefined, // onStatus
        undefined, // onThinking
        undefined, // onThinkingContent
        model // 传递模型ID
      );
    } catch (error) {
      const errorMessage = error instanceof Error ? error.message : "发送消息失败";
      set({ error: errorMessage, isStreaming: false, abortController: null });
    }
  },

  // 添加消息
  addMessage: (message: ChatMessage) => {
    set((state) => ({
      messages: [...state.messages, message],
    }));
  },

  // 更新消息
  updateMessage: (messageId: number, content: string) => {
    set((state) => ({
      messages: state.messages.map(msg =>
        msg.id === messageId ? { ...msg, content } : msg
      ),
    }));
  },

  // 设置加载状态
  setLoading: (loading: boolean) => {
    set({ isLoading: loading });
  },

  // 设置流式状态
  setStreaming: (streaming: boolean) => {
    set({ isStreaming: streaming });
  },

  // 设置错误
  setError: (error: string | null) => {
    set({ error });
  },

  // 清除错误
  clearError: () => {
    set({ error: null });
  },

  // 清除消息
  clearMessages: () => {
    set({ messages: [] });
  },

  // 编辑消息
  editMessage: async (messageId: number, newContent: string) => {
    try {
      await chatAPI.editMessage(messageId, newContent);
      set((state) => ({
        messages: state.messages.map(msg =>
          msg.id === messageId ? { ...msg, content: newContent, edited: true } : msg
        ),
      }));
    } catch (error) {
      console.error("编辑消息失败:", error);
      const errorMessage = error instanceof Error ? error.message : "编辑消息失败";
      set({ error: errorMessage });
    }
  },

  // 重新生成消息（流式）
  regenerateMessage: async (messageId: number, model?: string) => {
    const { currentSession, messages } = get();
    if (!currentSession) return;

    const messageIndex = messages.findIndex(msg => msg.id === messageId);
    if (messageIndex === -1) return;

    // 找到该 assistant 消息之前的 user 消息内容
    const userContent = messages
      .slice(0, messageIndex)
      .reverse()
      .find(msg => msg.role === "user")?.content;
    if (!userContent) return;

    // 删除该 assistant 消息及之后的所有消息，替换为流式占位符
    const trimmed = messages.slice(0, messageIndex);
    const assistantPlaceholder: ChatMessage = {
      id: Date.now() + 1,
      role: "assistant",
      content: "",
      created_at: new Date().toISOString(),
    };

    const abortController = new AbortController();
    set({
      messages: [...trimmed, assistantPlaceholder],
      isStreaming: true,
      error: null,
      abortController,
    });

    // 判断模式：检查之前的消息中是否有知识库相关内容
    const mode = "normal";

    let thinkingSteps: any[] = [];

    try {
      await chatAPI.streamMessage(
        userContent,
        currentSession.id,
        mode,
        undefined, // knowledgeBaseIds
        abortController.signal,
        (chunk: string) => {
          // 处理 chunk
          try {
            const data = JSON.parse(chunk);
            if (data.type === 'thinking') {
              thinkingSteps.push({
                stage: data.stage,
                message: data.message,
                doc_count: data.doc_count,
                time: data.time
              });
              requestAnimationFrame(() => {
                set((state) => ({
                  messages: state.messages.map(msg =>
                    msg.id === assistantPlaceholder.id
                      ? { ...msg, thinking: [...thinkingSteps] }
                      : msg
                  ),
                }));
              });
            } else if (data.type === 'chunk') {
              requestAnimationFrame(() => {
                set((state) => ({
                  messages: state.messages.map(msg =>
                    msg.id === assistantPlaceholder.id
                      ? { ...msg, content: msg.content + data.content }
                      : msg
                  ),
                }));
              });
            }
          } catch {
            // 纯文本 chunk
            requestAnimationFrame(() => {
              set((state) => ({
                messages: state.messages.map(msg =>
                  msg.id === assistantPlaceholder.id
                    ? { ...msg, content: msg.content + chunk }
                    : msg
                ),
              }));
            });
          }
        },
        (_sessionId: number, messageId?: number) => {
          // onComplete: 只替换 assistant 占位符 ID，不重新加载
          if (messageId) {
            set((state) => ({
              messages: state.messages.map(msg =>
                msg.id === assistantPlaceholder.id
                  ? { ...msg, id: messageId }
                  : msg
              ),
              isStreaming: false,
              abortController: null,
            }));
          } else {
            set({ isStreaming: false, abortController: null });
          }
        },
        (error: string) => {
          set({ error, isStreaming: false, abortController: null });
        },
        undefined, // onStatus
        undefined, // onThinking
        undefined, // onThinkingContent
        model
      );
    } catch (error) {
      const errorMessage = error instanceof Error ? error.message : "重新生成失败";
      set({ error: errorMessage, isStreaming: false, abortController: null });
    }
  },

  // 反馈消息
  feedbackMessage: async (messageId: number, feedback: "like" | "dislike") => {
    try {
      await chatAPI.feedbackMessage(messageId, feedback);
      set((state) => ({
        messages: state.messages.map(msg =>
          msg.id === messageId ? { ...msg, feedback } : msg
        ),
      }));
    } catch (error) {
      console.error("反馈消息失败:", error);
      const errorMessage = error instanceof Error ? error.message : "反馈消息失败";
      set({ error: errorMessage });
    }
  },

  // 停止生成
  stopGeneration: () => {
    const { abortController } = get();
    
    // 如果存在 AbortController，调用 abort 中止请求
    if (abortController) {
      abortController.abort();
    }

    // 清理状态
    set({ isStreaming: false, abortController: null });
  },
}));