AgentCoord/frontend/src/stores/modules/selection.ts

import { ref } from 'vue'
import { defineStore } from 'pinia'
import { v4 as uuidv4 } from 'uuid'
import { store } from '../index'
import type { IRawStepTask, IApiStepTask } from './agents'
import type { Node, Edge } from '@vue-flow/core'

/**
 * 分支数据接口
 * @description 用于存储流程图中的分支节点和边数据
 */
export interface IBranchData {
  /** 分支唯一 ID */
  id: string
  /** 父节点 ID（根节点或任务节点） */
  parentNodeId: string
  /** 分支需求内容 */
  branchContent: string
  /** 分支类型 */
  branchType: 'root' | 'task'
  /** 分支包含的所有节点 */
  nodes: Node[]
  /** 分支包含的所有边 */
  edges: Edge[]
  /** 分支的任务数据 */
  tasks: IRawStepTask[]
  /** 创建时间 */
  createdAt: number
}

export const useSelectionStore = defineStore('selection', () => {
  // ==================== 任务大纲探索分支数据存储 ====================

  /** 流程图分支列表 */
  const flowBranches = ref<IBranchData[]>([])

  /**
   * 添加流程图分支
   * @param data 分支数据
   * @returns 分支 ID
   */
  function addFlowBranch(data: {
    parentNodeId: string
    branchContent: string
    branchType: 'root' | 'task'
    nodes: Node[]
    edges: Edge[]
    tasks: IRawStepTask[]
  }): string {
    const branchId = `flow-branch-${uuidv4()}`
    const newBranch: IBranchData = {
      id: branchId,
      parentNodeId: data.parentNodeId,
      branchContent: data.branchContent,
      branchType: data.branchType,
      nodes: data.nodes,
      edges: data.edges,
      tasks: data.tasks,
      createdAt: Date.now(),
    }
    flowBranches.value.push(newBranch)
    return branchId
  }

  /**
   * 获取所有流程图分支
   * @returns 所有流程图分支数据
   */
  function getAllFlowBranches(): IBranchData[] {
    return flowBranches.value
  }

  /**
   * 根据父节点 ID 获取流程图分支
   * @param parentNodeId 父节点 ID
   * @returns 匹配的流程图分支列表
   */
  function getFlowBranchesByParent(parentNodeId: string): IBranchData[] {
    return flowBranches.value.filter((branch) => branch.parentNodeId === parentNodeId)
  }

  /**
   * 删除流程图分支
   * @param branchId 分支 ID
   * @returns 是否删除成功
   */
  function removeFlowBranch(branchId: string): boolean {
    const index = flowBranches.value.findIndex((branch) => branch.id === branchId)
    if (index > -1) {
      flowBranches.value.splice(index, 1)
      return true
    }
    return false
  }

  /** 清除所有流程图分支 */
  function clearFlowBranches() {
    flowBranches.value = []
  }

  /**
   * 根据父节点 ID 清除流程图分支
   * @param parentNodeId 父节点 ID
   */
  function clearFlowBranchesByParent(parentNodeId: string) {
    flowBranches.value = flowBranches.value.filter((branch) => branch.parentNodeId !== parentNodeId)
  }

  // ==================== 任务过程探索分支数据存储 ====================
  /**
   * 任务过程分支数据映射
   * @description 存储结构: Map<taskStepId, Map<agentGroupKey, IBranchData[]>>
   * - taskStepId: 任务步骤 ID
   * - agentGroupKey: agent 组合的唯一标识（排序后的 JSON 字符串）
   * - IBranchData[]: 该 agent 组合在该任务下的所有分支数据
   */
  const taskProcessBranchesMap = ref<Map<string, Map<string, IBranchData[]>>>(new Map())

  /**
   * 添加任务过程分支
   * @param taskStepId 任务步骤 ID
   * @param agents Agent 列表
   * @param data 分支数据
   * @returns 分支 ID
   */
  function addTaskProcessBranch(
    taskStepId: string,
    agents: string[],
    data: {
      parentNodeId: string
      branchContent: string
      branchType: 'root' | 'task'
      nodes: Node[]
      edges: Edge[]
      tasks: IRawStepTask[]
    },
  ): string {
    const branchId = `task-process-branch-${uuidv4()}`
    const agentGroupKey = getAgentGroupKey(agents)

    const newBranch: IBranchData = {
      id: branchId,
      parentNodeId: data.parentNodeId,
      branchContent: data.branchContent,
      branchType: data.branchType,
      nodes: data.nodes,
      edges: data.edges,
      tasks: data.tasks,
      createdAt: Date.now(),
    }

    // 获取或创建该任务步骤的 Map
    if (!taskProcessBranchesMap.value.has(taskStepId)) {
      taskProcessBranchesMap.value.set(taskStepId, new Map())
    }

    // 获取或创建该 agent 组合的分支列表
    const agentMap = taskProcessBranchesMap.value.get(taskStepId)!
    if (!agentMap.has(agentGroupKey)) {
      agentMap.set(agentGroupKey, [])
    }
    agentMap.get(agentGroupKey)!.push(newBranch)

    return branchId
  }

  /**
   * 获取指定任务步骤和 agent 组合的所有分支
   * @param taskStepId 任务步骤 ID
   * @param agents Agent 列表
   * @returns 匹配的任务过程分支列表
   */
  function getTaskProcessBranches(taskStepId: string, agents: string[]): IBranchData[] {
    const agentGroupKey = getAgentGroupKey(agents)
    return taskProcessBranchesMap.value.get(taskStepId)?.get(agentGroupKey) || []
  }

  /**
   * 获取所有任务过程分支
   * @returns 所有任务过程分支数据映射
   */
  function getAllTaskProcessBranches(): Map<string, Map<string, IBranchData[]>> {
    return taskProcessBranchesMap.value
  }

  /**
   * 根据父节点 ID 获取任务过程分支
   * @param taskStepId 任务步骤 ID
   * @param agents Agent 列表
   * @param parentNodeId 父节点 ID
   * @returns 匹配的任务过程分支列表
   */
  function getTaskProcessBranchesByParent(
    taskStepId: string,
    agents: string[],
    parentNodeId: string,
  ): IBranchData[] {
    const agentGroupKey = getAgentGroupKey(agents)
    const branches = taskProcessBranchesMap.value.get(taskStepId)?.get(agentGroupKey) || []
    return branches.filter((branch) => branch.parentNodeId === parentNodeId)
  }

  /**
   * 删除任务过程分支
   * @param taskStepId 任务步骤 ID
   * @param agents Agent 列表
   * @param branchId 分支 ID
   * @returns 是否删除成功
   */
  function removeTaskProcessBranch(
    taskStepId: string,
    agents: string[],
    branchId: string,
  ): boolean {
    const agentGroupKey = getAgentGroupKey(agents)
    const branches = taskProcessBranchesMap.value.get(taskStepId)?.get(agentGroupKey)
    if (branches) {
      const index = branches.findIndex((branch) => branch.id === branchId)
      if (index > -1) {
        branches.splice(index, 1)
        return true
      }
    }
    return false
  }

  /**
   * 清除指定任务步骤和 agent 组合的所有分支
   * @param taskStepId 任务步骤 ID
   * @param agents Agent 列表（可选，不传则清除该任务步骤的所有分支）
   */
  function clearTaskProcessBranches(taskStepId: string, agents?: string[]) {
    if (agents) {
      // 清除指定 agent 组合的分支
      const agentGroupKey = getAgentGroupKey(agents)
      const agentMap = taskProcessBranchesMap.value.get(taskStepId)
      if (agentMap) {
        agentMap.delete(agentGroupKey)
      }
    } else {
      // 清除该任务步骤的所有分支（所有 agent 组合）
      taskProcessBranchesMap.value.delete(taskStepId)
    }
  }

  // ==================== Agent 组合 TaskProcess 数据存储 ====================
  /**
   * Agent 组合 TaskProcess 数据映射
   * @description 用于存储 fill_stepTask_TaskProcess 接口返回的数据
   * 存储结构: Map<taskId, Map<agentGroupKey, IApiStepTask>>
   * - taskId: 任务 ID
   * - agentGroupKey: agent 组合的唯一标识（排序后的 JSON 字符串）
   * - IApiStepTask: 该 agent 组合的完整 TaskProcess 数据
   */
  const agentTaskProcessMap = ref<Map<string, Map<string, IApiStepTask>>>(new Map())

  /**
   * 生成 agent 组合的唯一 key
   * @description 排序后保证一致性
   * @param agents Agent 列表
   * @returns Agent 组合的唯一标识
   */
  function getAgentGroupKey(agents: string[]): string {
    // 处理 undefined 或 null 的情况
    if (!agents || !Array.isArray(agents)) {
      return JSON.stringify([])
    }
    return JSON.stringify([...agents].sort())
  }

  /**
   * 存储 agent 组合的 TaskProcess 数据
   * @param taskId 任务 ID
   * @param agents Agent 列表
   * @param taskProcess TaskProcess 数据
   */
  function setAgentTaskProcess(taskId: string, agents: string[], taskProcess: IApiStepTask) {
    const groupKey = getAgentGroupKey(agents)

    // 获取或创建该任务的 Map
    if (!agentTaskProcessMap.value.has(taskId)) {
      agentTaskProcessMap.value.set(taskId, new Map())
    }

    // 存储该 agent 组合的 TaskProcess 数据
    agentTaskProcessMap.value.get(taskId)!.set(groupKey, taskProcess)
  }

  /**
   * 获取 agent 组合的 TaskProcess 数据
   * @param taskId 任务 ID
   * @param agents Agent 列表
   * @returns TaskProcess 数据
   */
  function getAgentTaskProcess(taskId: string, agents: string[]): IApiStepTask | undefined {
    const groupKey = getAgentGroupKey(agents)
    return agentTaskProcessMap.value.get(taskId)?.get(groupKey)
  }

  /**
   * 检查 agent 组合是否已有 TaskProcess 数据
   * @param taskId 任务 ID
   * @param agents Agent 列表
   * @returns 是否存在数据
   */
  function hasAgentTaskProcess(taskId: string, agents: string[]): boolean {
    const groupKey = getAgentGroupKey(agents)
    return agentTaskProcessMap.value.get(taskId)?.has(groupKey) || false
  }

  /**
   * 清除指定任务的所有 agent 组合 TaskProcess 数据
   * @param taskId 任务 ID
   */
  function clearAgentTaskProcess(taskId: string) {
    agentTaskProcessMap.value.delete(taskId)
  }

  /** 清除所有任务的 agent 组合 TaskProcess 数据 */
  function clearAllAgentTaskProcess() {
    agentTaskProcessMap.value.clear()
  }

  // ==================== 当前生效的任务过程分支 ====================
  /**
   * 当前生效的任务过程分支映射
   * @description 记录每个任务步骤和 agent 组合当前生效的分支 ID（持久化选中状态）
   * 存储结构: Map<taskStepId, Map<agentGroupKey, branchId>>
   * - taskStepId: 任务步骤 ID
   * - agentGroupKey: agent 组合的唯一标识
   * - branchId: 当前选中的分支 ID
   */
  const activeTaskProcessBranchMap = ref<Map<string, Map<string, string>>>(new Map())

  /**
   * 当前生效的 TaskProcess 数据映射
   * @description 用于外部组件显示职责分配
   * 存储结构: Map<taskStepId, Map<agentGroupKey, TaskProcess[]>>
   */
  const activeTaskProcessDataMap = ref<Map<string, Map<string, any[]>>>(new Map())

  /**
   * 设置当前生效的分支
   * @param taskStepId 任务步骤 ID
   * @param agents Agent 列表
   * @param branchId 分支 ID
   */
  function setActiveTaskProcessBranch(taskStepId: string, agents: string[], branchId: string) {
    const agentGroupKey = getAgentGroupKey(agents)

    // 获取或创建该任务步骤的 Map
    if (!activeTaskProcessBranchMap.value.has(taskStepId)) {
      activeTaskProcessBranchMap.value.set(taskStepId, new Map())
    }

    activeTaskProcessBranchMap.value.get(taskStepId)!.set(agentGroupKey, branchId)
  }

  /**
   * 设置当前生效的 TaskProcess 数据
   * @param taskStepId 任务步骤 ID
   * @param agents Agent 列表
   * @param taskProcess TaskProcess 数据
   */
  function setActiveTaskProcessData(taskStepId: string, agents: string[], taskProcess: any[]) {
    const agentGroupKey = getAgentGroupKey(agents)

    // 获取或创建该任务步骤的 Map
    if (!activeTaskProcessDataMap.value.has(taskStepId)) {
      activeTaskProcessDataMap.value.set(taskStepId, new Map())
    }

    activeTaskProcessDataMap.value.get(taskStepId)!.set(agentGroupKey, taskProcess)
  }

  /**
   * 获取当前生效的分支 ID
   * @param taskStepId 任务步骤 ID
   * @param agents Agent 列表
   * @returns 分支 ID
   */
  function getActiveTaskProcessBranch(taskStepId: string, agents: string[]): string | undefined {
    const agentGroupKey = getAgentGroupKey(agents)
    return activeTaskProcessBranchMap.value.get(taskStepId)?.get(agentGroupKey)
  }

  /**
   * 获取当前生效的 TaskProcess 数据
   * @param taskStepId 任务步骤 ID
   * @param agents Agent 列表
   * @returns TaskProcess 数据
   */
  function getActiveTaskProcessData(taskStepId: string, agents: string[]): any[] | undefined {
    const agentGroupKey = getAgentGroupKey(agents)
    return activeTaskProcessDataMap.value.get(taskStepId)?.get(agentGroupKey)
  }

  /**
   * 清除生效分支
   * @param taskStepId 任务步骤 ID
   * @param agents Agent 列表（可选，不传则清除该任务步骤的所有生效分支）
   */
  function clearActiveTaskProcessBranch(taskStepId: string, agents?: string[]) {
    if (agents) {
      // 清除指定 agent 组合的生效分支
      const agentGroupKey = getAgentGroupKey(agents)
      activeTaskProcessBranchMap.value.get(taskStepId)?.delete(agentGroupKey)
      activeTaskProcessDataMap.value.get(taskStepId)?.delete(agentGroupKey)
    } else {
      // 清除该任务步骤的所有生效分支（所有 agent 组合）
      activeTaskProcessBranchMap.value.delete(taskStepId)
      activeTaskProcessDataMap.value.delete(taskStepId)
    }
  }

  return {
    // ==================== 状态 ====================
    flowBranches,
    taskProcessBranchesMap,
    agentTaskProcessMap,
    activeTaskProcessBranchMap,
    activeTaskProcessDataMap,

    // ==================== 任务大纲分支管理方法 ====================
    addFlowBranch,
    getAllFlowBranches,
    getFlowBranchesByParent,
    removeFlowBranch,
    clearFlowBranches,
    clearFlowBranchesByParent,

    // ==================== 任务过程分支管理方法 ====================
    addTaskProcessBranch,
    getTaskProcessBranches,
    getAllTaskProcessBranches,
    getTaskProcessBranchesByParent,
    removeTaskProcessBranch,
    clearTaskProcessBranches,

    // ==================== 任务过程分支生效状态管理方法 ====================
    setActiveTaskProcessBranch,
    setActiveTaskProcessData,
    getActiveTaskProcessBranch,
    getActiveTaskProcessData,
    clearActiveTaskProcessBranch,

    // ==================== Agent 组合 TaskProcess 数据管理方法 ====================
    getAgentGroupKey,
    setAgentTaskProcess,
    getAgentTaskProcess,
    hasAgentTaskProcess,
    clearAgentTaskProcess,
    clearAllAgentTaskProcess,
  }
})

/**
 * 用于在组件外部使用 Selection Store
 * @returns Selection Store 实例
 */
export function useSelectionStoreHook() {
  return useSelectionStore(store)
}