import { CompleteBatchResult } from "@internal/run-engine";
import { SpanKind } from "@internal/tracing";
import { tryCatch } from "@trigger.dev/core/utils";
import { createJsonErrorObject, sanitizeError } from "@trigger.dev/core/v3";
import { RunId } from "@trigger.dev/core/v3/isomorphic";
import { BatchTaskRunStatus, Prisma, RuntimeEnvironmentType } from "@trigger.dev/database";
import { TriggerFailedTaskService } from "~/runEngine/services/triggerFailedTask.server";
import { $replica, prisma } from "~/db.server";
import { env } from "~/env.server";
import { findEnvironmentById, findEnvironmentFromRun } from "~/models/runtimeEnvironment.server";
import { AuthenticatedEnvironment } from "~/services/apiAuth.server";
import { logger } from "~/services/logger.server";
import { updateMetadataService } from "~/services/metadata/updateMetadataInstance.server";
import { reportInvocationUsage } from "~/services/platform.v3.server";
import { MetadataTooLargeError } from "~/utils/packets";
import { QueueSizeLimitExceededError } from "~/v3/services/common.server";
import { TriggerTaskService } from "~/v3/services/triggerTask.server";
import { tracer } from "~/v3/tracer.server";
import { createExceptionPropertiesFromError } from "./eventRepository/common.server";
import { getEventRepositoryForStore, recordRunDebugLog } from "./eventRepository/index.server";
import { roomFromFriendlyRunId, socketIo } from "./handleSocketIo.server";
import { engine } from "./runEngine.server";
import { runStore } from "./runStore.server";
import { publishChangeRecord } from "~/services/realtime/runChangeNotifierInstance.server";
import { PerformTaskRunAlertsService } from "./services/alerts/performTaskRunAlerts.server";
import { TaskRunErrorCodes } from "@trigger.dev/core/v3";

export function registerRunEngineEventBusHandlers() {
  engine.eventBus.on("runSucceeded", async ({ time, run, organization, environment }) => {
    const [taskRunError, taskRun] = await tryCatch(
      runStore.findRunOrThrow(
        {
          id: run.id,
        },
        {
          select: {
            id: true,
            friendlyId: true,
            traceId: true,
            spanId: true,
            parentSpanId: true,
            createdAt: true,
            completedAt: true,
            taskIdentifier: true,
            projectId: true,
            runtimeEnvironmentId: true,
            environmentType: true,
            isTest: true,
            organizationId: true,
            taskEventStore: true,
            // Piggyback the realtime run-changed publish on this existing read so the
            // per-env channel carries the membership keys (no separate query). No-op when
            // the native backend is disabled.
            runTags: true,
            batchId: true,
          },
        },
        $replica
      )
    );

    if (taskRunError) {
      logger.error("[runSucceeded] Failed to find task run", {
        error: taskRunError,
        runId: run.id,
      });
      return;
    }

    publishChangeRecord({
      runId: taskRun.id,
      envId: environment.id,
      tags: taskRun.runTags,
      batchId: taskRun.batchId,
      updatedAtMs: run.updatedAt.getTime(),
    });

    const eventRepository = await getEventRepositoryForStore(
      run.taskEventStore,
      taskRun.organizationId ?? organization.id
    );

    const [completeSuccessfulRunEventError] = await tryCatch(
      eventRepository.completeSuccessfulRunEvent({
        run: taskRun,
        endTime: time,
      })
    );

    if (completeSuccessfulRunEventError) {
      logger.error("[runSucceeded] Failed to complete successful run event", {
        error: completeSuccessfulRunEventError,
        runId: run.id,
      });
    }
  });

  // Handle alerts
  engine.eventBus.on("runFailed", async ({ time, run }) => {
    try {
      await PerformTaskRunAlertsService.enqueue(run.id);
    } catch (error) {
      logger.error("[runFailed] Failed to enqueue alerts", {
        error: error instanceof Error ? error.message : error,
        runId: run.id,
        spanId: run.spanId,
      });
    }
  });

  // Handle events
  engine.eventBus.on("runFailed", async ({ time, run, organization, environment }) => {
    const sanitizedError = sanitizeError(run.error);
    const exception = createExceptionPropertiesFromError(sanitizedError);

    const [taskRunError, taskRun] = await tryCatch(
      runStore.findRunOrThrow(
        {
          id: run.id,
        },
        {
          select: {
            id: true,
            friendlyId: true,
            traceId: true,
            spanId: true,
            parentSpanId: true,
            createdAt: true,
            completedAt: true,
            taskIdentifier: true,
            projectId: true,
            runtimeEnvironmentId: true,
            environmentType: true,
            isTest: true,
            organizationId: true,
            taskEventStore: true,
            // Piggyback the realtime run-changed publish on this existing read (no-op when
            // the native backend is disabled).
            runTags: true,
            batchId: true,
          },
        },
        $replica
      )
    );

    if (taskRunError) {
      logger.error("[runFailed] Failed to find task run", {
        error: taskRunError,
        runId: run.id,
      });
      return;
    }

    publishChangeRecord({
      runId: taskRun.id,
      envId: environment.id,
      tags: taskRun.runTags,
      batchId: taskRun.batchId,
      updatedAtMs: run.updatedAt.getTime(),
    });

    const eventRepository = await getEventRepositoryForStore(
      run.taskEventStore,
      taskRun.organizationId ?? organization.id
    );

    const [completeFailedRunEventError] = await tryCatch(
      eventRepository.completeFailedRunEvent({
        run: taskRun,
        endTime: time,
        exception,
      })
    );

    if (completeFailedRunEventError) {
      logger.error("[runFailed] Failed to complete failed run event", {
        error: completeFailedRunEventError,
        runId: run.id,
      });
    }
  });

  engine.eventBus.on("runAttemptFailed", async ({ time, run }) => {
    const sanitizedError = sanitizeError(run.error);
    const exception = createExceptionPropertiesFromError(sanitizedError);

    const [taskRunError, taskRun] = await tryCatch(
      runStore.findRunOrThrow(
        {
          id: run.id,
        },
        {
          select: {
            id: true,
            friendlyId: true,
            traceId: true,
            spanId: true,
            parentSpanId: true,
            createdAt: true,
            completedAt: true,
            taskIdentifier: true,
            projectId: true,
            runtimeEnvironmentId: true,
            environmentType: true,
            isTest: true,
            organizationId: true,
            taskEventStore: true,
            // Piggyback the realtime run-changed publish on this existing read (no-op when
            // the native backend is disabled).
            runTags: true,
            batchId: true,
          },
        },
        $replica
      )
    );

    if (taskRunError) {
      logger.error("[runAttemptFailed] Failed to find task run", {
        error: taskRunError,
        runId: run.id,
      });
      return;
    }

    publishChangeRecord({
      runId: taskRun.id,
      envId: taskRun.runtimeEnvironmentId,
      tags: taskRun.runTags,
      batchId: taskRun.batchId,
      updatedAtMs: run.updatedAt.getTime(),
    });

    if (!taskRun.organizationId) {
      logger.error("[runAttemptFailed] Task run has no organization id", {
        runId: run.id,
      });
      return;
    }

    const eventRepository = await getEventRepositoryForStore(
      run.taskEventStore,
      taskRun.organizationId
    );

    const [createAttemptFailedRunEventError] = await tryCatch(
      eventRepository.createAttemptFailedRunEvent({
        run: taskRun,
        endTime: time,
        attemptNumber: run.attemptNumber,
        exception,
      })
    );

    if (createAttemptFailedRunEventError) {
      logger.error("[runAttemptFailed] Failed to create attempt failed run event", {
        error: createAttemptFailedRunEventError,
        runId: run.id,
      });
    }
  });

  engine.eventBus.on(
    "cachedRunCompleted",
    async ({ time, span, blockedRunId, hasError, cachedRunId }) => {
      const [parentSpanId, spanId] = span.id.split(":");

      if (!spanId || !parentSpanId) {
        logger.debug("[cachedRunCompleted] Invalid span id", {
          spanId,
          parentSpanId,
        });
        return;
      }

      const [cachedRunError, cachedRun] = await tryCatch(
        runStore.findRunOrThrow(
          {
            id: cachedRunId,
          },
          {
            select: {
              id: true,
              friendlyId: true,
              traceId: true,
              spanId: true,
              parentSpanId: true,
              createdAt: true,
              completedAt: true,
              taskIdentifier: true,
              projectId: true,
              runtimeEnvironmentId: true,
              environmentType: true,
              isTest: true,
              organizationId: true,
            },
          },
          $replica
        )
      );

      if (cachedRunError) {
        logger.error("[cachedRunCompleted] Failed to find cached run", {
          error: cachedRunError,
          cachedRunId,
        });
        return;
      }

      const [blockedRunError, blockedRun] = await tryCatch(
        runStore.findRun(
          {
            id: blockedRunId,
          },
          {
            select: {
              id: true,
              friendlyId: true,
              traceId: true,
              spanId: true,
              parentSpanId: true,
              createdAt: true,
              completedAt: true,
              taskIdentifier: true,
              projectId: true,
              runtimeEnvironmentId: true,
              environmentType: true,
              isTest: true,
              organizationId: true,
              taskEventStore: true,
            },
          },
          $replica
        )
      );

      if (blockedRunError) {
        logger.error("[cachedRunCompleted] Failed to find blocked run", {
          error: blockedRunError,
          blockedRunId,
        });
      }

      if (!blockedRun) {
        logger.error("[cachedRunCompleted] Blocked run not found", {
          blockedRunId,
        });
        return;
      }

      if (!blockedRun.organizationId) {
        logger.error("[cachedRunCompleted] Blocked run has no organization id", {
          blockedRunId,
        });
        return;
      }

      const eventRepository = await getEventRepositoryForStore(
        blockedRun.taskEventStore,
        blockedRun.organizationId
      );

      const [completeCachedRunEventError] = await tryCatch(
        eventRepository.completeCachedRunEvent({
          run: cachedRun,
          blockedRun,
          spanId,
          parentSpanId,
          spanCreatedAt: span.createdAt,
          isError: hasError,
          endTime: time,
        })
      );

      if (completeCachedRunEventError) {
        logger.error("[cachedRunCompleted] Failed to complete cached run event", {
          error: completeCachedRunEventError,
          cachedRunId,
        });
      }
    }
  );

  engine.eventBus.on("runExpired", async ({ time, run, organization, environment }) => {
    if (!run.ttl) {
      return;
    }

    const [taskRunError, taskRun] = await tryCatch(
      runStore.findRunOrThrow(
        {
          id: run.id,
        },
        {
          select: {
            id: true,
            friendlyId: true,
            traceId: true,
            spanId: true,
            parentSpanId: true,
            createdAt: true,
            completedAt: true,
            taskIdentifier: true,
            projectId: true,
            runtimeEnvironmentId: true,
            environmentType: true,
            isTest: true,
            organizationId: true,
            taskEventStore: true,
            // Piggyback the realtime run-changed publish on this existing read (no-op when
            // the native backend is disabled).
            runTags: true,
            batchId: true,
          },
        },
        $replica
      )
    );

    if (taskRunError) {
      logger.error("[runExpired] Failed to find task run", {
        error: taskRunError,
        runId: run.id,
      });
      return;
    }

    publishChangeRecord({
      runId: taskRun.id,
      envId: environment.id,
      tags: taskRun.runTags,
      batchId: taskRun.batchId,
      updatedAtMs: run.updatedAt.getTime(),
    });

    const eventRepository = await getEventRepositoryForStore(
      taskRun.taskEventStore,
      taskRun.organizationId ?? organization.id
    );

    const [completeExpiredRunEventError] = await tryCatch(
      eventRepository.completeExpiredRunEvent({
        run: taskRun,
        endTime: time,
        ttl: run.ttl,
      })
    );

    if (completeExpiredRunEventError) {
      logger.error("[runExpired] Failed to complete expired run event", {
        error: completeExpiredRunEventError,
        runId: run.id,
      });
    }
  });

  engine.eventBus.on("runCancelled", async ({ time, run, organization, environment }) => {
    const [taskRunError, taskRun] = await tryCatch(
      runStore.findRunOrThrow(
        {
          id: run.id,
        },
        {
          select: {
            id: true,
            friendlyId: true,
            traceId: true,
            spanId: true,
            parentSpanId: true,
            createdAt: true,
            completedAt: true,
            taskIdentifier: true,
            projectId: true,
            runtimeEnvironmentId: true,
            environmentType: true,
            isTest: true,
            organizationId: true,
            taskEventStore: true,
            // Piggyback the realtime run-changed publish on this existing read (no-op when
            // the native backend is disabled).
            runTags: true,
            batchId: true,
          },
        },
        $replica
      )
    );

    if (taskRunError) {
      logger.error("[runCancelled] Task run not found", {
        error: taskRunError,
        runId: run.id,
      });
      return;
    }

    publishChangeRecord({
      runId: taskRun.id,
      envId: environment.id,
      tags: taskRun.runTags,
      batchId: taskRun.batchId,
      updatedAtMs: run.updatedAt.getTime(),
    });

    const eventRepository = await getEventRepositoryForStore(
      taskRun.taskEventStore,
      taskRun.organizationId ?? organization.id
    );

    const error = createJsonErrorObject(run.error);

    const [cancelRunEventError] = await tryCatch(
      eventRepository.cancelRunEvent({
        reason: error.message,
        run: taskRun,
        cancelledAt: time,
      })
    );

    if (cancelRunEventError) {
      logger.error("[runCancelled] Failed to cancel run event", {
        error: cancelRunEventError,
        runId: run.id,
      });
    }
  });

  engine.eventBus.on(
    "runRetryScheduled",
    async ({ time, run, environment, retryAt, organization }) => {
      try {
        if (retryAt && time && time >= retryAt) {
          return;
        }

        let retryMessage = `Retry ${
          typeof run.attemptNumber === "number" ? `#${run.attemptNumber - 1}` : ""
        } delay`;

        if (run.nextMachineAfterOOM) {
          retryMessage += ` after OOM`;
        }

        const eventRepository = await getEventRepositoryForStore(
          run.taskEventStore ?? "taskEvent",
          organization.id
        );

        await eventRepository.recordEvent(retryMessage, {
          startTime: BigInt(time.getTime() * 1000000),
          taskSlug: run.taskIdentifier,
          environment,
          attributes: {
            properties: {
              retryAt: retryAt.toISOString(),
              nextMachine: run.nextMachineAfterOOM,
            },
            runId: run.friendlyId,
            style: {
              icon: "schedule-attempt",
            },
          },
          context: run.traceContext as Record<string, string | undefined>,
          endTime: retryAt,
        });
      } catch (error) {
        logger.error("[runRetryScheduled] Failed to record retry event", {
          error: error instanceof Error ? error.message : error,
          runId: run.id,
          spanId: run.spanId,
        });
      }
    }
  );

  engine.eventBus.on("runAttemptStarted", async ({ time, run, organization }) => {
    try {
      if (run.attemptNumber === 1 && run.baseCostInCents > 0) {
        await reportInvocationUsage(organization.id, run.baseCostInCents, { runId: run.id });
      }
    } catch (error) {
      logger.error("[runAttemptStarted] Failed to report invocation usage", {
        error: error instanceof Error ? error.message : error,
        runId: run.id,
        orgId: organization.id,
      });
    }
  });

  engine.eventBus.on("runMetadataUpdated", async ({ time, run }) => {
    const result = await findEnvironmentFromRun(run.id);

    if (!result) {
      logger.error("[runMetadataUpdated] Failed to find environment", { runId: run.id });
      return;
    }

    const { environment, runTags, batchId } = result;

    try {
      const updateResult = await updateMetadataService.call(run.id, run.metadata, environment);
      // Realtime run-changed publish, after the write so the router's hydrate sees the new
      // row. A full record (env + tags + batchId + the committed updatedAt watermark), so
      // feeds route by index. Nothing written here (no-op or buffered) = nothing to announce.
      if (updateResult?.updatedAtMs !== undefined) {
        publishChangeRecord({
          runId: run.id,
          envId: environment.id,
          tags: runTags,
          batchId,
          updatedAtMs: updateResult.updatedAtMs,
        });
      }
    } catch (e) {
      if (e instanceof MetadataTooLargeError) {
        logger.warn("[runMetadataUpdated] Failed to update metadata, too large", {
          taskRun: run.id,
          error:
            e instanceof Error
              ? {
                  name: e.name,
                  message: e.message,
                  stack: e.stack,
                }
              : e,
        });
      } else {
        logger.error("[runMetadataUpdated] Failed to update metadata", {
          taskRun: run.id,
          error:
            e instanceof Error
              ? {
                  name: e.name,
                  message: e.message,
                  stack: e.stack,
                }
              : e,
        });
      }
    }
  });

  engine.eventBus.on("executionSnapshotCreated", async ({ time, run, snapshot }) => {
    const eventResult = await recordRunDebugLog(
      run.id,
      `[engine] ${snapshot.executionStatus} - ${snapshot.description}`,
      {
        attributes: {
          properties: {
            snapshotId: snapshot.id,
            snapshotDescription: snapshot.description,
            snapshotStatus: snapshot.executionStatus,
            workerId: snapshot.workerId ?? undefined,
            runnerId: snapshot.runnerId ?? undefined,
          },
        },
        startTime: time,
      }
    );

    if (!eventResult.success) {
      logger.error("[executionSnapshotCreated] Failed to record event", {
        runId: run.id,
        snapshot,
        error: eventResult.error,
      });
    }
  });

  engine.eventBus.on("workerNotification", async ({ time, run, snapshot }) => {
    logger.debug("[workerNotification] Notifying worker", { time, runId: run.id, snapshot });

    // Notify the worker
    try {
      const runFriendlyId = RunId.toFriendlyId(run.id);
      const room = roomFromFriendlyRunId(runFriendlyId);

      //send the notification to connected workers
      socketIo.workerNamespace
        .to(room)
        .emit("run:notify", { version: "1", run: { friendlyId: runFriendlyId } });

      //send the notification to connected dev workers
      socketIo.devWorkerNamespace
        .to(room)
        .emit("run:notify", { version: "1", run: { friendlyId: runFriendlyId } });

      if (!env.RUN_ENGINE_DEBUG_WORKER_NOTIFICATIONS) {
        return;
      }

      // Record notification event
      const eventResult = await recordRunDebugLog(
        run.id,
        // don't prefix this with [engine] - "run:notify" is the correct prefix
        `run:notify platform -> supervisor: ${snapshot.executionStatus}`,
        {
          attributes: {
            properties: {
              snapshotId: snapshot.id,
              snapshotStatus: snapshot.executionStatus,
            },
          },
          startTime: time,
        }
      );

      if (!eventResult.success) {
        logger.error("[workerNotification] Failed to record event", {
          runId: run.id,
          snapshot,
          error: eventResult.error,
        });
      }
    } catch (error) {
      logger.error("[workerNotification] Failed to notify worker", {
        error: error instanceof Error ? error.message : error,
        runId: run.id,
        snapshot,
      });

      // Record notification event
      const eventResult = await recordRunDebugLog(
        run.id,
        // don't prefix this with [engine] - "run:notify" is the correct prefix
        `run:notify ERROR platform -> supervisor: ${snapshot.executionStatus}`,
        {
          attributes: {
            properties: {
              snapshotId: snapshot.id,
              snapshotStatus: snapshot.executionStatus,
              error: error instanceof Error ? error.message : String(error),
            },
          },
          startTime: time,
        }
      );

      if (!eventResult.success) {
        logger.error("[workerNotification] Failed to record event", {
          runId: run.id,
          snapshot,
          error: eventResult.error,
        });
      }
    }
  });

  engine.eventBus.on("incomingCheckpointDiscarded", async ({ time, run, snapshot, checkpoint }) => {
    const eventResult = await recordRunDebugLog(
      run.id,
      `[engine] Checkpoint discarded: ${checkpoint.discardReason}`,
      {
        attributes: {
          properties: {
            snapshotId: snapshot.id,
            ...checkpoint.metadata,
          },
        },
        startTime: time,
      }
    );

    if (!eventResult.success) {
      logger.error("[incomingCheckpointDiscarded] Failed to record event", {
        runId: run.id,
        snapshot,
        error: eventResult.error,
      });
    }
  });
}

/**
 * errorCode returned by the batch process-item callback when the trigger was
 * rejected because the environment's queue is at its maximum size. The
 * BatchQueue (via `skipRetries`) short-circuits retries for this code, and the
 * batch completion callback collapses per-item errors into a single aggregate
 * `BatchTaskRunError` row instead of writing one per item.
 */
const QUEUE_SIZE_LIMIT_EXCEEDED_ERROR_CODE = "QUEUE_SIZE_LIMIT_EXCEEDED";

/**
 * Set up the BatchQueue processing callbacks.
 * These handle creating runs from batch items and completing batches.
 *
 * Payload handling:
 * - If payloadType is "application/store", the payload is an R2 path (already offloaded)
 * - DefaultPayloadProcessor in TriggerTaskService will pass it through without re-offloading
 * - The run engine will download from R2 when the task executes
 */
export function setupBatchQueueCallbacks() {
  // Item processing callback - creates a run for each batch item
  engine.setBatchProcessItemCallback(
    async ({ batchId, friendlyId, itemIndex, item, meta, attempt, isFinalAttempt }) => {
      return tracer.startActiveSpan(
        "batch.processItem",
        {
          kind: SpanKind.INTERNAL,
          attributes: {
            "batch.id": friendlyId,
            "batch.item_index": itemIndex,
            "batch.task": item.task,
            "batch.environment_id": meta.environmentId,
            "batch.parent_run_id": meta.parentRunId ?? "",
            "batch.attempt": attempt,
            "batch.is_final_attempt": isFinalAttempt,
          },
        },
        async (span) => {
          const triggerFailedTaskService = new TriggerFailedTaskService({
            prisma,
            engine,
            replicaPrisma: $replica,
          });

          // Check for pre-marked error items (e.g. oversized payloads)
          const itemError = item.options?.__error as string | undefined;
          if (itemError) {
            const errorCode = (item.options?.__errorCode as string) ?? "ITEM_ERROR";

            let environment: AuthenticatedEnvironment | undefined;
            try {
              environment = (await findEnvironmentById(meta.environmentId)) ?? undefined;
            } catch {
              // Best-effort environment lookup
            }

            if (environment) {
              const failedRunId = await triggerFailedTaskService.call({
                taskId: item.task,
                environment,
                payload: item.payload ?? "{}",
                payloadType: item.payloadType as string,
                errorMessage: itemError,
                errorCode: errorCode as TaskRunErrorCodes,
                parentRunId: meta.parentRunId,
                resumeParentOnCompletion: meta.resumeParentOnCompletion,
                batch: { id: batchId, index: itemIndex },
                traceContext: meta.traceContext as Record<string, unknown> | undefined,
                spanParentAsLink: meta.spanParentAsLink,
              });

              if (failedRunId) {
                span.setAttribute("batch.result.pre_failed", true);
                span.setAttribute("batch.result.run_id", failedRunId);
                span.end();
                return { success: true as const, runId: failedRunId };
              }
            }

            // Fallback if TriggerFailedTaskService or environment lookup fails
            span.end();
            return { success: false as const, error: itemError, errorCode };
          }

          let environment: AuthenticatedEnvironment | undefined;
          try {
            environment = (await findEnvironmentById(meta.environmentId)) ?? undefined;

            if (!environment) {
              span.setAttribute("batch.result.error", "Environment not found");
              span.end();

              return {
                success: false as const,
                error: "Environment not found",
                errorCode: "ENVIRONMENT_NOT_FOUND",
              };
            }

            const triggerTaskService = new TriggerTaskService();

            // Normalize payload - for application/store (R2 paths), this passes through as-is
            const payload = normalizePayload(item.payload, item.payloadType);

            const result = await triggerTaskService.call(
              item.task,
              environment,
              {
                payload,
                options: {
                  ...(item.options as Record<string, unknown>),
                  payloadType: item.payloadType,
                  parentRunId: meta.parentRunId,
                  resumeParentOnCompletion: meta.resumeParentOnCompletion,
                  parentBatch: batchId,
                },
              },
              {
                triggerVersion: meta.triggerVersion,
                traceContext: meta.traceContext as Record<string, unknown> | undefined,
                spanParentAsLink: meta.spanParentAsLink,
                batchId,
                batchIndex: itemIndex,
                realtimeStreamsVersion: meta.realtimeStreamsVersion,
                planType: meta.planType,
                triggerSource: meta.parentRunId ? "sdk" : (meta.triggerSource ?? "api"),
                triggerAction: "trigger",
              },
              "V2"
            );

            if (result) {
              span.setAttribute("batch.result.run_id", result.run.friendlyId);
              span.end();
              return { success: true as const, runId: result.run.friendlyId };
            } else {
              logger.error("[BatchQueue] TriggerTaskService returned undefined", {
                batchId,
                friendlyId,
                itemIndex,
                task: item.task,
                environmentId: meta.environmentId,
                attempt,
                isFinalAttempt,
              });

              span.setAttribute("batch.result.error", "TriggerTaskService returned undefined");

              // Only create a pre-failed run on the final attempt; otherwise let the retry mechanism handle it
              if (isFinalAttempt) {
                const failedRunId = await triggerFailedTaskService.call({
                  taskId: item.task,
                  environment,
                  payload: item.payload,
                  payloadType: item.payloadType as string,
                  errorMessage: "TriggerTaskService returned undefined",
                  parentRunId: meta.parentRunId,
                  resumeParentOnCompletion: meta.resumeParentOnCompletion,
                  batch: { id: batchId, index: itemIndex },
                  options: item.options as Record<string, unknown>,
                  traceContext: meta.traceContext as Record<string, unknown> | undefined,
                  spanParentAsLink: meta.spanParentAsLink,
                  errorCode: TaskRunErrorCodes.BATCH_ITEM_COULD_NOT_TRIGGER,
                });

                span.end();

                if (failedRunId) {
                  return { success: true as const, runId: failedRunId };
                }
              } else {
                span.end();
              }

              return {
                success: false as const,
                error: "TriggerTaskService returned undefined",
                errorCode: "TRIGGER_FAILED",
              };
            }
          } catch (error) {
            const errorMessage = error instanceof Error ? error.message : String(error);

            // Queue-size-limit rejections are a customer-overload scenario (the
            // env's queue is at its configured max). Retrying is pointless — the
            // same item will fail again — and creating pre-failed TaskRuns for
            // every item of every retried batch is exactly what chews through
            // DB capacity when a noisy tenant fills their queue. Signal the
            // BatchQueue to skip retries and skip pre-failed run creation, and
            // let the completion callback collapse the per-item errors into a
            // single summary row.
            if (error instanceof QueueSizeLimitExceededError) {
              logger.warn("[BatchQueue] Batch item rejected: queue size limit reached", {
                batchId,
                friendlyId,
                itemIndex,
                task: item.task,
                environmentId: meta.environmentId,
                maximumSize: error.maximumSize,
              });

              span.setAttribute("batch.result.error", errorMessage);
              span.setAttribute("batch.result.errorCode", QUEUE_SIZE_LIMIT_EXCEEDED_ERROR_CODE);
              span.setAttribute("batch.result.skipRetries", true);
              span.end();

              return {
                success: false as const,
                error: errorMessage,
                errorCode: QUEUE_SIZE_LIMIT_EXCEEDED_ERROR_CODE,
                skipRetries: true,
              };
            }

            logger.error("[BatchQueue] Failed to trigger batch item", {
              batchId,
              friendlyId,
              itemIndex,
              task: item.task,
              environmentId: meta.environmentId,
              attempt,
              isFinalAttempt,
              error,
            });

            span.setAttribute("batch.result.error", errorMessage);
            span.recordException(error instanceof Error ? error : new Error(String(error)));

            // Only create a pre-failed run on the final attempt; otherwise let the retry mechanism handle it
            if (isFinalAttempt && environment) {
              const failedRunId = await triggerFailedTaskService.call({
                taskId: item.task,
                environment,
                payload: item.payload,
                payloadType: item.payloadType as string,
                errorMessage,
                parentRunId: meta.parentRunId,
                resumeParentOnCompletion: meta.resumeParentOnCompletion,
                batch: { id: batchId, index: itemIndex },
                options: item.options as Record<string, unknown>,
                traceContext: meta.traceContext as Record<string, unknown> | undefined,
                spanParentAsLink: meta.spanParentAsLink,
                errorCode: TaskRunErrorCodes.BATCH_ITEM_COULD_NOT_TRIGGER,
              });

              span.end();

              if (failedRunId) {
                return { success: true as const, runId: failedRunId };
              }
            } else {
              span.end();
            }

            return {
              success: false as const,
              error: errorMessage,
              errorCode: "TRIGGER_ERROR",
            };
          }
        }
      );
    }
  );

  // Batch completion callback - updates Postgres with results
  engine.setBatchCompletionCallback(async (result: CompleteBatchResult) => {
    const { batchId, runIds, successfulRunCount, failedRunCount, failures } = result;

    // Determine final status
    let status: BatchTaskRunStatus;
    if (failedRunCount > 0 && successfulRunCount === 0) {
      status = "ABORTED";
    } else if (failedRunCount > 0) {
      status = "PARTIAL_FAILED";
    } else {
      status = "PENDING"; // All runs created, waiting for completion
    }

    try {
      // Use a transaction to ensure atomicity of batch update and error record creation
      // skipDuplicates handles idempotency when callback is retried (relies on unique constraint)
      await prisma.$transaction(async (tx) => {
        // Update BatchTaskRun
        await tx.batchTaskRun.update({
          where: { id: batchId },
          data: {
            status,
            runIds,
            successfulRunCount,
            failedRunCount,
            completedAt: status === "ABORTED" ? new Date() : undefined,
            processingCompletedAt: new Date(),
          },
        });

        // Create error records if there were failures.
        //
        // Fast-path for queue-size-limit overload: when every failure is the
        // same QUEUE_SIZE_LIMIT_EXCEEDED error, collapse them into a single
        // aggregate row instead of writing one per item. This keeps the DB
        // write volume bounded to O(batches) instead of O(items) when a noisy
        // tenant fills their queue and all of their batches start bouncing.
        if (failures.length > 0) {
          const allQueueSizeLimit = failures.every(
            (f) => f.errorCode === QUEUE_SIZE_LIMIT_EXCEEDED_ERROR_CODE
          );

          if (allQueueSizeLimit) {
            const sample = failures[0]!;
            await tx.batchTaskRunError.createMany({
              data: [
                {
                  batchTaskRunId: batchId,
                  // Use the first item's index as a stable anchor for the
                  // (batchTaskRunId, index) unique constraint so callback
                  // retries remain idempotent.
                  index: sample.index,
                  taskIdentifier: sample.taskIdentifier,
                  payload: sample.payload,
                  options: sample.options as Prisma.InputJsonValue | undefined,
                  error: `${sample.error} (${failures.length} items in this batch failed with the same error)`,
                  errorCode: sample.errorCode,
                },
              ],
              skipDuplicates: true,
            });
          } else {
            await tx.batchTaskRunError.createMany({
              data: failures.map((failure) => ({
                batchTaskRunId: batchId,
                index: failure.index,
                taskIdentifier: failure.taskIdentifier,
                payload: failure.payload,
                options: failure.options as Prisma.InputJsonValue | undefined,
                error: failure.error,
                errorCode: failure.errorCode,
              })),
              skipDuplicates: true,
            });
          }
        }
      });

      // Try to complete the batch (handles waitpoint completion if all runs are done)
      if (status !== "ABORTED") {
        await engine.tryCompleteBatch({ batchId });
      }

      logger.info("Batch completion handled", {
        batchId,
        status,
        successfulRunCount,
        failedRunCount,
      });
    } catch (error) {
      logger.error("Failed to handle batch completion", {
        batchId,
        error: error instanceof Error ? error.message : String(error),
      });
      // Re-throw to preserve Redis data for retry (BatchQueue expects errors to propagate)
      throw error;
    }
  });

  logger.info("BatchQueue callbacks configured");
}

/**
 * Normalize the payload from BatchQueue.
 *
 * Handles different payload types:
 * - "application/store": Already offloaded to R2, payload is the path - pass through as-is
 * - "application/json": May be a pre-serialized JSON string - parse to avoid double-stringification
 * - Other types: Pass through as-is
 *
 * @param payload - The raw payload from the batch item
 * @param payloadType - The payload type (e.g., "application/json", "application/store")
 */
function normalizePayload(payload: unknown, payloadType?: string): unknown {
  // Only process "application/json" payloads
  // For all other types (including undefined), return as-is
  if (payloadType !== "application/json") {
    return payload;
  }

  // For JSON payloads, if payload is a string, try to parse it
  // This handles pre-serialized JSON from the SDK
  if (typeof payload === "string") {
    try {
      return JSON.parse(payload);
    } catch {
      // If it's not valid JSON, return as-is
      return payload;
    }
  }

  return payload;
}
