mky-vent-base/src/views/ventAI/manageAssistent/hooks/useChatApproval.ts

import { computed, ref } from 'vue';
import { message } from 'ant-design-vue';
import { chatResumeStream, StreamError } from '/@/views/ventAI/manageAssistent/api';
import type { Message, SseEvent, Task } from '/@/views/ventAI/manageAssistent/types/index';

/**
 * 待审批（ask_user 中断）流程：按任务登记审批信息、提交答案续跑、取消提问。
 *
 * 审批状态按任务 id 存放；提交/取消都要走一次 resume 流，因此需要注入 SSE 事件处理、
 * AbortController 管理与滚动能力，保持与发送流程一致的行为。
 */
export interface ChatApprovalDeps {
  getTaskList: () => Task[];
  getCurrentTaskId: () => string;
  switchTask: (taskId: string) => void | Promise<void>;
  setTaskAbortController: (taskId: string, controller: AbortController | null) => void;
  handleStreamEvent: (data: SseEvent, aiMsg: Message, taskId: string) => void;
  scrollToBottom: (force?: boolean) => Promise<void>;
  scrollToBottomThrottled: (force?: boolean) => void;
  refreshUsage: () => void;
  resetTodos: () => void;
  setCurrentSessionId: (sessionId: string) => void;
  /** 补齐消息主键：续跑完成后的回答同样依赖它才有删除入口，同时把这一轮被中断的前半段挂到消息上 */
  hydrateMessageIds: (sessionId: string, msgs: Message[]) => Promise<void>;
}

export const useChatApproval = (deps: ChatApprovalDeps) => {
  const {
    getTaskList,
    getCurrentTaskId,
    switchTask,
    setTaskAbortController,
    handleStreamEvent,
    scrollToBottom,
    scrollToBottomThrottled,
    refreshUsage,
    resetTodos,
    setCurrentSessionId,
    hydrateMessageIds,
  } = deps;
  // 按任务维度的状态存储
  const taskPendingApprovals = ref<Map<string, { sessionId: string; threadId: string }>>(new Map());

  // 获取当前任务的待审批状态
  const currentPendingApproval = computed(() => {
    return taskPendingApprovals.value.get(getCurrentTaskId());
  });

  // 设置任务的审批状态
  const setTaskPendingApproval = (taskId: string, sessionId: string, threadId: string) => {
    const map = new Map(taskPendingApprovals.value);
    map.set(taskId, { sessionId, threadId });
    taskPendingApprovals.value = map;
  };

  // 清除任务的审批状态
  const clearTaskPendingApproval = (taskId: string) => {
    const map = new Map(taskPendingApprovals.value);
    map.delete(taskId);
    taskPendingApprovals.value = map;
  };

  // 任务 ID 迁移（流式期间临时 id → 后端 session_id）时同步审批状态的 key，否则审批会跟丢任务
  const migrateTaskApproval = (fromTaskId: string, toTaskId: string) => {
    const entry = taskPendingApprovals.value.get(fromTaskId);
    if (!entry) return;
    const map = new Map(taskPendingApprovals.value);
    map.delete(fromTaskId);
    map.set(toTaskId, entry);
    taskPendingApprovals.value = map;
  };

  const findPendingApprovalTask = (): { taskId: string; approval: { sessionId: string; threadId?: string } } | null => {
    // 优先用当前任务 ID 查找
    const approval = currentPendingApproval.value;
    if (approval) return { taskId: getCurrentTaskId(), approval };
    // 遍历所有审批（可能任务 ID 已变，用 sessionId 做回退匹配）
    for (const [taskId, a] of taskPendingApprovals.value.entries()) {
      return { taskId, approval: a };
    }
    return null;
  };

  const handleAskUserSubmit = async (answers: string[]) => {
    // 防重复提交：如果当前任务正在流式中，忽略
    const currentTask = getTaskList().find((t) => t.id === getCurrentTaskId());
    if (currentTask?.isStreaming) return;

    const found = findPendingApprovalTask();

    // 优先从 taskPendingApprovals 获取，若被清除则从消息的 pendingApprovalData 回退
    let sessionId = found?.approval.sessionId;
    let threadId = found?.approval.threadId;
    const foundTaskId = found?.taskId;

    // 查找任务：优先用 foundTaskId，回退到 currentTaskId
    const task = (foundTaskId ? getTaskList().find((t) => t.id === foundTaskId) : null) || getTaskList().find((t) => t.id === getCurrentTaskId());
    if (!task) return;

    // 切换到目标任务（如果不在当前视图）
    if (task.id !== getCurrentTaskId()) {
      await switchTask(task.id);
    }

    // 单一数据源：消息只存在于任务自己的 messages 数组里（当前任务的 messages.value 就是它）
    const sourceArr = task.messages;
    const lastMsg = sourceArr[sourceArr.length - 1];
    if (!lastMsg || !lastMsg.isPendingApproval) return;

    // 若 taskPendingApprovals 已被清除，从消息的 pendingApprovalData 回退取值
    if (!sessionId && lastMsg.pendingApprovalData) {
      sessionId = lastMsg.pendingApprovalData.sessionId;
      threadId = lastMsg.pendingApprovalData.threadId;
    }
    if (!sessionId) return;

    lastMsg.isPendingApproval = false;
    lastMsg.pendingApprovalData = undefined;
    lastMsg.isLoading = true;
    lastMsg.generateStartTime = Date.now();

    task.isStreaming = true;
    resetTodos();

    const controller = new AbortController();
    setTaskAbortController(task.id, controller);

    try {
      const result = await chatResumeStream(
        {
          session_id: sessionId,
          thread_id: threadId || undefined,
          status: 'answered',
          answers,
          signal: controller.signal,
        },
        (data: SseEvent) => {
          handleStreamEvent(data, lastMsg, task.id);
          if (getCurrentTaskId() === task.id) scrollToBottomThrottled();
        }
      );

      if (result.session_id) {
        setCurrentSessionId(result.session_id);
      }

      // 续跑完成，这一轮的回答刚由后端持久化：补主键（删除按钮依赖它），
      // 并把本次打断落库的前半段挂到这条消息上（删除时一起删，详见 hydrateMessageIds）
      await hydrateMessageIds(sessionId, task.messages);
    } catch (error) {
      if ((error as Error)?.name === 'AbortError') {
        lastMsg.isLoading = false;
        lastMsg.durationMs = (lastMsg.baseDurationMs || 0) + (lastMsg.generateStartTime ? Date.now() - lastMsg.generateStartTime : 0);
        if (!lastMsg.content) {
          lastMsg.content = '已停止生成。';
        }
      } else {
        console.error('提交答案失败:', error);
        const isStreamError = error instanceof StreamError;
        message.error(isStreamError ? error.message : '提交失败，请重试');
        if (isStreamError) refreshUsage();
        lastMsg.isLoading = false;
        lastMsg.durationMs = (lastMsg.baseDurationMs || 0) + (lastMsg.generateStartTime ? Date.now() - lastMsg.generateStartTime : 0);
        if (!lastMsg.content) {
          lastMsg.content = isStreamError ? error.message : '抱歉，提交请求失败，请稍后重试。';
        }
      }
    } finally {
      setTaskAbortController(task.id, null);
      task.isStreaming = false;
      // 仅当消息不再处于审批状态时才清除（避免 resume 流返回新 interrupt 被误清）
      if (!lastMsg.isPendingApproval) {
        clearTaskPendingApproval(task.id);
      }
      if (getCurrentTaskId() === task.id) {
        await scrollToBottom();
      }
    }
  };

  const handleAskUserCancel = async () => {
    // 防重复操作：如果当前任务正在流式中，忽略
    const currentTask = getTaskList().find((t) => t.id === getCurrentTaskId());
    if (currentTask?.isStreaming) return;

    const found = findPendingApprovalTask();

    // 优先从 taskPendingApprovals 获取，若被清除则从消息的 pendingApprovalData 回退
    let sessionId = found?.approval.sessionId;
    let threadId = found?.approval.threadId;
    const foundTaskId = found?.taskId;

    // 查找任务：优先用 foundTaskId，回退到 currentTaskId
    const task = (foundTaskId ? getTaskList().find((t) => t.id === foundTaskId) : null) || getTaskList().find((t) => t.id === getCurrentTaskId());
    if (!task) return;

    // 切换到目标任务（如果不在当前视图）
    if (task.id !== getCurrentTaskId()) {
      await switchTask(task.id);
    }

    // 单一数据源：消息只存在于任务自己的 messages 数组里
    const sourceArr = task.messages;
    const lastMsg = sourceArr[sourceArr.length - 1];
    if (!lastMsg || !lastMsg.isPendingApproval) return;

    // 若 taskPendingApprovals 已被清除，从消息的 pendingApprovalData 回退取值
    if (!sessionId && lastMsg.pendingApprovalData) {
      sessionId = lastMsg.pendingApprovalData.sessionId;
      threadId = lastMsg.pendingApprovalData.threadId;
    }
    if (!sessionId) return;

    // 冻结累计耗时到 durationMs，避免取消后计时消失
    lastMsg.durationMs = (lastMsg.baseDurationMs || 0) + (lastMsg.generateStartTime ? Date.now() - lastMsg.generateStartTime : 0);
    lastMsg.isPendingApproval = false;
    lastMsg.pendingApprovalData = undefined;
    lastMsg.isLoading = false;
    lastMsg.content = (lastMsg.content ? lastMsg.content + '\n\n' : '') + '已取消提问。';

    clearTaskPendingApproval(task.id);

    try {
      await chatResumeStream(
        {
          session_id: sessionId,
          thread_id: threadId || undefined,
          status: 'cancelled',
        },
        () => {}
      );
    } catch (error) {
      console.error('取消操作失败:', error);
    }

    task.isStreaming = false;

    if (getCurrentTaskId() === task.id) {
      await scrollToBottom();
    }
  };
  return {
    setTaskPendingApproval,
    clearTaskPendingApproval,
    migrateTaskApproval,
    currentPendingApproval,
    findPendingApprovalTask,
    handleAskUserSubmit,
    handleAskUserCancel,
  };
};