simstudioai--sim/apps/sim/lib/copilot/tools/handlers/workflow/queries.ts

import { toError } from '@sim/utils/errors'
import { mergeSubblockStateWithValues } from '@sim/workflow-persistence/subblocks'
import type { ExecutionContext, ToolCallResult } from '@/lib/copilot/request/types'
import { formatNormalizedWorkflowForCopilot } from '@/lib/copilot/tools/shared/workflow-utils'
import { mcpService } from '@/lib/mcp/service'
import { listWorkspaceFiles } from '@/lib/uploads/contexts/workspace'
import { getEffectiveBlockOutputPaths } from '@/lib/workflows/blocks/block-outputs'
import { BlockPathCalculator } from '@/lib/workflows/blocks/block-path-calculator'
import { getBlockReferenceTags } from '@/lib/workflows/blocks/block-reference-tags'
import { listCustomTools } from '@/lib/workflows/custom-tools/operations'
import {
  loadDeployedWorkflowState,
  loadWorkflowFromNormalizedTables,
} from '@/lib/workflows/persistence/utils'
import { resolveTriggerRunOptions, toPublicRunOption } from '@/lib/workflows/triggers/run-options'
import { hasTriggerCapability } from '@/lib/workflows/triggers/trigger-utils'
import { getWorkflowById } from '@/lib/workflows/utils'
import { listUserWorkspaces } from '@/lib/workspaces/utils'
import { getBlock } from '@/blocks/registry'
import { normalizeName } from '@/executor/constants'
import type { Loop, Parallel } from '@/stores/workflows/workflow/types'
import { ensureWorkflowAccess } from '../access'
import type {
  GetBlockOutputsParams,
  GetBlockUpstreamReferencesParams,
  GetDeployedWorkflowStateParams,
  GetWorkflowDataParams,
  GetWorkflowRunOptionsParams,
} from '../param-types'

export async function executeListUserWorkspaces(
  context: ExecutionContext
): Promise<ToolCallResult> {
  try {
    const workspaces = await listUserWorkspaces(context.userId)

    return { success: true, output: { workspaces } }
  } catch (error) {
    return { success: false, error: toError(error).message }
  }
}

export async function executeGetWorkflowRunOptions(
  params: GetWorkflowRunOptionsParams,
  context: ExecutionContext
): Promise<ToolCallResult> {
  try {
    const workflowId = params.workflowId || context.workflowId
    if (!workflowId) {
      return { success: false, error: 'workflowId is required' }
    }

    await ensureWorkflowAccess(workflowId, context.userId)

    const normalized = await loadWorkflowFromNormalizedTables(workflowId)
    if (!normalized) {
      return { success: false, error: `Workflow ${workflowId} has no saved state` }
    }

    const merged = mergeSubblockStateWithValues(normalized.blocks)
    const options = resolveTriggerRunOptions(merged, normalized.edges)

    if (options.length === 0) {
      return {
        success: true,
        output: {
          workflowId,
          default: null,
          triggers: [],
          message:
            'No runnable trigger blocks found. Add a Start/API/Input/Chat trigger or an external (webhook/integration) trigger before running.',
        },
      }
    }

    const guidanceFor = (kind: string): string => {
      switch (kind) {
        case 'fields':
          return 'Build workflow_input matching inputSchema. Copy mockPayload only if you have no better values.'
        case 'event_payload':
          return 'Construct an event payload matching inputSchema, or run with useMockPayload: true if you cannot build one.'
        case 'chat':
          return 'Provide workflow_input shaped like { "input": "<message>" }.'
        default:
          return 'No input required.'
      }
    }

    const triggers = options.map((option) => {
      const pub = toPublicRunOption(option)
      const callExample =
        pub.inputKind === 'none'
          ? { triggerBlockId: pub.triggerBlockId }
          : { triggerBlockId: pub.triggerBlockId, workflow_input: pub.mockPayload }
      return { ...pub, guidance: guidanceFor(pub.inputKind), callExample }
    })

    const defaultOption = options.find((option) => option.isDefault)

    return {
      success: true,
      output: {
        workflowId,
        default: defaultOption
          ? {
              triggerBlockId: defaultOption.triggerBlockId,
              reason: `Highest-priority trigger (${defaultOption.blockName})`,
            }
          : null,
        triggers,
      },
    }
  } catch (error) {
    return { success: false, error: toError(error).message }
  }
}

export async function executeGetWorkflowData(
  params: GetWorkflowDataParams,
  context: ExecutionContext
): Promise<ToolCallResult> {
  try {
    const workflowId = params.workflowId || context.workflowId
    const dataType = params.data_type || params.dataType || ''
    if (!workflowId) {
      return { success: false, error: 'workflowId is required' }
    }
    if (!dataType) {
      return { success: false, error: 'data_type is required' }
    }

    const { workflow: workflowRecord, workspaceId } = await ensureWorkflowAccess(
      workflowId,
      context.userId
    )

    if (dataType === 'global_variables') {
      const variablesRecord = (workflowRecord.variables as Record<string, unknown>) || {}
      const variables = Object.values(variablesRecord).map((v) => {
        const variable = v as Record<string, unknown> | null
        return {
          id: String(variable?.id || ''),
          name: String(variable?.name || ''),
          value: variable?.value,
        }
      })
      return { success: true, output: { variables } }
    }

    if (dataType === 'custom_tools') {
      if (!workspaceId) {
        return { success: false, error: 'workspaceId is required' }
      }
      const toolsRows = await listCustomTools({
        userId: context.userId,
        workspaceId,
      })

      const customToolsData = toolsRows.map((tool) => {
        const schema = tool.schema as Record<string, unknown> | null
        const fn = (schema?.function ?? {}) as Record<string, unknown>
        return {
          id: String(tool.id || ''),
          title: String(tool.title || ''),
          functionName: String(fn.name || ''),
          description: String(fn.description || ''),
          parameters: fn.parameters,
        }
      })

      return { success: true, output: { customTools: customToolsData } }
    }

    if (dataType === 'mcp_tools') {
      if (!workspaceId) {
        return { success: false, error: 'workspaceId is required' }
      }
      const tools = await mcpService.discoverTools(context.userId, workspaceId, false)
      const mcpTools = tools.map((tool) => ({
        name: String(tool.name || ''),
        serverId: String(tool.serverId || ''),
        serverName: String(tool.serverName || ''),
        description: String(tool.description || ''),
        inputSchema: tool.inputSchema,
      }))
      return { success: true, output: { mcpTools } }
    }

    if (dataType === 'files') {
      if (!workspaceId) {
        return { success: false, error: 'workspaceId is required' }
      }
      const files = await listWorkspaceFiles(workspaceId)
      const fileResults = files.map((file) => ({
        id: String(file.id || ''),
        name: String(file.name || ''),
        key: String(file.key || ''),
        path: String(file.path || ''),
        size: Number(file.size || 0),
        type: String(file.type || ''),
        uploadedAt: String(file.uploadedAt || ''),
      }))
      return { success: true, output: { files: fileResults } }
    }

    return { success: false, error: `Unknown data_type: ${dataType}` }
  } catch (error) {
    return { success: false, error: toError(error).message }
  }
}

export async function executeGetBlockOutputs(
  params: GetBlockOutputsParams,
  context: ExecutionContext
): Promise<ToolCallResult> {
  try {
    const workflowId = params.workflowId || context.workflowId
    if (!workflowId) {
      return { success: false, error: 'workflowId is required' }
    }
    await ensureWorkflowAccess(workflowId, context.userId)

    const normalized = await loadWorkflowFromNormalizedTables(workflowId)
    if (!normalized) {
      return { success: false, error: 'Workflow has no normalized data' }
    }

    const blocks = normalized.blocks || {}
    const loops = normalized.loops || {}
    const parallels = normalized.parallels || {}
    const blockIds =
      Array.isArray(params.blockIds) && params.blockIds.length > 0
        ? params.blockIds
        : Object.keys(blocks)

    const results: Array<{
      blockId: string
      blockName: string
      blockType: string
      outputs: string[]
      relativeOutputs?: string[]
      insideSubflowOutputs?: string[]
      outsideSubflowOutputs?: string[]
      relativeInsideSubflowOutputs?: string[]
      relativeOutsideSubflowOutputs?: string[]
      triggerMode?: boolean
    }> = []

    for (const blockId of blockIds) {
      const block = blocks[blockId]
      if (!block?.type) continue
      const blockName = block.name || block.type

      if (block.type === 'loop' || block.type === 'parallel') {
        const insidePaths = getSubflowInsidePaths(block.type, blockId, loops, parallels)
        results.push({
          blockId,
          blockName,
          blockType: block.type,
          outputs: [],
          relativeOutputs: [],
          insideSubflowOutputs: formatOutputsForDisplay(insidePaths, blockName),
          outsideSubflowOutputs: formatOutputsForDisplay(['results'], blockName),
          relativeInsideSubflowOutputs: insidePaths,
          relativeOutsideSubflowOutputs: ['results'],
          triggerMode: block.triggerMode,
        })
        continue
      }

      const blockConfig = getBlock(block.type)
      const isTriggerCapable = blockConfig ? hasTriggerCapability(blockConfig) : false
      const triggerMode = Boolean(block.triggerMode && isTriggerCapable)
      const outputs = getEffectiveBlockOutputPaths(block.type, block.subBlocks, {
        triggerMode,
        preferToolOutputs: !triggerMode,
      })
      results.push({
        blockId,
        blockName,
        blockType: block.type,
        outputs: formatOutputsForDisplay(outputs, blockName),
        relativeOutputs: outputs,
        triggerMode: block.triggerMode,
      })
    }

    const variables = await getWorkflowVariablesForTool(workflowId)

    const payload = { blocks: results, variables }
    return { success: true, output: payload }
  } catch (error) {
    return { success: false, error: toError(error).message }
  }
}

export async function executeGetBlockUpstreamReferences(
  params: GetBlockUpstreamReferencesParams,
  context: ExecutionContext
): Promise<ToolCallResult> {
  try {
    const workflowId = params.workflowId || context.workflowId
    if (!workflowId) {
      return { success: false, error: 'workflowId is required' }
    }
    if (!Array.isArray(params.blockIds) || params.blockIds.length === 0) {
      return { success: false, error: 'blockIds array is required' }
    }
    await ensureWorkflowAccess(workflowId, context.userId)

    const normalized = await loadWorkflowFromNormalizedTables(workflowId)
    if (!normalized) {
      return { success: false, error: 'Workflow has no normalized data' }
    }

    const blocks = normalized.blocks || {}
    const edges = normalized.edges || []
    const loops = normalized.loops || {}
    const parallels = normalized.parallels || {}

    const graphEdges = edges.map((edge) => ({ source: edge.source, target: edge.target }))
    const variableOutputs = await getWorkflowVariablesForTool(workflowId)

    interface AccessibleBlockEntry {
      blockId: string
      blockName: string
      blockType: string
      outputs: string[]
      triggerMode?: boolean
      accessContext?: 'inside' | 'outside'
    }

    interface UpstreamReferenceResult {
      blockId: string
      blockName: string
      blockType: string
      accessibleBlocks: AccessibleBlockEntry[]
      insideSubflows: Array<{ blockId: string; blockName: string; blockType: string }>
      variables: Array<{ id: string; name: string; type: string; tag: string }>
    }

    const results: UpstreamReferenceResult[] = []

    for (const blockId of params.blockIds) {
      const targetBlock = blocks[blockId]
      if (!targetBlock) continue

      const insideSubflows: Array<{ blockId: string; blockName: string; blockType: string }> = []
      const containingLoopIds = new Set<string>()
      const containingParallelIds = new Set<string>()

      Object.values(loops).forEach((loop) => {
        if (loop?.nodes?.includes(blockId)) {
          containingLoopIds.add(loop.id)
          const loopBlock = blocks[loop.id]
          if (loopBlock) {
            insideSubflows.push({
              blockId: loop.id,
              blockName: loopBlock.name || loopBlock.type,
              blockType: 'loop',
            })
          }
        }
      })

      Object.values(parallels).forEach((parallel) => {
        if (parallel?.nodes?.includes(blockId)) {
          containingParallelIds.add(parallel.id)
          const parallelBlock = blocks[parallel.id]
          if (parallelBlock) {
            insideSubflows.push({
              blockId: parallel.id,
              blockName: parallelBlock.name || parallelBlock.type,
              blockType: 'parallel',
            })
          }
        }
      })

      const ancestorIds = BlockPathCalculator.findAllPathNodes(graphEdges, blockId)
      const accessibleIds = new Set<string>(ancestorIds)
      accessibleIds.add(blockId)

      containingLoopIds.forEach((loopId) => accessibleIds.add(loopId))

      containingParallelIds.forEach((parallelId) => accessibleIds.add(parallelId))

      const accessibleBlocks: AccessibleBlockEntry[] = []

      for (const accessibleBlockId of accessibleIds) {
        const block = blocks[accessibleBlockId]
        if (!block?.type) continue
        const canSelfReference = block.type === 'approval' || block.type === 'human_in_the_loop'
        if (accessibleBlockId === blockId && !canSelfReference) continue

        const blockName = block.name || block.type
        let accessContext: 'inside' | 'outside' | undefined

        let formattedOutputs: string[]
        if (block.type === 'loop' || block.type === 'parallel') {
          const isInside =
            (block.type === 'loop' && containingLoopIds.has(accessibleBlockId)) ||
            (block.type === 'parallel' && containingParallelIds.has(accessibleBlockId))
          accessContext = isInside ? 'inside' : 'outside'
          const outputPaths = isInside
            ? getSubflowInsidePaths(block.type, accessibleBlockId, loops, parallels)
            : ['results']
          formattedOutputs = formatOutputsForDisplay(outputPaths, blockName)
        } else {
          formattedOutputs = getBlockReferenceTags({
            block: {
              id: accessibleBlockId,
              type: block.type,
              name: block.name,
              triggerMode: block.triggerMode,
              subBlocks: block.subBlocks,
            },
            currentBlockId: blockId,
          })
        }
        const entry: AccessibleBlockEntry = {
          blockId: accessibleBlockId,
          blockName,
          blockType: block.type,
          outputs: formattedOutputs,
          ...(block.triggerMode ? { triggerMode: true } : {}),
          ...(accessContext ? { accessContext } : {}),
        }
        accessibleBlocks.push(entry)
      }

      results.push({
        blockId,
        blockName: targetBlock.name || targetBlock.type,
        blockType: targetBlock.type,
        accessibleBlocks,
        insideSubflows,
        variables: variableOutputs,
      })
    }

    const payload = { results }
    return { success: true, output: payload }
  } catch (error) {
    return { success: false, error: toError(error).message }
  }
}

async function getWorkflowVariablesForTool(
  workflowId: string
): Promise<Array<{ id: string; name: string; type: string; tag: string }>> {
  const workflowRecord = await getWorkflowById(workflowId)

  const variablesRecord = (workflowRecord?.variables as Record<string, unknown>) || {}
  return Object.values(variablesRecord)
    .filter((v): v is Record<string, unknown> => {
      if (!v || typeof v !== 'object') return false
      const variable = v as Record<string, unknown>
      return !!variable.name && String(variable.name).trim() !== ''
    })
    .map((v) => ({
      id: String(v.id || ''),
      name: String(v.name || ''),
      type: String(v.type || 'plain'),
      tag: `variable.${normalizeName(String(v.name || ''))}`,
    }))
}

function getSubflowInsidePaths(
  blockType: 'loop' | 'parallel',
  blockId: string,
  loops: Record<string, Loop>,
  parallels: Record<string, Parallel>
): string[] {
  const paths = ['index']
  if (blockType === 'loop') {
    const loopType = loops[blockId]?.loopType || 'for'
    if (loopType === 'forEach') {
      paths.push('currentItem', 'items')
    }
  } else {
    const parallelType = parallels[blockId]?.parallelType || 'count'
    if (parallelType === 'collection') {
      paths.push('currentItem', 'items')
    }
  }
  return paths
}

function formatOutputsForDisplay(paths: string[], blockName: string): string[] {
  const normalizedName = normalizeName(blockName)
  return paths.map((path) => `${normalizedName}.${path}`)
}

export async function executeGetDeployedWorkflowState(
  params: GetDeployedWorkflowStateParams,
  context: ExecutionContext
): Promise<ToolCallResult> {
  try {
    const workflowId = params.workflowId || context.workflowId
    if (!workflowId) {
      return { success: false, error: 'workflowId is required' }
    }

    const { workflow: workflowRecord } = await ensureWorkflowAccess(workflowId, context.userId)

    try {
      const deployedState = await loadDeployedWorkflowState(workflowId)
      const formatted = formatNormalizedWorkflowForCopilot({
        blocks: deployedState.blocks,
        edges: deployedState.edges,
        loops: deployedState.loops as Record<string, Loop>,
        parallels: deployedState.parallels as Record<string, Parallel>,
      })

      return {
        success: true,
        output: {
          workflowId,
          workflowName: workflowRecord.name || '',
          isDeployed: true,
          deploymentVersionId: deployedState.deploymentVersionId,
          deployedState: formatted,
        },
      }
    } catch {
      return {
        success: true,
        output: {
          workflowId,
          workflowName: workflowRecord.name || '',
          isDeployed: false,
          message: 'Workflow has not been deployed yet.',
        },
      }
    }
  } catch (error) {
    return { success: false, error: toError(error).message }
  }
}