import { AUTOMATION_DELAY_ENTRYPOINT_PREFIX } from "@automate.ax/api-contract/runtime"
import { automationInvokedTriggerDefinition } from "@automate.ax/catalog/triggers/core-invocation"
import * as z from "zod"
import { defineAction } from "../../automation/actions"
import { getCurrentHookScopePath } from "../../automation/runtime"
import {
  correlate,
  getCurrentContextPrerequisite,
  getCurrentSignalPrerequisites,
  group,
  race,
  scope,
  withoutSignalPrerequisites,
  withContextPrerequisite,
  withPrerequisites,
  type Duration,
} from "../../automation/signal-operators"
import {
  inheritSignalCoordination,
  isSignal,
  transform,
  type InferSignal,
  type InheritSignalCoordination,
  type Signal,
} from "../../automation/signal-protocol"
import { createSubscription } from "../../automation/subscription"

const DELAY_ACTION_OUTPUT_SCHEMA = z.object({ key: z.string() })

const scheduleDelayAction = defineAction("Delay")
  .input(
    z.object({
      entrypoint: z.string().startsWith(AUTOMATION_DELAY_ENTRYPOINT_PREFIX),
      runIn: z.union([z.string().min(1), z.number().nonnegative()]),
    }),
  )
  .output(DELAY_ACTION_OUTPUT_SCHEMA)
  .retry({ replaySafety: "safe" })
  .handler(async ({ input, runtime }) => {
    const key = `${runtime.contextId}:${input.entrypoint}`
    await runtime.invokeAutomation({
      automation: runtime.automationId,
      entrypoint: input.entrypoint,
      payload: key,
      runIn: input.runIn,
    })
    return { key }
  })

/**
 * Emits null in a correlated child context after a durable relative delay.
 *
 * @param duration - Compact duration string, milliseconds, or compatible
 *   signal.
 */
export function delay(duration: Duration): Signal<null>

/**
 * Re-emits a signal value in a correlated child context after a durable delay.
 *
 * @param duration - Compact duration string, milliseconds, or compatible
 *   signal.
 * @param value - Signal whose value starts the delay and is emitted afterward.
 */
export function delay<TSignal extends Signal>(
  duration: Duration,
  value: TSignal,
): InheritSignalCoordination<TSignal, InferSignal<TSignal>>

/**
 * Traverses an automation section whose durable operations wait for a delay.
 *
 * @param duration - Compact duration string, milliseconds, or compatible
 *   signal.
 * @param section - Synchronous automation section to traverse.
 */
export function delay<TResult extends Signal | void>(
  duration: Duration,
  section: () => TResult,
): TResult

export function delay<T>(
  duration: Duration,
  valueOrSection?: Signal<T> | (() => Signal | void),
): unknown {
  const prerequisites = getCurrentSignalPrerequisites()
  const contextPrerequisite = getCurrentContextPrerequisite()
  return group({ name: "Delay", presentation: "hidden" }, () =>
    withoutSignalPrerequisites(() => {
      const value = isSignal(valueOrSection) ? valueOrSection : undefined
      const entrypoint = `${AUTOMATION_DELAY_ENTRYPOINT_PREFIX}${getCurrentHookScopePath().join(".")}`
      // The delivery trigger must reserve its hook before the scheduling action
      // so planning and later execution traverse the same durable identities.
      const elapsed = createSubscription<string>(
        automationInvokedTriggerDefinition,
        { entrypoint },
        undefined,
        { inferActionBoundary: false },
      )
      const schedule = () =>
        scheduleDelayAction({
          entrypoint,
          runIn: duration,
        })
      const scheduleWithContext = () =>
        contextPrerequisite
          ? withContextPrerequisite(contextPrerequisite, schedule)
          : schedule()
      const schedulePrerequisite =
        value && prerequisites
          ? transform([value, prerequisites], () => null)
          : (value ?? prerequisites)
      const scheduled = schedulePrerequisite
        ? scope(() =>
            withPrerequisites(schedulePrerequisite, scheduleWithContext),
          )
        : scheduleWithContext()
      const source = value
        ? transform(
            [value, scheduled],
            (delayedValue, { key }): { key: string; value: T | null } => ({
              key,
              value: delayedValue,
            }),
          )
        : transform(
            [scheduled],
            ({ key }): { key: string; value: T | null } => ({
              key,
              value: null,
            }),
          )
      const delayed = transform(
        [
          correlate([
            source.keyBy(({ key }) => key),
            elapsed.keyBy((key) => key),
          ]),
          source,
        ],
        (_resumed, { value: delayedValue }) => delayedValue,
      )
      return valueOrSection && !isSignal(valueOrSection)
        ? scope(() => withPrerequisites(delayed, valueOrSection))
        : value
          ? inheritSignalCoordination(value, delayed)
          : delayed
    }),
  )
}

/**
 * Fails unless a signal emits or fails before a durable relative deadline.
 *
 * The deadline can span execution jobs. It does not interrupt an action handler
 * that is already running.
 *
 * @param signal - Signal competing against the deadline.
 * @param duration - Compact duration string, milliseconds, or compatible
 *   signal.
 */
export function timeout<TSignal extends Signal>(
  signal: TSignal,
  duration: Duration,
): InheritSignalCoordination<TSignal, InferSignal<TSignal>> {
  return inheritSignalCoordination(
    signal,
    withContextPrerequisite(signal, () =>
      race([
        signal,
        delay(duration).transform(() => {
          const error = new Error("Signal timed out.")
          error.name = "TimeoutError"
          throw error
        }),
      ]),
    ),
  )
}