AgentCoord/frontend/src/api/index.ts

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[]
}

/**
 * WebSocket 流式事件类型
 */
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 {
  // 提取响应数据的公共方法
  private extractResponse<T>(raw: any): T {
    return (raw.data || raw) as T
  }

  // 颜色向量转 HSL 字符串
  private vec2Hsl = (color: number[]): string => {
    const [h, s, l] = color
    return `hsl(${h}, ${s}%, ${l}%)`
  }

  setAgents = (data: Pick<Agent, 'Name' | 'Profile' | 'apiUrl' | 'apiKey' | 'apiModel'>[]) =>
    websocket.send('set_agents', data)

  getAgents = (user_id: string = 'default_user') => websocket.send('get_agents', { user_id })

  generateBasePlan = (data: {
    goal: string
    inputs: string[]
    apiUrl?: string
    apiKey?: string
    apiModel?: string
    onProgress?: (progress: {
      status: string
      stage?: string
      message?: string
      [key: string]: any
    }) => void
  }) =>
    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,
    )

  /**
   * 优化版流式执行计划（支持动态追加步骤）
   * 步骤级流式 + 动作级智能并行 + 动态追加步骤
   */
  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[],
    TaskID?: string,
  ) => {
    // eslint-disable-next-line @typescript-eslint/no-unused-vars
    void _useWebSocket // 保留参数位置以保持兼容性
    const data = {
      RehearsalLog: rehearsalLog || [], // 使用传递的 RehearsalLog
      num_StepToRun: null,
      existingKeyObjects: existingKeyObjects || {},
      enable_dynamic: enableDynamic || false,
      execution_id: executionId || null,
      restart_from_step_index: restartFromStepIndex ?? null, // 新增：传递重新执行索引
      task_id: TaskID || 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.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)
      },
    )
  }

  /**
   * 分支任务大纲
   */
  branchPlanOutline = (data: {
    branch_Number: number
    Modification_Requirement: string
    Existing_Steps: IRawStepTask[]
    Baseline_Completion: number
    initialInputs: string[]
    goal: string
    onProgress?: (progress: {
      status: string
      stage?: string
      message?: string
      [key: string]: any
    }) => void
  }) =>
    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,
    )

  /**
   * 分支任务流程
   */
  branchTaskProcess = (data: {
    branch_Number: number
    Modification_Requirement: string
    Existing_Steps: BranchAction[]
    Baseline_Completion: number
    stepTaskExisting: any
    goal: string
    onProgress?: (progress: {
      status: string
      stage?: string
      message?: string
      [key: string]: any
    }) => void
  }) =>
    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,
    )

  fillStepTask = async (data: {
    goal: string
    stepTask: any
    generation_id?: string
    TaskID?: string
    onProgress?: (progress: {
      status: string
      stage?: string
      message?: string
      [key: string]: any
    }) => void
  }): Promise<IRawStepTask> => {
    const rawResponse = await withRetry(
      () =>
        websocket.send(
          'fill_step_task',
          {
            'General Goal': data.goal,
            stepTask: data.stepTask,
            generation_id: data.generation_id || '',
            task_id: data.TaskID || '',
          },
          undefined,
          data.onProgress,
        ),
      {
        maxRetries: 3,
        initialDelayMs: 2000,
        onRetry: (error, attempt, delay) => {
          console.warn(` [fillStepTask] 第${attempt}次重试，等待 ${delay}ms...`, error?.message)
        },
      },
    )

    let response = this.extractResponse<any>(rawResponse)
    if (response?.filled_stepTask) {
      response = response.filled_stepTask
    }

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

    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[]
    TaskID?: string
    onProgress?: (progress: {
      status: string
      stage?: string
      message?: string
      [key: string]: any
    }) => void
  }): Promise<IApiStepTask> => {
    const rawResponse = await withRetry(
      () =>
        websocket.send(
          'fill_step_task_process',
          {
            'General Goal': data.goal,
            task_id: data.TaskID || undefined,
            stepTask_lackTaskProcess: {
              StepName: data.stepTask.name,
              TaskContent: data.stepTask.content,
              InputObject_List: data.stepTask.inputs,
              OutputObject: data.stepTask.output,
              AgentSelection: data.agents,
            },
            agents: data.agents,
          },
          undefined,
          data.onProgress,
        ),
      {
        maxRetries: 3,
        initialDelayMs: 2000,
        onRetry: (error, attempt, delay) => {
          console.warn(
            `[fillStepTaskTaskProcess] 第${attempt}次重试，等待 ${delay}ms...`,
            error?.message,
          )
        },
      },
    )

    let response = this.extractResponse<any>(rawResponse)
    if (response?.filled_stepTask) {
      response = response.filled_stepTask
    }

    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: this.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
    TaskID?: string
    onProgress?: (progress: {
      status: string
      stage?: string
      message?: string
      [key: string]: any
    }) => void
  }): Promise<Record<string, Record<string, { reason: string; score: number }>>> => {
    const requestPayload = {
      'General Goal': data.goal,
      stepTask: {
        Id: data.stepTask.Id || data.stepTask.id,
        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,
      },
      task_id: data.TaskID || '',
    }

    const rawResponse = await withRetry(
      () =>
        websocket.send(
          'agent_select_modify_init',
          requestPayload,
          undefined,
          data.onProgress,
        ),
      {
        maxRetries: 3,
        initialDelayMs: 2000,
        onRetry: (error, attempt, delay) => {
          console.warn(
            `[agentSelectModifyInit] 第${attempt}次重试，等待 ${delay}ms...`,
            error?.message,
          )
        },
      },
    )

    let response = this.extractResponse<any>(rawResponse)
    if (response?.scoreTable) {
      response = response.scoreTable
    }

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

    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[]
    stepTask?: {
      Id?: string
      StepName?: string
      TaskContent?: string
      InputObject_List?: string[]
      OutputObject?: string
    }
    TaskID?: string
    onProgress?: (progress: {
      status: string
      stage?: string
      message?: string
      [key: string]: any
    }) => void
  }): Promise<{
    aspectName: string
    agentScores: Record<string, { score: number; reason: string }>
  }> => {
    const rawResponse = await websocket.send(
      'agent_select_modify_add_aspect',
      {
        aspectList: data.aspectList,
        stepTask: data.stepTask,
        task_id: data.TaskID || '',
      },
      undefined,
      data.onProgress,
    )

    const response = this.extractResponse<any>(rawResponse)

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

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

    if (newAspectAgents) {
      for (const [agentName, scoreInfo] of Object.entries(newAspectAgents as Record<string, { Score?: number; score?: number; Reason?: string; reason?: string }>)) {
        agentScores[agentName] = {
          score: scoreInfo.Score || scoreInfo.score || 0,
          reason: scoreInfo.Reason || scoreInfo.reason || '',
        }
      }
    }

    return {
      aspectName: newAspect,
      agentScores,
    }
  }

  /**
   * 删除评估维度
   * @param taskId 任务ID
   * @param stepId 步骤ID（可选，不传则删除所有步骤中的该维度）
   * @param aspectName 要删除的维度名称
   * @returns 是否删除成功
   */
  agentSelectModifyDeleteAspect = async (
    taskId: string,
    aspectName: string,
    stepId?: string
  ): Promise<boolean> => {
    if (!websocket.connected) {
      throw new Error('WebSocket未连接')
    }

    try {
      const rawResponse = await websocket.send(
        'agent_select_modify_delete_aspect',
        {
          task_id: taskId,
          step_id: stepId || '',
          aspect_name: aspectName,
        },
        undefined,
        undefined
      )

      const response = this.extractResponse<{ status: string }>(rawResponse)
      return response?.status === 'success' || false
    } catch (error) {
      console.error('删除维度失败:', error)
      return false
    }
  }

  /**
   * 向正在执行的任务追加新步骤
   * @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,
        })),
      })),
    })

    const response = this.extractResponse<{ added_count: number }>(rawResponse)

    return response?.added_count || 0
  }

  /**
   * 保存任务分支数据
   * @param taskId 任务ID
   * @param branches 分支数据数组
   * @returns 是否保存成功
   */
  saveBranches = async (taskId: string, branches: any[]): Promise<boolean> => {
    if (!websocket.connected) {
      throw new Error('WebSocket未连接')
    }

    try {
      const rawResponse = await websocket.send('save_branches', {
        task_id: taskId,
        branches,
      })

      const response = this.extractResponse<{ status: string }>(rawResponse)
      return response?.status === 'success' || false
    } catch (error) {
      console.error('保存分支数据失败:', error)
      return false
    }
  }

  /**
   * 保存任务过程分支数据
   * @param TaskID 大任务ID（数据库主键）
   * @param branches 任务过程分支数据（Map结构转对象）
   * @returns 是否保存成功
   */
  saveTaskProcessBranches = async (
    TaskID: string,
    branches: Record<string, Record<string, any[]>>,
  ): Promise<boolean> => {
    if (!websocket.connected) {
      throw new Error('WebSocket未连接')
    }

    try {
      const rawResponse = await websocket.send('save_task_process_branches', {
        task_id: TaskID,
        branches,
      })

      const response = this.extractResponse<{ status: string }>(rawResponse)
      return response?.status === 'success' || false
    } catch (error) {
      console.error('保存任务过程分支数据失败:', error)
      return false
    }
  }

  /**
   * 更新任务大纲数据
   * @param taskId 任务ID
   * @param taskOutline 完整的大纲数据
   * @returns 是否更新成功
   */
  updateTaskOutline = async (taskId: string, taskOutline: any): Promise<boolean> => {
    if (!websocket.connected) {
      throw new Error('WebSocket未连接')
    }

    try {
      const rawResponse = await websocket.send('save_task_outline', {
        task_id: taskId,
        task_outline: taskOutline,
      })

      const response = this.extractResponse<{ status: string }>(rawResponse)
      return response?.status === 'success' || false
    } catch (error) {
      console.error('更新任务大纲失败:', error)
      return false
    }
  }

  /**
   * 更新指定步骤的 assigned_agents
   * @param params 参数对象
   * @returns 是否更新成功
   */
  updateAssignedAgents = async (params: {
    task_id: string // 大任务ID（数据库主键）
    step_id: string // 步骤级ID（小任务UUID）
    agents: string[] // 选中的 agent 列表
    confirmed_groups?: string[][] // 可选：确认的 agent 组合列表
    agent_combinations?: Record<string, { process: any; brief: any }> // 可选：agent 组合的 TaskProcess 数据
  }): Promise<boolean> => {
    if (!websocket.connected) {
      throw new Error('WebSocket未连接')
    }

    try {
      const rawResponse = await websocket.send('update_assigned_agents', {
        task_id: params.task_id,
        step_id: params.step_id,
        agents: params.agents,
        confirmed_groups: params.confirmed_groups,
        agent_combinations: params.agent_combinations,
      })

      const response = this.extractResponse<{ status: string; error?: string }>(rawResponse)
      if (response?.status === 'success') {
        console.log('更新 assigned_agents 成功:', params)
        return true
      }
      console.warn('更新 assigned_agents 失败:', response?.error)
      return false
    } catch (error) {
      console.error('更新 assigned_agents 失败:', error)
      return false
    }
  }
}

export default new Api()