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 { withRetry } from '@/utils/retry'

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

export interface BranchAction {
  ID: string
  ActionType: string
  AgentName: string
  Description: string
  ImportantInput: 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,
            })),
          })),
        },
      },
    })
  }

  /**
   * 优化版流式执行计划（支持动态追加步骤）
   * 步骤级流式 + 动作级智能并行 + 动态追加步骤
   */
  executePlanOptimized = (
    plan: IRawPlanResponse,
    onMessage: (event: StreamingEvent) => void,
    onError?: (error: Error) => void,
    onComplete?: () => void,
    useWebSocket?: boolean,
    existingKeyObjects?: Record<string, any>,
    enableDynamic?: boolean,
    onExecutionStarted?: (executionId: string) => void,
    executionId?: string,
    restartFromStepIndex?: number, // 新增：从指定步骤重新执行的索引
    rehearsalLog?: any[], // 新增：传递截断后的 RehearsalLog
  ) => {
    const useWs = useWebSocket !== undefined ? useWebSocket : this.useWebSocketDefault

    const data = {
      RehearsalLog: rehearsalLog || [], // ✅ 使用传递的 RehearsalLog
      num_StepToRun: null,
      existingKeyObjects: existingKeyObjects || {},
      enable_dynamic: enableDynamic || false,
      execution_id: executionId || null,
      restart_from_step_index: restartFromStepIndex ?? 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 {
            let event: StreamingEvent

            // 处理不同类型的progress数据
            if (typeof progressData === 'string') {
              event = JSON.parse(progressData)
            } else {
              event = progressData as StreamingEvent
            }

            // 处理特殊事件类型
            if (event && typeof event === 'object') {
              // 检查是否是execution_started事件
              if ('status' in event && event.status === 'execution_started') {
                if ('execution_id' in event && onExecutionStarted) {
                  onExecutionStarted(event.execution_id as string)
                }
                return
              }
            }

            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
    generation_id?: string
    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

    // 定义实际的 API 调用逻辑
    const executeRequest = async (): Promise<any> => {
      if (useWs && websocket.connected) {
        return await websocket.send(
          'fill_step_task',
          {
            'General Goal': data.goal,
            stepTask: data.stepTask,
            generation_id: data.generation_id || '',
          },
          undefined,
          data.onProgress,
        )
      }
      // 否则使用REST API
      return 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 rawResponse = await withRetry(executeRequest, {
      maxRetries: 3,
      initialDelayMs: 2000,
      onRetry: (error, attempt, delay) => {
        console.warn(`⚠️ [fillStepTask] 第${attempt}次重试，等待 ${delay}ms...`, error?.message)
      },
    })

    // WebSocket 返回格式: { data: {...}, generation_id, execution_id }
    // REST API 返回格式: {...}
    const response = rawResponse.data || rawResponse

    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 as { text: string; color: number[] }).text,
          style: {
            background: vec2Hsl((value as { text: string; color: number[] }).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

    // 定义实际的 API 调用逻辑
    const executeRequest = async (): Promise<any> => {
      if (useWs && websocket.connected) {
        return 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,
        )
      }
      // 否则使用REST API
      return 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 rawResponse = await withRetry(executeRequest, {
      maxRetries: 3,
      initialDelayMs: 2000,
      onRetry: (error, attempt, delay) => {
        console.warn(
          `⚠️ [fillStepTaskTaskProcess] 第${attempt}次重试，等待 ${delay}ms...`,
          error?.message,
        )
      },
    })

    // WebSocket 返回格式: { data: {...}, generation_id, execution_id }
    // REST API 返回格式: {...}
    const response = rawResponse.data || rawResponse

    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 as { text: string; color: number[] }).text,
          style: {
            background: vec2Hsl((value as { text: string; color: number[] }).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

    // 调试日志：打印请求参数
    const requestPayload = {
      '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,
      },
    }
    console.log('🔍 [agentSelectModifyInit] 请求参数:', {
      goal: requestPayload['General Goal'],
      stepTaskName: requestPayload.stepTask.StepName,
      stepTaskContentLength: requestPayload.stepTask.TaskContent?.length,
      useWebSocket: useWs && websocket.connected,
      wsConnected: websocket.connected,
    })

    // 定义实际的 API 调用逻辑
    const executeRequest = async (): Promise<
      Record<string, Record<string, { Reason: string; Score: number }>>
    > => {
      if (useWs && websocket.connected) {
        return await websocket.send(
          'agent_select_modify_init',
          requestPayload,
          undefined,
          data.onProgress,
        )
      }
      // 否则使用REST API
      return await request<
        {
          'General Goal': string
          stepTask: any
        },
        Record<string, Record<string, { Reason: string; Score: number }>>
      >({
        url: '/agentSelectModify_init',
        method: 'POST',
        data: requestPayload,
      })
    }

    // 使用重试机制执行请求
    const rawResponse = await withRetry(executeRequest, {
      maxRetries: 3,
      initialDelayMs: 2000,
      onRetry: (error, attempt, delay) => {
        console.warn(
          `⚠️ [agentSelectModifyInit] 第${attempt}次重试，等待 ${delay}ms...`,
          error?.message,
        )
      },
    })

    // WebSocket 返回格式: { data: {...}, generation_id, execution_id }
    // REST API 返回格式: {...}
    const response = rawResponse.data || rawResponse

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

    // 确保 response 存在且是有效对象
    if (!response || typeof response !== 'object' || Array.isArray(response)) {
      console.warn('[agentSelectModifyInit] 后端返回数据格式异常:', response)
      return transformedData
    }

    for (const [aspect, agents] of Object.entries(response)) {
      for (const [agentName, scoreInfo] of Object.entries(agents as Record<string, { Reason: string; Score: number }> || {})) {
        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) {
      const rawResponse = await websocket.send(
        'agent_select_modify_add_aspect',
        {
          aspectList: data.aspectList,
        },
        undefined,
        data.onProgress,
      )
      // WebSocket 返回格式: { data: {...}, generation_id, execution_id }
      response = rawResponse.data || rawResponse
    } 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,
    }
  }
  /**
   * 向正在执行的任务追加新步骤
   * @param executionId 执行ID
   * @param newSteps 新步骤列表
   * @returns 追加的步骤数量
   */
  addStepsToExecution = async (executionId: string, newSteps: IRawStepTask[]): Promise<number> => {
    if (!websocket.connected) {
      throw new Error('WebSocket未连接')
    }

    const rawResponse = await websocket.send('add_steps_to_execution', {
      execution_id: executionId,
      new_steps: newSteps.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 返回格式: { data: {...}, generation_id, execution_id }
    // REST API 返回格式: {...}
    const response = (rawResponse.data || rawResponse) as { added_count: number }

    return response?.added_count || 0
  }
}

export default new Api()