import { Client } from '@botpress/client'
import { Cognitive, Models, type BotpressClientLike } from '@botpress/cognitive'
import { z } from '@bpinternal/zui'

import { clamp, isEqual, isPlainObject, omit } from 'lodash-es'
import ms from 'ms'

import { ulid } from 'ulid'
import { createJoinedAbortController } from './abort-signal.js'
import { Chat } from './chat.js'
import { Context, Iteration } from './context.js'
import {
  AssignmentError,
  CodeExecutionError,
  CognitiveError,
  InvalidCodeError,
  LoopExceededError,
  Signals,
  SnapshotSignal,
  ThinkSignal,
  VMSignal,
} from './errors.js'
import { parseExit } from './exit-parser.js'
import { Exit, ExitResult } from './exit.js'
import { ValueOrGetter } from './getter.js'

import { type ObjectInstance } from './objects.js'

import { ErrorExecutionResult, ExecutionResult, PartialExecutionResult, SuccessExecutionResult } from './result.js'
import { Snapshot } from './snapshots.js'
import { cleanStackTrace } from './stack-traces.js'
import { type Tool } from './tool.js'

import { truncateWrappedContent } from './truncator.js'
import { Trace } from './types.js'

import { init, stripInvalidIdentifiers } from './utils.js'
import { runAsyncFunction } from './vm.js'

const getErrorMessage = (err: unknown) => (err instanceof Error ? err.message : JSON.stringify(err))

const SLOW_TOOL_WARNING = ms('15s')

const RESPONSE_LENGTH_BUFFER = {
  MIN_TOKENS: 1_000,
  MAX_TOKENS: 16_000,
  PERCENTAGE: 0.1,
} as const

const getModelOutputLimit = (inputLength: number) =>
  clamp(
    RESPONSE_LENGTH_BUFFER.PERCENTAGE * inputLength,
    RESPONSE_LENGTH_BUFFER.MIN_TOKENS,
    RESPONSE_LENGTH_BUFFER.MAX_TOKENS
  )

export type ExecutionHooks = {
  /**
   * NON-BLOCKING HOOK
   *   This hook will not block the execution of the iteration.
   * NON-MUTATION HOOK
   *   This hook can't mutate traces.
   *
   * This hook is called for each trace that is generated during the iteration.
   * It is useful for logging, debugging, or monitoring the execution of the iteration.
   */
  onTrace?: (event: { trace: Trace; iteration: number }) => void

  /**
   * BLOCKING HOOK
   *   This hook will block the execution of the iteration until it resolves.
   *
   * This hook will be called before each iteration starts, regardless of the status.
   * This is useful for logging or dynamically change model arguments
   */
  onIterationStart?: (
    iteration: Iteration,
    controller: AbortController,
    context: Context
  ) => Promise<void | Partial<Iteration>> | void | Partial<Iteration>

  /**
   * BLOCKING HOOK
   *   This hook will block the execution of the iteration until it resolves.
   * NON-MUTATION HOOK
   *   This hook can't mutate the result or status of the iteration.
   *
   * This hook will be called after each iteration ends, regardless of the status.
   * This is useful for logging, cleanup or to prevent or delay the next iteration from starting.
   */
  onIterationEnd?: (iteration: Iteration, controller: AbortController) => Promise<void> | void

  /**
   * BLOCKING HOOK
   *   This hook will block the execution of the iteration until it resolves.
   * NON-MUTATION HOOK
   *   This hook can't mutate the exit result.
   *
   * This hook is called when an exit is reached in the iteration.
   * It is useful for logging, sending notifications, or performing actions based on the exit.
   * It can also be used to throw an error and preventing the exit from being successful.
   * If this hook throws an error, the execution will keep iterating with the error as context.
   */
  onExit?: <T = unknown>(result: ExitResult<T>) => Promise<void> | void

  /**
   * BLOCKING HOOK
   *   This hook will block the execution of the iteration until it resolves.
   * MUTATION HOOK
   *   This hook can mutate the code to run in the iteration.
   *
   * This hook is called after the LLM generates the code for the iteration, but before it is executed.
   * It is useful for modifying the code to run, or for guarding against certain code patterns.
   */
  onBeforeExecution?: (iteration: Iteration, controller: AbortController) => Promise<{ code?: string } | void>

  /**
   * BLOCKING HOOK
   *   This hook will block the execution of the iteration until it resolves.
   * MUTATION HOOK
   *   This hook can mutate the input to the tool.
   *
   * This hook is called before any tool is executed.
   * It is useful for modifying the input to a tool or prevent a tool from executing.
   */
  onBeforeTool?: (event: {
    iteration: Iteration
    tool: Tool
    input: any
    controller: AbortController
    toolCallId: string
    object?: string
  }) => Promise<{ input?: any } | void>

  /**
   * BLOCKING HOOK
   *   This hook will block the execution of the iteration until it resolves.
   * MUTATION HOOK
   *   This hook can mutate the output of the tool.
   *
   * This hook is called after a tool is executed.
   * It is useful for modifying the output of a tool or for logging purposes.
   */
  onAfterTool?: (event: {
    iteration: Iteration
    tool: Tool
    input: any
    output: any
    controller: AbortController
    toolCallId: string
    object?: string
  }) => Promise<{ output?: any } | void>
}

type Options = Partial<Pick<Context, 'loop' | 'timeout'>>

export type ExecutionProps = {
  /**
   * If provided, the execution will be run in "Chat Mode".
   * In this mode, the execution will be able to send messages to the chat and will also have access to a chat transcript.
   * The execution can still end with a custom Exit, but a special ListenExit will be added to give back the chat control to the user.
   *
   * If `chat` is not provided, the execution will run in "Worker Mode", where it will not have access to a chat transcript.
   * In Worker Mode, the execution will iterate until it reaches an Exit or runs out of iterations.
   */
  chat?: Chat

  /**
   * Instructions for the LLM to follow.
   * This is a system prompt that will be used to guide the LLM's behavior.
   * It can be a simple string or a function that returns a string based on the current context (dynamic instructions).
   * Dynamic instructions are evaluated at the start of each iteration, allowing for context-aware instructions.
   */
  instructions?: ValueOrGetter<string, Context>

  /**
   * Objects available in the context.
   * Objects are useful to scope related tools together and to provide data to the VM.
   * Objects can contain "properties" that can be read and written to, as well as "tools" that can be executed.
   * Properties are type-safe and can be defined using a Zui schema.
   * Properties can be marked as read-only or writable.
   * The sandbox will ensure that properties are only modified if they are writable, and will also ensure that the values are valid according to the schema.
   *
   * Objects can be a static array of objects or a function that returns an array based on the current context.
   * Dynamic objects are evaluated at the start of each iteration, allowing for context-aware objects.
   *
   * Example:
   * An object "user" with properties "name" (string, writable) and "age" (number, read-only) will allow the LLM to see both properties and their values,
   * and executing the code `user.name = 'John'` will succeed, while `user.age = 30` will throw an error as the property is read-only.
   * Similarly, `user.name = 123` will throw an error as the value is not a string.
   */
  objects?: ValueOrGetter<ObjectInstance[], Context>

  /**
   * Tools available in the context.
   * Tools are functions that can be executed by the LLM to perform actions or retrieve data.
   * Tools can be defined with input and output schemas using Zui, and can be scoped to an object.
   * Tools can also have aliases, which are alternative names for the tool that can be used to call it.
   *
   * Tools can be a static array of tools or a function that returns an array based on the current context.
   * Dynamic tools are evaluated at the start of each iteration, allowing for context-aware tools.
   */
  tools?: ValueOrGetter<Tool[], Context>

  /**
   * Exits available in the context.
   * Exits define the possible endpoints for the execution. Every execution will either end with an exit, or run out of iterations.
   *
   * When `chat` is provided, the built-in "ListenExit" is automatically added.
   * When `exits` is not provided, the built-in "DefaultExit" is automatically added.
   *
   * Each exit has a name and can have aliases, which are alternative names for the exit that can be used to call it.
   * Exits can also have a Zui schema to validate the return value when the exit is reached.
   *
   * Exits can be a static array of exits or a function that returns an array based on the current context.
   * Dynamic exits are evaluated at the start of each iteration, allowing for context-aware exits.
   */
  exits?: ValueOrGetter<Exit[], Context>
  options?: Options

  /**
   * An instance of a Botpress Client, or an instance of Cognitive Client (@botpress/cognitive).
   * This is used to generate content using the LLM and to access the Botpress API.
   * If not provided, a default client will be created using environment variables.
   */
  client?: Cognitive | BotpressClientLike

  /**
   * When provided, the execution will immediately stop when the signal is aborted.
   * This will stop the LLM generation, as well as kill the VM sandbox execution.
   * Aborted iterations will end with IterationStatuses.Aborted and the execution will be marked as failed.
   */
  signal?: AbortSignal

  /**
   * A snapshot is a saved state of the execution context.
   * It can be used to resume the execution of a context at a later time.
   * This is useful for long-running executions that may need to be paused and resumed later.
   * The snapshot MUST be settled, which means it has to be resolved or rejected.
   * Providing an unsettled snapshot will throw an error.
   */
  snapshot?: Snapshot

  /**
   * The model to use for the LLM.
   * This can be a static model name or a function that returns a model name based on the current context.
   */
  model?: ValueOrGetter<Models | Models[], Context>

  /**
   * The temperature to use for the LLM.
   * This can be a static temperature or a function that returns a temperature based on the current context.
   * The temperature must be between 0 and 2.
   * If the temperature is outside this range, it will be clamped to the nearest valid value.
   * If no temperature is provided, the default temperature of 0.7 will be used.
   */
  temperature?: ValueOrGetter<number, Context>

  /**
   * The reasoning effort to use for models that support reasoning.
   * - "none": Disable reasoning (for models with optional reasoning)
   * - "low" | "medium" | "high": Fixed reasoning effort levels
   * - "dynamic": Let the provider automatically determine the reasoning effort
   * If not provided, the model will not use reasoning for models with optional reasoning.
   */
  reasoningEffort?: ValueOrGetter<'low' | 'medium' | 'high' | 'dynamic' | 'none', Context>
} & ExecutionHooks

export const executeContext = async (props: ExecutionProps): Promise<ExecutionResult> => {
  await init()
  const result = await _executeContext(props)
  try {
    result.context.chat?.onExecutionDone?.(result)
  } catch {}
  return result
}

export const _executeContext = async (props: ExecutionProps): Promise<ExecutionResult> => {
  const controller = createJoinedAbortController([props.signal])
  const { onIterationStart, onIterationEnd, onTrace, onExit, onBeforeExecution, onAfterTool, onBeforeTool } = props

  const client = props.client ?? new Client()

  const cognitive = Cognitive.isCognitiveClient(client) ? client : new Cognitive({ client, __experimental_beta: true })
  const cleanups: (() => void)[] = []

  const ctx = new Context({
    chat: props.chat,
    instructions: props.instructions,
    objects: props.objects,
    tools: props.tools,
    loop: props.options?.loop,
    timeout: props.options?.timeout,
    exits: props.exits,
    snapshot: props.snapshot,
    model: props.model,
    temperature: props.temperature,
    reasoningEffort: props.reasoningEffort,
  })

  try {
    while (true) {
      if (ctx.iterations.length >= ctx.loop) {
        return new ErrorExecutionResult(ctx, new LoopExceededError())
      }

      const iteration = await ctx.nextIteration()

      try {
        await executeOnIterationStartHook({
          iteration,
          ctx,
          onIterationStart,
          controller,
          onIterationEnd,
        })
      } catch (err) {
        if (err instanceof ThinkSignal) {
          continue
        }
      }

      if (controller.signal.aborted) {
        iteration.end({
          type: 'aborted',
          aborted: {
            reason: controller.signal.reason ?? 'The operation was aborted',
          },
        })

        return new ErrorExecutionResult(ctx, controller.signal.reason ?? 'The operation was aborted')
      }

      cleanups.push(
        iteration.traces.onPush((traces) => {
          for (const trace of traces) {
            onTrace?.({ trace, iteration: ctx.iterations.length })
          }
        })
      )

      try {
        await executeIteration({
          iteration,
          ctx,
          cognitive,
          controller,
          onExit,
          onBeforeExecution,
          onAfterTool,
          onBeforeTool,
        })
      } catch (err) {
        if (err instanceof CognitiveError) {
          // early return when problems with cognitive service
          return new ErrorExecutionResult(ctx, err)
        }

        iteration.end({
          type: 'execution_error',
          execution_error: {
            message: 'An unexpected error occurred: ' + getErrorMessage(err),
            stack: cleanStackTrace((err as Error).stack ?? 'No stack trace available'),
          },
        })
      }

      try {
        await onIterationEnd?.(iteration, controller)
      } catch (err) {
        console.error(err)
      }

      // Successful states
      if (iteration.status.type === 'exit_success') {
        const exitName = iteration.status.exit_success.exit_name

        return new SuccessExecutionResult(ctx, {
          exit: iteration.exits.find((x) => x.name === exitName)!,
          result: iteration.status.exit_success.return_value,
        })
      }

      if (iteration.status.type === 'callback_requested') {
        return new PartialExecutionResult(
          ctx,
          iteration.status.callback_requested.signal,
          Snapshot.fromSignal(iteration.status.callback_requested.signal)
        )
      }

      // Retryable errors
      if (
        iteration.status.type === 'thinking_requested' ||
        iteration.status.type === 'exit_error' ||
        iteration.status.type === 'execution_error' ||
        iteration.status.type === 'invalid_code_error'
      ) {
        continue
      }

      // Fatal errors
      return new ErrorExecutionResult(ctx, iteration.error ?? `Unknown error. Status: ${iteration.status.type}`)
    }
  } catch (error) {
    return new ErrorExecutionResult(ctx, error ?? 'Unknown error')
  } finally {
    for (const cleanup of cleanups) {
      try {
        cleanup()
      } catch {}
    }
  }
}

const executeIteration = async ({
  iteration,
  ctx,
  cognitive,
  controller,
  onExit,
  onBeforeExecution,
  onBeforeTool,
  onAfterTool,
}: {
  ctx: Context
  iteration: Iteration
  cognitive: Cognitive
  controller: AbortController
} & ExecutionHooks): Promise<void> => {
  let startedAt = Date.now()
  const traces = iteration.traces

  // When an array is provided, additional entries act as ordered fallbacks at the
  // cognitive layer. Token budget and stop-token computation only need the primary
  // model's details, so we resolve those off the first entry.
  const modelRef = Array.isArray(iteration.model) ? iteration.model[0]! : iteration.model
  const model = await cognitive.getModelDetails(modelRef).catch((thrown) => {
    throw new CognitiveError(`Failed to fetch model details for model "${modelRef}": ${getErrorMessage(thrown)}`)
  })
  const modelLimit = Math.max(model.input.maxTokens, 8_000)
  const responseLengthBuffer = getModelOutputLimit(modelLimit)

  const messages = truncateWrappedContent({
    messages: iteration.messages,
    tokenLimit: modelLimit - responseLengthBuffer,
    throwOnFailure: true,
  }).filter(
    (x) =>
      // Filter out empty messages, as they are not valid inputs for the LLM
      // This can happen when a message is truncated and the content is empty
      typeof x.content !== 'string' || x.content.trim().length > 0
  )
  iteration.messages = messages

  traces.push({
    type: 'llm_call_started',
    started_at: startedAt,
    ended_at: startedAt,
    model: model.ref,
  })

  const output = await cognitive
    .generateContent({
      signal: controller.signal,
      systemPrompt: messages.find((x) => x.role === 'system')?.content,
      // Validated upstream in Context (see context.ts isValidModel). Cast narrows
      // Models' `({} & string)` escape hatch down to cognitive's stricter InputModel.
      model: iteration.model as Required<Parameters<Cognitive['generateContent']>[0]>['model'],
      temperature: iteration.temperature,
      responseFormat: 'text',
      reasoningEffort: iteration.reasoningEffort,
      messages: messages.filter((x) => x.role !== 'system'),
      stopSequences: ctx.version.getStopTokens(),
    })
    .catch((thrown) => {
      throw new CognitiveError(`LLM generation failed: ${getErrorMessage(thrown)}`)
    })

  const out = typeof output.output.choices?.[0]?.content === 'string' ? output.output.choices[0].content : null

  if (!out) {
    throw new CognitiveError('LLM did not return any text output')
  }

  const assistantResponse = ctx.version.parseAssistantResponse(out)

  iteration.code = assistantResponse.code.trim()

  if (typeof onBeforeExecution === 'function') {
    try {
      const hookRes = await onBeforeExecution(iteration, controller)
      if (typeof hookRes?.code === 'string' && hookRes.code.trim().length > 0) {
        iteration.code = hookRes.code.trim()
      }
    } catch (err) {
      if (err instanceof ThinkSignal) {
        return iteration.end({
          type: 'thinking_requested',
          thinking_requested: {
            variables: err.context,
            reason: err.reason,
          },
        })
      }

      return iteration.end({
        type: 'execution_error',
        execution_error: {
          message: `Error in onBeforeExecution hook: ${getErrorMessage(err)}`,
          stack: cleanStackTrace((err as Error).stack ?? 'No stack trace available'),
        },
      })
    }
  }

  iteration.llm = {
    cached: output.meta.cached || false,
    ended_at: Date.now(),
    started_at: startedAt,
    status: 'success',
    tokens: output.meta.tokens.input + output.meta.tokens.output,
    spend: output.meta.cost.input + output.meta.cost.output,
    output: assistantResponse.raw,
    model: `${output.meta.model.integration}:${output.meta.model.model}`,
    usage: output.output.usage,
  }

  traces.push({
    type: 'llm_call_success',
    started_at: startedAt,
    ended_at: iteration.llm.ended_at,
    model: model.ref,
    code: iteration.code,
  })

  const vmContext = { ...stripInvalidIdentifiers(iteration.variables) }

  for (const obj of iteration.objects) {
    const internalValues: Record<string, string> = {}
    const instance: Record<string, any> = {}

    for (const { name, value, writable, type } of obj.properties ?? []) {
      internalValues[name] = value

      const initialValue = value
      const schema = (type ?? z.any()) as z.ZodType

      Object.defineProperty(instance, name, {
        enumerable: true,
        configurable: true,
        get() {
          return internalValues[name]
        },
        set(value) {
          if (isEqual(value, internalValues[name])) {
            return
          }

          if (!writable) {
            throw new AssignmentError(`Property ${obj.name}.${name} is read-only and cannot be modified`)
          }

          if (value === internalValues[name]) {
            return
          }

          const parsed = schema.safeParse(value)

          if (!parsed.success) {
            throw new AssignmentError(
              `Invalid value for Object property ${obj.name}.${name}: ${getErrorMessage(parsed.error)}`
            )
          }

          internalValues[name] = parsed.data

          traces.push({
            type: 'property',
            started_at: Date.now(),
            object: obj.name,
            property: name,
            value: parsed.data,
          })

          iteration.trackMutation({ object: obj.name, property: name, before: initialValue, after: parsed.data })
        },
      })
    }

    for (const tool of obj.tools ?? []) {
      instance[tool.name] = wrapTool({
        tool,
        traces,
        object: obj.name,
        iteration,
        beforeHook: onBeforeTool,
        afterHook: onAfterTool,
        controller,
      })
    }

    Object.preventExtensions(instance)
    Object.seal(instance)

    vmContext[obj.name] = instance
  }

  for (const tool of iteration.tools) {
    const wrapped = wrapTool({ tool, traces, iteration, beforeHook: onBeforeTool, afterHook: onAfterTool, controller })
    for (const key of [tool.name, ...(tool.aliases ?? [])]) {
      vmContext[key] = wrapped
    }
  }

  if (controller.signal.aborted) {
    traces.push({
      type: 'abort_signal',
      started_at: Date.now(),
      reason: 'The operation was aborted by user.',
    })
    // TODO: handle abortions better.. there's more than one path possible

    return iteration.end({
      type: 'aborted',
      aborted: {
        reason: controller.signal.reason ?? 'The operation was aborted',
      },
    })
  }

  type Result = Awaited<ReturnType<typeof runAsyncFunction>>

  startedAt = Date.now()
  const result: Result = await runAsyncFunction(
    vmContext,
    iteration.code,
    traces,
    controller.signal,
    ctx.timeout
  ).catch((err) => {
    return {
      success: false,
      error: err as Error,
      lines_executed: [],
      traces: [],
      variables: {},
    } satisfies Result
  })

  if (result.error && result.error instanceof InvalidCodeError) {
    return iteration.end({
      type: 'invalid_code_error',
      invalid_code_error: {
        message: result.error.message,
      },
    })
  }

  traces.push({
    type: 'code_execution',
    lines_executed: result.lines_executed ?? 0,
    started_at: startedAt,
    ended_at: Date.now(),
  })

  // Check for abort first - takes precedence over other errors
  if (controller.signal.aborted) {
    return iteration.end({
      type: 'aborted',
      aborted: {
        reason: controller.signal.reason ?? 'The operation was aborted',
      },
    })
  }

  if (result.error && result.error instanceof CodeExecutionError) {
    return iteration.end({
      type: 'execution_error',
      execution_error: {
        message: result.error.message,
        stack: cleanStackTrace(result.error.stacktrace ?? result.error.stack ?? 'No stack trace available'),
      },
    })
  }

  if (!result.success) {
    return iteration.end({
      type: 'execution_error',
      execution_error: {
        message: result?.error?.message ?? 'Unknown error occurred',
        stack: cleanStackTrace(result.error.stack ?? 'No stack trace available'),
      },
    })
  }

  if (result.signal instanceof ThinkSignal) {
    return iteration.end({
      type: 'thinking_requested',
      thinking_requested: {
        variables: result.signal.context,
        reason: result.signal.reason,
        metadata: result.signal.metadata,
      },
    })
  }

  if (result.signal instanceof SnapshotSignal) {
    return iteration.end({
      type: 'callback_requested',
      callback_requested: {
        signal: result.signal,
      },
    })
  }

  let returnValue: { action: string; value?: unknown } | null =
    result.success && result.return_value ? result.return_value : null

  const returnAction = returnValue?.action

  // Handle think action separately
  if (returnAction === 'think') {
    const variables = omit(returnValue ?? {}, 'action')
    if (isPlainObject(variables) && Object.keys(variables).length > 0) {
      return iteration.end({
        type: 'thinking_requested',
        thinking_requested: {
          variables,
          reason: 'Thinking requested',
        },
      })
    }

    return iteration.end({
      type: 'thinking_requested',
      thinking_requested: {
        reason: 'Thinking requested',
        variables: iteration.variables,
      },
    })
  }

  // Parse the exit using the extracted function
  const parsedExit = parseExit(returnValue, iteration.exits)

  if (!parsedExit.success) {
    return iteration.end({
      type: 'exit_error',
      exit_error: {
        exit: returnValue?.action ?? 'n/a',
        message: parsedExit.error,
        return_value: returnValue,
      },
    })
  }

  const returnExit = parsedExit.exit
  returnValue = { action: returnExit.name, value: parsedExit.value }

  try {
    await onExit?.({
      exit: returnExit,
      result: returnValue?.value,
    })
  } catch (err) {
    return iteration.end({
      type: 'exit_error',
      exit_error: {
        exit: returnExit.name,
        message: `Error executing exit ${returnExit.name}: ${getErrorMessage(err)}`,
        return_value: returnValue,
      },
    })
  }

  return iteration.end({
    type: 'exit_success',
    exit_success: {
      exit_name: returnExit.name,
      return_value: returnValue?.value,
    },
  })
}

type Props = {
  tool: Tool
  object?: string
  traces: Trace[]
  iteration: Iteration
  beforeHook?: ExecutionHooks['onBeforeTool']
  afterHook?: ExecutionHooks['onAfterTool']
  controller: AbortController
}

function wrapTool({ tool, traces, object, iteration, beforeHook, afterHook, controller }: Props) {
  const getToolInput = (input: any) => (tool.zInput as any).safeParse(input).data ?? input

  return function (input: any) {
    const toolCallId = `tcall_${ulid()}`

    const alertSlowTool = setTimeout(
      () =>
        traces.push({
          type: 'tool_slow',
          tool_name: tool.name,
          tool_call_id: toolCallId,
          started_at: Date.now(),
          input: getToolInput(input),
          object,
          duration: SLOW_TOOL_WARNING,
        }),
      SLOW_TOOL_WARNING
    )

    const cancelSlowTool = () => clearTimeout(alertSlowTool)

    const toolStart = Date.now()
    let output: any
    let error: any
    let success = true

    const handleSignals = (error: any) => {
      if (output === error) {
        return true
      }

      if (error instanceof SnapshotSignal) {
        error.toolCall = {
          name: tool.name,
          inputSchema: tool.input,
          outputSchema: tool.output,
          input,
        }
        error.message = Signals.serializeError(error)
      }

      if (error instanceof ThinkSignal) {
        traces.push({
          type: 'think_signal',
          started_at: Date.now(),
          line: 0,
          ended_at: Date.now(),
        })
        success = true
        output = error
        return true
      }

      return false
    }

    try {
      const withHooks = async (input: any): Promise<any> => {
        const beforeRes = await beforeHook?.({
          iteration,
          tool,
          input,
          controller,
          object,
          toolCallId,
        })

        if (typeof beforeRes?.input !== 'undefined') {
          input = beforeRes.input
        }

        let output = await tool.execute(input, {
          callId: toolCallId,
        })

        const afterRes = await afterHook?.({
          iteration,
          tool,
          input,
          output,
          controller,
          object,
          toolCallId,
        })

        if (typeof afterRes?.output !== 'undefined') {
          output = afterRes.output
        }

        return output
      }

      const result = withHooks(input)

      if (result instanceof Promise || ((result as any)?.then && (result as any)?.catch)) {
        return result
          .then((res: any) => {
            output = res
            success = true
            return res
          })
          .catch(async (err: any) => {
            if (!handleSignals(err)) {
              success = false
              error = err
            } else {
              const afterRes = await afterHook?.({
                iteration,
                tool,
                input,
                output,
                controller,
                object,
                toolCallId,
              })

              if (typeof afterRes?.output !== 'undefined') {
                output = afterRes.output
              }
            }

            // Important: we want to re-throw signals so that the VM can handle them
            throw err
          })
          .finally(() => {
            cancelSlowTool()
            traces.push({
              type: 'tool_call',
              tool_call_id: toolCallId,
              started_at: toolStart,
              ended_at: Date.now(),
              tool_name: tool.name,
              object,
              input: getToolInput(input),
              output,
              error,
              success,
            })
          })
      }

      success = true
      output = result
    } catch (err) {
      if (!handleSignals(err)) {
        success = false
        error = err
      }
    }

    cancelSlowTool()
    traces.push({
      type: 'tool_call',
      tool_call_id: toolCallId,
      started_at: toolStart,
      ended_at: Date.now(),
      tool_name: tool.name,
      object,
      input: getToolInput(input),
      output,
      error,
      success,
    })

    if (!success) {
      throw error
    }

    if (output instanceof VMSignal) {
      // Important: we want to re-throw signals so that the VM can handle them
      throw output
    }

    return output
  }
}

const executeOnIterationStartHook = async (props: {
  iteration: Iteration
  ctx: Context
  onIterationStart?: ExecutionHooks['onIterationStart']
  onIterationEnd?: ExecutionHooks['onIterationEnd']
  controller: AbortController
}) => {
  const { iteration, ctx, onIterationStart, controller, onIterationEnd } = props

  try {
    const hookRes = await onIterationStart?.(iteration, controller, ctx)
    if (hookRes) {
      Object.assign(iteration, hookRes)
    }
  } catch (err) {
    if (err instanceof ThinkSignal) {
      iteration.end({
        type: 'thinking_requested',
        thinking_requested: {
          variables: err.context,
          reason: err.reason,
        },
      })

      try {
        await onIterationEnd?.(iteration, controller)
      } catch (err) {
        console.error(err)
      }
    } else {
      iteration.end({
        type: 'execution_error',
        execution_error: {
          message: `Error in onIterationStart hook: ${getErrorMessage(err)}`,
          stack: cleanStackTrace((err as Error).stack ?? 'No stack trace available'),
        },
      })
    }

    throw err
  }
}
