AgentCoord/frontend/src/api/index.ts

import request from '@/utils/request'
import websocket from '@/utils/websocket'
import type { Agent, IApiStepTask, IRawPlanResponse, IRawStepTask } from '@/stores'
import {
  mockBackendAgentSelectModifyInit,
  mockBackendAgentSelectModifyAddAspect,
  type BackendAgentScoreResponse,
} from '@/layout/components/Main/TaskTemplate/TaskSyllabus/components/mock/AgentAssignmentBackendMock'
import {
  mockBackendFillAgentTaskProcess,
  type RawAgentTaskProcessResponse,
} from '@/layout/components/Main/TaskTemplate/TaskProcess/components/mock/AgentTaskProcessBackendMock'
import { mockBranchPlanOutlineAPI } from '@/layout/components/Main/TaskTemplate/TaskSyllabus/Branch/mock/branchPlanOutlineMock'
import { mockFillStepTaskAPI } from '@/layout/components/Main/TaskTemplate/TaskSyllabus/Branch/mock/fill-step-task-mock'
import {
  mockBranchTaskProcessAPI,
  type BranchAction,
} from '@/layout/components/Main/TaskTemplate/TaskSyllabus/Branch/mock/branchTaskProcessMock'

export interface ActionHistory {
  ID: string
  ActionType: string
  AgentName: string
  Description: string
  ImportantInput: string[]
  Action_Result: string
}

export type IExecuteRawResponse = {
  LogNodeType: string
  NodeId: string
  InputName_List?: string[] | null
  OutputName?: string
  content?: string
  ActionHistory: ActionHistory[]
}

/**
 * SSE 流式事件类型
 */
export type StreamingEvent =
  | {
      type: 'step_start'
      step_index: number
      total_steps: number
      step_name: string
      task_description?: string
    }
  | {
      type: 'action_complete'
      step_index: number
      step_name: string
      action_index: number
      total_actions: number
      completed_actions: number
      action_result: ActionHistory
      batch_info?: {
        batch_index: number
        batch_size: number
        is_parallel: boolean
      }
    }
  | {
      type: 'step_complete'
      step_index: number
      step_name: string
      step_log_node: any
      object_log_node: any
    }
  | {
      type: 'execution_complete'
      total_steps: number
    }
  | {
      type: 'error'
      message: string
    }

export interface IFillAgentSelectionRequest {
  goal: string
  stepTask: IApiStepTask
  agents: string[]
}

class Api {
  // 默认使用WebSocket
  private useWebSocketDefault = true

  setAgents = (data: Pick<Agent, 'Name' | 'Profile' | 'apiUrl' | 'apiKey' | 'apiModel'>[], useWebSocket: boolean = this.useWebSocketDefault) => {
    // 如果启用WebSocket且已连接，使用WebSocket
    if (useWebSocket && websocket.connected) {
      return websocket.send('set_agents', data)
    }

    // 否则使用REST API
    return request({
      url: '/setAgents',
      data,
      method: 'POST',
    })
  }

  generateBasePlan = (data: {
    goal: string
    inputs: string[]
    apiUrl?: string
    apiKey?: string
    apiModel?: string
    useWebSocket?: boolean
    onProgress?: (progress: { status: string; stage?: string; message?: string; [key: string]: any }) => void
  }) => {
    const useWs = data.useWebSocket !== undefined ? data.useWebSocket : this.useWebSocketDefault

    // 如果启用WebSocket且已连接，使用WebSocket
    if (useWs && websocket.connected) {
      return websocket.send('generate_base_plan', {
        'General Goal': data.goal,
        'Initial Input Object': data.inputs,
        apiUrl: data.apiUrl,
        apiKey: data.apiKey,
        apiModel: data.apiModel,
      }, undefined, data.onProgress)
    }

    // 否则使用REST API
    return request<unknown, IRawPlanResponse>({
      url: '/generate_basePlan',
      method: 'POST',
      data: {
        'General Goal': data.goal,
        'Initial Input Object': data.inputs,
        apiUrl: data.apiUrl,
        apiKey: data.apiKey,
        apiModel: data.apiModel,
      },
    })
  }

  executePlan = (plan: IRawPlanResponse) => {
    return request<unknown, IExecuteRawResponse[]>({
      url: '/executePlan',
      method: 'POST',
      data: {
        RehearsalLog: [],
        num_StepToRun: null,
        plan: {
          'Initial Input Object': plan['Initial Input Object'],
          'General Goal': plan['General Goal'],
          'Collaboration Process': plan['Collaboration Process']?.map((step) => ({
            StepName: step.StepName,
            TaskContent: step.TaskContent,
            InputObject_List: step.InputObject_List,
            OutputObject: step.OutputObject,
            AgentSelection: step.AgentSelection,
            Collaboration_Brief_frontEnd: step.Collaboration_Brief_frontEnd,
            TaskProcess: step.TaskProcess.map((action) => ({
              ActionType: action.ActionType,
              AgentName: action.AgentName,
              Description: action.Description,
              ID: action.ID,
              ImportantInput: action.ImportantInput,
            })),
          })),
        },
      },
    })
  }

  /**
   * 优化版流式执行计划（阶段1+2：步骤级流式 + 动作级智能并行）
   * 无依赖关系的动作并行执行，有依赖关系的动作串行执行
   *
   * 默认使用WebSocket，如果连接失败则降级到SSE
   */
  executePlanOptimized = (
    plan: IRawPlanResponse,
    onMessage: (event: StreamingEvent) => void,
    onError?: (error: Error) => void,
    onComplete?: () => void,
    useWebSocket?: boolean,
  ) => {
    const useWs = useWebSocket !== undefined ? useWebSocket : this.useWebSocketDefault

    const data = {
      RehearsalLog: [],
      num_StepToRun: null,
      plan: {
        'Initial Input Object': plan['Initial Input Object'],
        'General Goal': plan['General Goal'],
        'Collaboration Process': plan['Collaboration Process']?.map((step) => ({
          StepName: step.StepName,
          TaskContent: step.TaskContent,
          InputObject_List: step.InputObject_List,
          OutputObject: step.OutputObject,
          AgentSelection: step.AgentSelection,
          Collaboration_Brief_frontEnd: step.Collaboration_Brief_frontEnd,
          TaskProcess: step.TaskProcess.map((action) => ({
            ActionType: action.ActionType,
            AgentName: action.AgentName,
            Description: action.Description,
            ID: action.ID,
            ImportantInput: action.ImportantInput,
          })),
        })),
      },
    }

    // 如果启用WebSocket且已连接，使用WebSocket
    if (useWs && websocket.connected) {
      websocket.subscribe(
        'execute_plan_optimized',
        data,
        // onProgress
        (progressData) => {
          try {
            // progressData 应该已经是解析后的对象了
            // 如果是字符串，说明后端发送的是 JSON 字符串，需要解析
            let event: StreamingEvent
            if (typeof progressData === 'string') {
              event = JSON.parse(progressData)
            } else {
              event = progressData as StreamingEvent
            }
            onMessage(event)
          } catch (e) {
            // Failed to parse WebSocket data
          }
        },
        // onComplete
        () => {
          onComplete?.()
        },
        // onError
        (error) => {
          onError?.(error)
        }
      )
      return
    }

    // 否则使用原有的SSE方式

    fetch('/api/executePlanOptimized', {
      method: 'POST',
      headers: {
        'Content-Type': 'application/json',
      },
      body: JSON.stringify(data),
    })
      .then(async (response) => {
        if (!response.ok) {
          throw new Error(`HTTP error! status: ${response.status}`)
        }

        const reader = response.body?.getReader()
        const decoder = new TextDecoder()

        if (!reader) {
          throw new Error('Response body is null')
        }

        let buffer = ''

        while (true) {
          const { done, value } = await reader.read()

          if (done) {
            onComplete?.()
            break
          }

          buffer += decoder.decode(value, { stream: true })

          const lines = buffer.split('\n')
          buffer = lines.pop() || ''

          for (const line of lines) {
            if (line.startsWith('data: ')) {
              const data = line.slice(6)
              try {
                const event = JSON.parse(data)
                onMessage(event)
              } catch (e) {
                // Failed to parse SSE data
              }
            }
          }
        }
      })
      .catch((error) => {
        onError?.(error)
      })
  }

  /**
   * 分支任务大纲
   */
  branchPlanOutline = (data: {
    branch_Number: number
    Modification_Requirement: string
    Existing_Steps: IRawStepTask[]
    Baseline_Completion: number
    initialInputs: string[]
    goal: string
    useWebSocket?: boolean
    onProgress?: (progress: { status: string; stage?: string; message?: string; [key: string]: any }) => void
  }) => {
    const useWs = data.useWebSocket !== undefined ? data.useWebSocket : this.useWebSocketDefault

    // 如果启用WebSocket且已连接，使用WebSocket
    if (useWs && websocket.connected) {
      return websocket.send('branch_plan_outline', {
        branch_Number: data.branch_Number,
        Modification_Requirement: data.Modification_Requirement,
        Existing_Steps: data.Existing_Steps,
        Baseline_Completion: data.Baseline_Completion,
        'Initial Input Object': data.initialInputs,
        'General Goal': data.goal,
      }, undefined, data.onProgress)
    }

    // 否则使用REST API
    return request<unknown, IRawPlanResponse>({
      url: '/branch_PlanOutline',
      method: 'POST',
      data: {
        branch_Number: data.branch_Number,
        Modification_Requirement: data.Modification_Requirement,
        Existing_Steps: data.Existing_Steps,
        Baseline_Completion: data.Baseline_Completion,
        'Initial Input Object': data.initialInputs,
        'General Goal': data.goal,
      },
    })
  }

  /**
   * 分支任务流程
   */
  branchTaskProcess = (data: {
    branch_Number: number
    Modification_Requirement: string
    Existing_Steps: BranchAction[]
    Baseline_Completion: number
    stepTaskExisting: any
    goal: string
    useWebSocket?: boolean
    onProgress?: (progress: { status: string; stage?: string; message?: string; [key: string]: any }) => void
  }) => {
    const useWs = data.useWebSocket !== undefined ? data.useWebSocket : this.useWebSocketDefault

    // 如果启用WebSocket且已连接，使用WebSocket
    if (useWs && websocket.connected) {
      return websocket.send('branch_task_process', {
        branch_Number: data.branch_Number,
        Modification_Requirement: data.Modification_Requirement,
        Existing_Steps: data.Existing_Steps,
        Baseline_Completion: data.Baseline_Completion,
        stepTaskExisting: data.stepTaskExisting,
        'General Goal': data.goal,
      }, undefined, data.onProgress)
    }

    // 否则使用REST API
    return request<unknown, BranchAction[][]>({
      url: '/branch_TaskProcess',
      method: 'POST',
      data: {
        branch_Number: data.branch_Number,
        Modification_Requirement: data.Modification_Requirement,
        Existing_Steps: data.Existing_Steps,
        Baseline_Completion: data.Baseline_Completion,
        stepTaskExisting: data.stepTaskExisting,
        'General Goal': data.goal,
      },
    })
  }

  fillStepTask = async (data: {
    goal: string
    stepTask: any
    useWebSocket?: boolean
    onProgress?: (progress: { status: string; stage?: string; message?: string; [key: string]: any }) => void
  }): Promise<IRawStepTask> => {
    const useWs = data.useWebSocket !== undefined ? data.useWebSocket : this.useWebSocketDefault
    let response: any

    // 如果启用WebSocket且已连接，使用WebSocket
    if (useWs && websocket.connected) {
      response = await websocket.send('fill_step_task', {
        'General Goal': data.goal,
        stepTask: data.stepTask,
      }, undefined, data.onProgress)
    } else {
      // 否则使用REST API
      response = await request<
        {
          'General Goal': string
          stepTask: any
        },
        {
          AgentSelection?: string[]
          Collaboration_Brief_FrontEnd?: {
            template: string
            data: Record<string, { text: string; color: number[] }>
          }
          InputObject_List?: string[]
          OutputObject?: string
          StepName?: string
          TaskContent?: string
          TaskProcess?: Array<{
            ID: string
            ActionType: string
            AgentName: string
            Description: string
            ImportantInput: string[]
          }>
        }
      >({
        url: '/fill_stepTask',
        method: 'POST',
        data: {
          'General Goal': data.goal,
          stepTask: data.stepTask,
        },
      })
    }

    const vec2Hsl = (color: number[]): string => {
      const [h, s, l] = color
      return `hsl(${h}, ${s}%, ${l}%)`
    }

    const briefData: Record<string, { text: string; style?: Record<string, string> }> = {}
    if (response.Collaboration_Brief_FrontEnd?.data) {
      for (const [key, value] of Object.entries(response.Collaboration_Brief_FrontEnd.data)) {
        briefData[key] = {
          text: value.text,
          style: {
            background: vec2Hsl(value.color),
          },
        }
      }
    }

    /**
     * 构建前端格式的 IRawStepTask
     */
    return {
      StepName: response.StepName || '',
      TaskContent: response.TaskContent || '',
      InputObject_List: response.InputObject_List || [],
      OutputObject: response.OutputObject || '',
      AgentSelection: response.AgentSelection || [],
      Collaboration_Brief_frontEnd: {
        template: response.Collaboration_Brief_FrontEnd?.template || '',
        data: briefData,
      },
      TaskProcess: response.TaskProcess || [],
    }
  }

  fillStepTaskTaskProcess = async (data: {
    goal: string
    stepTask: IApiStepTask
    agents: string[]
    useWebSocket?: boolean
    onProgress?: (progress: { status: string; stage?: string; message?: string; [key: string]: any }) => void
  }): Promise<IApiStepTask> => {
    const useWs = data.useWebSocket !== undefined ? data.useWebSocket : this.useWebSocketDefault
    let response: any

    // 如果启用WebSocket且已连接，使用WebSocket
    if (useWs && websocket.connected) {
      response = await websocket.send('fill_step_task_process', {
        'General Goal': data.goal,
        stepTask_lackTaskProcess: {
          StepName: data.stepTask.name,
          TaskContent: data.stepTask.content,
          InputObject_List: data.stepTask.inputs,
          OutputObject: data.stepTask.output,
          AgentSelection: data.agents,
        },
      }, undefined, data.onProgress)
    } else {
      // 否则使用REST API
      response = await request<
        {
          'General Goal': string
          stepTask_lackTaskProcess: {
            StepName: string
            TaskContent: string
            InputObject_List: string[]
            OutputObject: string
            AgentSelection: string[]
          }
        },
        {
          StepName?: string
          TaskContent?: string
          InputObject_List?: string[]
          OutputObject?: string
          AgentSelection?: string[]
          TaskProcess?: Array<{
            ID: string
            ActionType: string
            AgentName: string
            Description: string
            ImportantInput: string[]
          }>
          Collaboration_Brief_FrontEnd?: {
            template: string
            data: Record<string, { text: string; color: number[] }>
          }
        }
      >({
        url: '/fill_stepTask_TaskProcess',
        method: 'POST',
        data: {
          'General Goal': data.goal,
          stepTask_lackTaskProcess: {
            StepName: data.stepTask.name,
            TaskContent: data.stepTask.content,
            InputObject_List: data.stepTask.inputs,
            OutputObject: data.stepTask.output,
            AgentSelection: data.agents,
          },
        },
      })
    }

    const vec2Hsl = (color: number[]): string => {
      const [h, s, l] = color
      return `hsl(${h}, ${s}%, ${l}%)`
    }

    const briefData: Record<string, { text: string; style: { background: string } }> = {}
    if (response.Collaboration_Brief_FrontEnd?.data) {
      for (const [key, value] of Object.entries(response.Collaboration_Brief_FrontEnd.data)) {
        briefData[key] = {
          text: value.text,
          style: {
            background: vec2Hsl(value.color),
          },
        }
      }
    }

    const process = (response.TaskProcess || []).map((action: any) => ({
      id: action.ID,
      type: action.ActionType,
      agent: action.AgentName,
      description: action.Description,
      inputs: action.ImportantInput,
    }))

    return {
      name: response.StepName || '',
      content: response.TaskContent || '',
      inputs: response.InputObject_List || [],
      output: response.OutputObject || '',
      agents: response.AgentSelection || [],
      brief: {
        template: response.Collaboration_Brief_FrontEnd?.template || '',
        data: briefData,
      },
      process,
    }
  }

  /**
   * 为每个智能体评分
   */
  agentSelectModifyInit = async (data: {
    goal: string
    stepTask: any
    useWebSocket?: boolean
    onProgress?: (progress: { status: string; stage?: string; message?: string; [key: string]: any }) => void
  }): Promise<Record<string, Record<string, { reason: string; score: number }>>> => {
    const useWs = data.useWebSocket !== undefined ? data.useWebSocket : this.useWebSocketDefault
    let response: Record<string, Record<string, { Reason: string; Score: number }>>

    // 如果启用WebSocket且已连接，使用WebSocket
    if (useWs && websocket.connected) {
      response = await websocket.send('agent_select_modify_init', {
        'General Goal': data.goal,
        stepTask: {
          StepName: data.stepTask.StepName || data.stepTask.name,
          TaskContent: data.stepTask.TaskContent || data.stepTask.content,
          InputObject_List: data.stepTask.InputObject_List || data.stepTask.inputs,
          OutputObject: data.stepTask.OutputObject || data.stepTask.output,
        },
      }, undefined, data.onProgress)
    } else {
      // 否则使用REST API
      response = await request<
        {
          'General Goal': string
          stepTask: any
        },
        Record<string, Record<string, { Reason: string; Score: number }>>
      >({
        url: '/agentSelectModify_init',
        method: 'POST',
        data: {
          'General Goal': data.goal,
          stepTask: {
            StepName: data.stepTask.StepName || data.stepTask.name,
            TaskContent: data.stepTask.TaskContent || data.stepTask.content,
            InputObject_List: data.stepTask.InputObject_List || data.stepTask.inputs,
            OutputObject: data.stepTask.OutputObject || data.stepTask.output,
          },
        },
      })
    }

    const transformedData: Record<string, Record<string, { reason: string; score: number }>> = {}

    for (const [aspect, agents] of Object.entries(response)) {
      for (const [agentName, scoreInfo] of Object.entries(agents)) {
        if (!transformedData[agentName]) {
          transformedData[agentName] = {}
        }
        transformedData[agentName][aspect] = {
          reason: scoreInfo.Reason,
          score: scoreInfo.Score,
        }
      }
    }

    return transformedData
  }

  /**
   * 添加新的评估维度
   */
  agentSelectModifyAddAspect = async (data: {
    aspectList: string[]
    useWebSocket?: boolean
    onProgress?: (progress: { status: string; stage?: string; message?: string; [key: string]: any }) => void
  }): Promise<{
    aspectName: string
    agentScores: Record<string, { score: number; reason: string }>
  }> => {
    const useWs = data.useWebSocket !== undefined ? data.useWebSocket : this.useWebSocketDefault
    let response: Record<string, Record<string, { Reason: string; Score: number }>>

    // 如果启用WebSocket且已连接，使用WebSocket
    if (useWs && websocket.connected) {
      response = await websocket.send('agent_select_modify_add_aspect', {
        aspectList: data.aspectList,
      }, undefined, data.onProgress)
    } else {
      // 否则使用REST API
      response = await request<
        {
          aspectList: string[]
        },
        Record<string, Record<string, { Reason: string; Score: number }>>
      >({
        url: '/agentSelectModify_addAspect',
        method: 'POST',
        data: {
          aspectList: data.aspectList,
        },
      })
    }

    /**
     * 获取新添加的维度
     */
    const newAspect = data.aspectList[data.aspectList.length - 1]
    if (!newAspect) {
      throw new Error('aspectList is empty')
    }

    const newAspectAgents = response[newAspect]
    const agentScores: Record<string, { score: number; reason: string }> = {}

    if (newAspectAgents) {
      for (const [agentName, scoreInfo] of Object.entries(newAspectAgents)) {
        agentScores[agentName] = {
          score: scoreInfo.Score,
          reason: scoreInfo.Reason,
        }
      }
    }

    return {
      aspectName: newAspect,
      agentScores,
    }
  }
  /**
   * ==================== Mock API（开发阶段使用）====================
   *为每个智能体评分
   */
  mockAgentSelectModifyInit = async (): Promise<
    Record<string, Record<string, { reason: string; score: number }>>
  > => {
    const response: BackendAgentScoreResponse = await mockBackendAgentSelectModifyInit()

    const transformedData: Record<string, Record<string, { reason: string; score: number }>> = {}

    for (const [aspect, agents] of Object.entries(response)) {
      for (const [agentName, scoreInfo] of Object.entries(agents)) {
        if (!transformedData[agentName]) {
          transformedData[agentName] = {}
        }
        transformedData[agentName][aspect] = {
          reason: scoreInfo.Reason,
          score: scoreInfo.Score,
        }
      }
    }

    return transformedData
  }

  mockAgentSelectModifyAddAspect = async (data: {
    aspectList: string[]
  }): Promise<{
    aspectName: string
    agentScores: Record<string, { score: number; reason: string }>
  }> => {
    const response: BackendAgentScoreResponse = await mockBackendAgentSelectModifyAddAspect(
      data.aspectList,
    )

    const newAspect = data.aspectList[data.aspectList.length - 1]
    if (!newAspect) {
      throw new Error('aspectList is empty')
    }

    const newAspectAgents = response[newAspect]
    const agentScores: Record<string, { score: number; reason: string }> = {}

    if (newAspectAgents) {
      for (const [agentName, scoreInfo] of Object.entries(newAspectAgents)) {
        agentScores[agentName] = {
          score: scoreInfo.Score,
          reason: scoreInfo.Reason,
        }
      }
    }

    return {
      aspectName: newAspect,
      agentScores,
    }
  }

  mockFillStepTaskTaskProcess = async (data: {
    goal: string
    stepTask: IApiStepTask
    agents: string[]
  }): Promise<IApiStepTask> => {
    const response: RawAgentTaskProcessResponse = await mockBackendFillAgentTaskProcess(
      data.goal,
      data.stepTask,
      data.agents,
    )

    const vec2Hsl = (color: number[]): string => {
      const [h, s, l] = color
      return `hsl(${h}, ${s}%, ${l}%)`
    }

    const briefData: Record<string, { text: string; style: { background: string } }> = {}
    if (response.Collaboration_Brief_frontEnd?.data) {
      for (const [key, value] of Object.entries(response.Collaboration_Brief_frontEnd.data)) {
        briefData[key] = {
          text: value.text,
          style: {
            background: vec2Hsl(value.color),
          },
        }
      }
    }

    const process = (response.TaskProcess || []).map((action) => ({
      id: action.ID,
      type: action.ActionType,
      agent: action.AgentName,
      description: action.Description,
      inputs: action.ImportantInput,
    }))

    return {
      name: response.StepName || '',
      content: response.TaskContent || '',
      inputs: response.InputObject_List || [],
      output: response.OutputObject || '',
      agents: response.AgentSelection || [],
      brief: {
        template: response.Collaboration_Brief_frontEnd?.template || '',
        data: briefData,
      },
      process,
    }
  }

  mockBranchPlanOutline = async (data: {
    branch_Number: number
    Modification_Requirement: string
    Existing_Steps: IRawStepTask[]
    Baseline_Completion: number
    initialInputs: string[]
    goal: string
  }): Promise<IRawPlanResponse> => {
    const response = await mockBranchPlanOutlineAPI({
      branch_Number: data.branch_Number,
      Modification_Requirement: data.Modification_Requirement,
      Existing_Steps: data.Existing_Steps,
      Baseline_Completion: data.Baseline_Completion,
      InitialObject_List: data.initialInputs,
      General_Goal: data.goal,
    })

    return response
  }

  mockFillStepTask = async (data: { goal: string; stepTask: any }): Promise<any> => {
    const response = await mockFillStepTaskAPI({
      General_Goal: data.goal,
      stepTask: data.stepTask,
    })

    return response
  }

  mockBranchTaskProcess = async (data: {
    branch_Number: number
    Modification_Requirement: string
    Existing_Steps: BranchAction[]
    Baseline_Completion: number
    stepTaskExisting: any
    goal: string
  }): Promise<BranchAction[][]> => {
    const response = await mockBranchTaskProcessAPI({
      branch_Number: data.branch_Number,
      Modification_Requirement: data.Modification_Requirement,
      Existing_Steps: data.Existing_Steps,
      Baseline_Completion: data.Baseline_Completion,
      stepTaskExisting: data.stepTaskExisting,
      General_Goal: data.goal,
    })

    return response
  }
}

export default new Api()