import type {
  BatchTriggerTaskV2RequestBody,
  BatchTriggerTaskV2Response,
  IOPacket,
} from "@trigger.dev/core/v3";
import { packetRequiresOffloading, parsePacket } from "@trigger.dev/core/v3";
import type { BatchTaskRun, TaskRunAttempt } from "@trigger.dev/database";
import {
  isPrismaRaceConditionError,
  isPrismaRetriableError,
  isUniqueConstraintError,
  Prisma,
} from "@trigger.dev/database";
import type { RunStore } from "@internal/run-store";
import { z } from "zod";
import type { PrismaClientOrTransaction } from "~/db.server";
import { prisma } from "~/db.server";
import { runStore as defaultRunStore } from "~/v3/runStore.server";
import { env } from "~/env.server";
import { findEnvironmentById } from "~/models/runtimeEnvironment.server";
import { batchTaskRunItemStatusForRunStatus } from "~/models/taskRun.server";
import type { AuthenticatedEnvironment } from "~/services/apiAuth.server";
import { dependentAttemptWhere } from "./dependentAttemptScope";
import { logger } from "~/services/logger.server";
import { getEntitlement } from "~/services/platform.v3.server";
import { controlPlaneResolver } from "~/v3/runOpsMigration/controlPlaneResolver.server";
import { resolveRunIdMintKind, type RunIdMintKind } from "~/v3/engineVersion.server";
import { resolveInheritedMintKind } from "~/v3/runOpsMigration/resolveInheritedMintKind.server";
import { mintFriendlyIdForKind } from "~/v3/runOpsMigration/mintAnchoredRunFriendlyId.server";
import { mintBatchFriendlyId } from "~/v3/runOpsMigration/mintBatchFriendlyId.server";
import { batchTriggerWorker } from "../batchTriggerWorker.server";
import { legacyRunEngineWorker } from "../legacyRunEngineWorker.server";
import { marqs } from "../marqs/index.server";
import { guardQueueSizeLimitsForEnv } from "../queueSizeLimits.server";
import { downloadPacketFromObjectStore, uploadPacketToObjectStore } from "../objectStore.server";
import { isFinalAttemptStatus, isFinalRunStatus } from "../taskStatus";
import { startActiveSpan } from "../tracer.server";
import { BaseService, ServiceValidationError } from "./baseService.server";
import { ResumeBatchRunService } from "./resumeBatchRun.server";
import { OutOfEntitlementError, TriggerTaskService } from "./triggerTask.server";

const PROCESSING_BATCH_SIZE = 50;
const ASYNC_BATCH_PROCESS_SIZE_THRESHOLD = 20;
const MAX_ATTEMPTS = 10;

export const BatchProcessingStrategy = z.enum(["sequential", "parallel"]);
export type BatchProcessingStrategy = z.infer<typeof BatchProcessingStrategy>;

export const BatchProcessingOptions = z.object({
  batchId: z.string(),
  processingId: z.string(),
  range: z.object({ start: z.number().int(), count: z.number().int() }),
  attemptCount: z.number().int(),
  strategy: BatchProcessingStrategy,
});

export type BatchProcessingOptions = z.infer<typeof BatchProcessingOptions>;

export type BatchTriggerTaskServiceOptions = {
  idempotencyKey?: string;
  idempotencyKeyExpiresAt?: Date;
  triggerVersion?: string;
  traceContext?: Record<string, string | undefined>;
  spanParentAsLink?: boolean;
  oneTimeUseToken?: string;
  realtimeStreamsVersion?: "v1" | "v2";
  triggerSource?: string;
  triggerAction?: string;
};

type RunItemData = {
  id: string;
  isCached: boolean;
  idempotencyKey: string | undefined;
  taskIdentifier: string;
};

/**
 * ### V3
 *
 * BatchTrigger v3 doesn't have any changes from v2, other than a different system for tracking if the
 * batch is completed.
 *
 * v3 BatchTaskRun's now must be "sealed" before they could be considered completed. Being "sealed" means
 * that all the items in the batch have been processed and the batch is ready to be considered completed.
 *
 * We also now track the expected count of items in the batch, and then as each BatchTaskRunItem is set to COMPLETED,
 * we increment the BatchTaskRun's completed count. Once the completed count is equal to the expected count, and the
 * batch is sealed, we can consider the batch completed.
 *
 * So now when the v3 batch is considered completed, we will enqueue the ResumeBatchRunService to resume the dependent
 * task attempt if there is one. This is in contrast to v2 batches where every time a task was completed, we would schedule
 * the ResumeBatchRunService to check if the batch was completed and set it to completed if it was.
 *
 * We've also introduced a new column "resumedAt" that will be set when the batch is resumed. Previously in v2 batches, the status == "COMPLETED" was overloaded
 * to mean that the batch was completed and resumed. Now we have a separate column to track when the batch was resumed (and to make sure it's only resumed once).
 *
 * ### V2
 *
 * Batch v2 added the ability to trigger more than 100 tasks in a single batch. This was done by offloading the payload to the object store and
 * then processing the batch in chunks of 50 tasks at a time in the background.
 *
 * The other main difference from v1 is that a single batch in v2 could trigger multiple different tasks, whereas in v1 a batch could only trigger a single task.
 */
export class BatchTriggerV3Service extends BaseService {
  private _batchProcessingStrategy: BatchProcessingStrategy;
  private _asyncBatchProcessSizeThreshold: number;

  constructor(
    batchProcessingStrategy?: BatchProcessingStrategy,
    asyncBatchProcessSizeThreshold: number = ASYNC_BATCH_PROCESS_SIZE_THRESHOLD,
    protected readonly _prisma: PrismaClientOrTransaction = prisma,
    protected readonly runStore: RunStore = defaultRunStore,
    // Injected so tests force the env-default branch deterministically; defaults
    // to the live per-env mint resolver.
    private readonly resolveMintKind: (environment: {
      organizationId: string;
      id: string;
      orgFeatureFlags?: unknown;
    }) => Promise<RunIdMintKind> = resolveRunIdMintKind
  ) {
    super(_prisma);

    this._batchProcessingStrategy = batchProcessingStrategy ?? "parallel";
    this._asyncBatchProcessSizeThreshold = asyncBatchProcessSizeThreshold;
  }

  public async call(
    environment: AuthenticatedEnvironment,
    body: BatchTriggerTaskV2RequestBody,
    options: BatchTriggerTaskServiceOptions = {}
  ): Promise<BatchTriggerTaskV2Response> {
    try {
      return await this.traceWithEnv<BatchTriggerTaskV2Response>(
        "call()",
        environment,
        async (span) => {
          if (!body.items || body.items.length === 0) {
            throw new ServiceValidationError("A batch trigger must have at least one item");
          }

          // BatchTaskRun.runtimeEnvironmentId no longer has an FK into RuntimeEnvironment;
          // validate env existence app-side before any create arm (passthrough when split is off).
          await controlPlaneResolver.assertEnvExists(environment.id);

          const existingBatch = options.idempotencyKey
            ? await this.runStore.findBatchTaskRunByIdempotencyKey(
                environment.id,
                options.idempotencyKey
              )
            : undefined;

          if (existingBatch) {
            if (
              existingBatch.idempotencyKeyExpiresAt &&
              existingBatch.idempotencyKeyExpiresAt < new Date()
            ) {
              logger.debug("[BatchTriggerV2][call] Idempotency key has expired", {
                idempotencyKey: options.idempotencyKey,
                batch: {
                  id: existingBatch.id,
                  friendlyId: existingBatch.friendlyId,
                  runCount: existingBatch.runCount,
                  idempotencyKeyExpiresAt: existingBatch.idempotencyKeyExpiresAt,
                  idempotencyKey: existingBatch.idempotencyKey,
                },
              });

              // Update the existing batch to remove the idempotency key
              await this.runStore.updateBatchTaskRun({
                where: { id: existingBatch.id },
                data: { idempotencyKey: null },
                select: { id: true },
              });

              // Don't return, just continue with the batch trigger
            } else {
              span.setAttribute("batchId", existingBatch.friendlyId);

              return this.#respondWithExistingBatch(existingBatch, environment);
            }
          }

          const { id: batchInternalId, friendlyId: batchId } = await mintBatchFriendlyId({
            environment: {
              organizationId: environment.organizationId,
              id: environment.id,
              orgFeatureFlags: environment.organization.featureFlags,
            },
            parentRunFriendlyId: body.parentRunId,
          });

          span.setAttribute("batchId", batchId);

          const dependentAttempt = body?.dependentAttempt
            ? await this._prisma.taskRunAttempt.findFirst({
                // Scope to the caller's environment (see dependentAttemptWhere).
                where: dependentAttemptWhere(body.dependentAttempt, environment.id),
                include: {
                  taskRun: {
                    select: {
                      id: true,
                      status: true,
                    },
                  },
                },
              })
            : undefined;

          if (
            dependentAttempt &&
            (isFinalAttemptStatus(dependentAttempt.status) ||
              isFinalRunStatus(dependentAttempt.taskRun.status))
          ) {
            logger.debug("[BatchTriggerV2][call] Dependent attempt or run is in a terminal state", {
              dependentAttempt: dependentAttempt,
              batchId,
            });

            throw new ServiceValidationError(
              "Cannot process batch as the parent run is already in a terminal state"
            );
          }

          if (environment.type !== "DEVELOPMENT") {
            const result = await getEntitlement(environment.organizationId);
            if (result && result.hasAccess === false) {
              throw new OutOfEntitlementError();
            }
          }

          const runs = await this.#prepareRunData(environment, body, batchId);

          const newRunCount = runs.filter((r) => !r.isCached).length;

          if (newRunCount === 0) {
            logger.debug("[BatchTriggerV2][call] All runs are cached", {
              batchId,
            });

            await this.runStore.createBatchTaskRun({
              id: batchInternalId,
              friendlyId: batchId,
              runtimeEnvironmentId: environment.id,
              idempotencyKey: options.idempotencyKey,
              idempotencyKeyExpiresAt: options.idempotencyKeyExpiresAt,
              dependentTaskAttemptId: dependentAttempt?.id,
              runCount: body.items.length,
              runIds: runs.map((r) => r.id),
              status: "COMPLETED",
              batchVersion: "v3",
              oneTimeUseToken: options.oneTimeUseToken,
            });

            return {
              id: batchId,
              isCached: false,
              idempotencyKey: options.idempotencyKey ?? undefined,
              runs,
            };
          }

          const queueSizeGuard = await guardQueueSizeLimitsForEnv(environment, marqs, newRunCount);

          logger.debug("Queue size guard result", {
            newRunCount,
            queueSizeGuard,
            environment: {
              id: environment.id,
              type: environment.type,
              organization: environment.organization,
              project: environment.project,
            },
          });

          if (!queueSizeGuard.isWithinLimits) {
            throw new ServiceValidationError(
              `Cannot trigger ${newRunCount} tasks as the queue size limit for this environment has been reached. The maximum size is ${queueSizeGuard.maximumSize}`,
              undefined,
              "warn"
            );
          }

          // Upload to object store
          const payloadPacket = await this.#handlePayloadPacket(
            body.items,
            `batch/${batchId}`,
            environment
          );

          const batch = await this.#createAndProcessBatchTaskRun(
            batchId,
            batchInternalId,
            runs,
            payloadPacket,
            newRunCount,
            environment,
            body,
            options,
            dependentAttempt ?? undefined
          );

          if (!batch) {
            throw new Error("Failed to create batch");
          }

          return {
            id: batch.friendlyId,
            isCached: false,
            idempotencyKey: batch.idempotencyKey ?? undefined,
            runs,
          };
        }
      );
    } catch (error) {
      // Detect a prisma transaction Unique constraint violation
      if (error instanceof Prisma.PrismaClientKnownRequestError) {
        logger.debug("BatchTriggerV2: Prisma transaction error", {
          code: error.code,
          message: error.message,
          meta: error.meta,
        });

        if (error.code === "P2002") {
          const target = error.meta?.target;

          if (
            Array.isArray(target) &&
            target.length > 0 &&
            typeof target[0] === "string" &&
            target[0].includes("oneTimeUseToken")
          ) {
            throw new ServiceValidationError(
              "Cannot batch trigger with a one-time use token as it has already been used."
            );
          } else {
            throw new ServiceValidationError(
              "Cannot batch trigger as it has already been triggered with the same idempotency key."
            );
          }
        }
      }

      throw error;
    }
  }

  // Mint a child run's friendlyId so it lands in the SAME physical store as its
  // residency anchor. The caller passes the batch's friendlyId, so a run-ops id
  // (NEW) anchor yields a run-ops id (NEW) child and a cuid anchor yields a cuid
  // (LEGACY) child. With no anchor it falls back to the env's cutover setting.
  // Mirrors RunEngineTriggerTaskService.mintRunFriendlyId.
  private async mintChildFriendlyId(
    environment: AuthenticatedEnvironment,
    anchorFriendlyId?: string,
    region?: string
  ): Promise<string> {
    const mintKind = anchorFriendlyId
      ? resolveInheritedMintKind(anchorFriendlyId)
      : await this.resolveMintKind({
          organizationId: environment.organizationId,
          id: environment.id,
          orgFeatureFlags: environment.organization.featureFlags,
        });

    return mintFriendlyIdForKind(mintKind, region);
  }

  async #prepareRunData(
    environment: AuthenticatedEnvironment,
    body: BatchTriggerTaskV2RequestBody,
    batchFriendlyId: string
  ): Promise<Array<RunItemData>> {
    // Anchor every child to the batch's residency: the batch friendlyId is
    // minted once, so deriving each child's id-kind from it — rather than re-resolving
    // the per-org flag, which can flip mid-batch — keeps batch + children co-resident.
    const childAnchor = batchFriendlyId;

    // batchTriggerAndWait cannot have cached runs because that does not work in run engine v1 and is not available in the client
    if (body?.dependentAttempt) {
      return Promise.all(
        body.items.map(async (item) => ({
          id: await this.mintChildFriendlyId(environment, childAnchor, item.options?.region),
          isCached: false,
          idempotencyKey: undefined,
          taskIdentifier: item.task,
        }))
      );
    }

    // Group items by taskIdentifier
    const itemsByTask = body.items.reduce(
      (acc, item) => {
        if (!item.options?.idempotencyKey) return acc;

        if (!acc[item.task]) {
          acc[item.task] = [];
        }
        acc[item.task].push(item);
        return acc;
      },
      {} as Record<string, typeof body.items>
    );

    logger.debug("[BatchTriggerV2][call] Grouped items by task identifier", {
      itemsByTask,
    });

    // Fetch cached runs for each task identifier separately to make use of the index
    const cachedRuns = await Promise.all(
      Object.entries(itemsByTask).map(([taskIdentifier, items]) =>
        this.runStore.findRuns(
          {
            where: {
              runtimeEnvironmentId: environment.id,
              taskIdentifier,
              idempotencyKey: {
                in: items.map((i) => i.options?.idempotencyKey).filter(Boolean),
              },
            },
            select: {
              friendlyId: true,
              idempotencyKey: true,
              idempotencyKeyExpiresAt: true,
            },
          },
          this._prisma
        )
      )
    ).then((results) => results.flat());

    // Build the run IDs in order: reuse an unexpired cached id, else mint a new id (and record any
    // expired cached id so its idempotency key can be cleared below).
    const expiredRunIds = new Set<string>();

    const runs = await Promise.all(
      body.items.map(async (item) => {
        const cachedRun = cachedRuns.find((r) => r.idempotencyKey === item.options?.idempotencyKey);

        if (cachedRun) {
          if (cachedRun.idempotencyKeyExpiresAt && cachedRun.idempotencyKeyExpiresAt < new Date()) {
            expiredRunIds.add(cachedRun.friendlyId);

            return {
              id: await this.mintChildFriendlyId(environment, childAnchor, item.options?.region),
              isCached: false,
              idempotencyKey: item.options?.idempotencyKey ?? undefined,
              taskIdentifier: item.task,
            };
          }

          return {
            id: cachedRun.friendlyId,
            isCached: true,
            idempotencyKey: item.options?.idempotencyKey ?? undefined,
            taskIdentifier: item.task,
          };
        }

        return {
          id: await this.mintChildFriendlyId(environment, childAnchor, item.options?.region),
          isCached: false,
          idempotencyKey: item.options?.idempotencyKey ?? undefined,
          taskIdentifier: item.task,
        };
      })
    );

    // Expire the cached runs that are no longer valid
    if (expiredRunIds.size) {
      await this.runStore.clearIdempotencyKey(
        { byFriendlyIds: Array.from(expiredRunIds) },
        this._prisma
      );
    }

    return runs;
  }

  async #createAndProcessBatchTaskRun(
    batchId: string,
    batchInternalId: string,
    runs: Array<RunItemData>,
    payloadPacket: IOPacket,
    newRunCount: number,
    environment: AuthenticatedEnvironment,
    body: BatchTriggerTaskV2RequestBody,
    options: BatchTriggerTaskServiceOptions = {},
    dependentAttempt?: TaskRunAttempt
  ) {
    if (runs.length <= this._asyncBatchProcessSizeThreshold) {
      const batch = await this.runStore.createBatchTaskRun({
        id: batchInternalId,
        friendlyId: batchId,
        runtimeEnvironmentId: environment.id,
        idempotencyKey: options.idempotencyKey,
        idempotencyKeyExpiresAt: options.idempotencyKeyExpiresAt,
        dependentTaskAttemptId: dependentAttempt?.id,
        runCount: runs.length,
        runIds: runs.map((r) => r.id),
        payload: payloadPacket.data,
        payloadType: payloadPacket.dataType,
        options,
        batchVersion: "v3",
        oneTimeUseToken: options.oneTimeUseToken,
      });

      const result = await this.#processBatchTaskRunItems(
        batch,
        environment,
        0,
        PROCESSING_BATCH_SIZE,
        body.items,
        options
      );

      if (result.error) {
        logger.error("[BatchTriggerV2][call] Batch inline processing error", {
          batchId: batch.friendlyId,
          currentIndex: result.workingIndex,
          error: result.error,
        });

        await this.runStore.updateBatchTaskRun({
          where: { id: batch.id },
          data: {
            status: "ABORTED",
            completedAt: new Date(),
          },
          select: { id: true },
        });

        throw result.error;
      }

      await this.runStore.updateBatchTaskRun({
        where: { id: batch.id },
        data: { sealed: true, sealedAt: new Date() },
        select: { id: true },
      });

      return batch;
    } else {
      const batch = await this.runStore.createBatchTaskRun({
        id: batchInternalId,
        friendlyId: batchId,
        runtimeEnvironmentId: environment.id,
        idempotencyKey: options.idempotencyKey,
        idempotencyKeyExpiresAt: options.idempotencyKeyExpiresAt,
        dependentTaskAttemptId: dependentAttempt?.id,
        runCount: body.items.length,
        runIds: runs.map((r) => r.id),
        payload: payloadPacket.data,
        payloadType: payloadPacket.dataType,
        options,
        batchVersion: "v3",
        oneTimeUseToken: options.oneTimeUseToken,
      });

      switch (this._batchProcessingStrategy) {
        case "sequential": {
          await this.#enqueueBatchTaskRun({
            batchId: batch.id,
            processingId: batchId,
            range: { start: 0, count: PROCESSING_BATCH_SIZE },
            attemptCount: 0,
            strategy: this._batchProcessingStrategy,
          });

          break;
        }
        case "parallel": {
          const ranges = Array.from({
            length: Math.ceil(newRunCount / PROCESSING_BATCH_SIZE),
          }).map((_, index) => ({
            start: index * PROCESSING_BATCH_SIZE,
            count: PROCESSING_BATCH_SIZE,
          }));

          await this.runStore.updateBatchTaskRun({
            where: { id: batch.id },
            data: {
              processingJobsExpectedCount: ranges.length,
            },
            select: { id: true },
          });

          await Promise.all(
            ranges.map((range, index) =>
              this.#enqueueBatchTaskRun({
                batchId: batch.id,
                processingId: `${index}`,
                range,
                attemptCount: 0,
                strategy: this._batchProcessingStrategy,
              })
            )
          );

          break;
        }
      }

      return batch;
    }
  }

  async #respondWithExistingBatch(
    batch: BatchTaskRun,
    environment: AuthenticatedEnvironment
  ): Promise<BatchTriggerTaskV2Response> {
    // Resolve the payload
    const payloadPacket = await downloadPacketFromObjectStore(
      {
        data: batch.payload ?? undefined,
        dataType: batch.payloadType,
      },
      environment
    );

    const payload = await parsePacket(payloadPacket).then(
      (p) => p as BatchTriggerTaskV2RequestBody["items"]
    );

    const runs = batch.runIds.map((id, index) => {
      const item = payload[index];

      return {
        id,
        taskIdentifier: item.task,
        isCached: true,
        idempotencyKey: item.options?.idempotencyKey ?? undefined,
      };
    });

    return {
      id: batch.friendlyId,
      idempotencyKey: batch.idempotencyKey ?? undefined,
      isCached: true,
      runs,
    };
  }

  async processBatchTaskRun(options: BatchProcessingOptions) {
    logger.debug("[BatchTriggerV2][processBatchTaskRun] Processing batch", {
      options,
    });

    const $attemptCount = options.attemptCount + 1;

    if ($attemptCount > MAX_ATTEMPTS) {
      logger.error("[BatchTriggerV2][processBatchTaskRun] Max attempts reached", {
        options,
        attemptCount: $attemptCount,
      });
      return;
    }

    const batch = await this.runStore.findBatchTaskRunById(options.batchId);

    if (!batch) {
      return;
    }

    // BatchTaskRun -> RuntimeEnvironment FK is dropped; resolve the env from the scalar id.
    const environment = await findEnvironmentById(batch.runtimeEnvironmentId);
    if (!environment) {
      logger.error("[BatchTriggerV2][processBatchTaskRun] Environment not found", {
        batchId: batch.id,
        runtimeEnvironmentId: batch.runtimeEnvironmentId,
      });
      return;
    }

    if (options.range.start >= batch.runCount) {
      logger.debug("[BatchTriggerV2][processBatchTaskRun] currentIndex is greater than runCount", {
        options,
        batchId: batch.friendlyId,
        runCount: batch.runCount,
        attemptCount: $attemptCount,
      });

      return;
    }

    // Resolve the payload
    const payloadPacket = await downloadPacketFromObjectStore(
      {
        data: batch.payload ?? undefined,
        dataType: batch.payloadType,
      },
      environment
    );

    const payload = await parsePacket(payloadPacket);

    if (!payload) {
      logger.debug("[BatchTriggerV2][processBatchTaskRun] Failed to parse payload", {
        options,
        batchId: batch.friendlyId,
        attemptCount: $attemptCount,
      });

      throw new Error("Failed to parse payload");
    }

    // Skip zod parsing
    const $payload = payload as BatchTriggerTaskV2RequestBody["items"];
    const $options = batch.options as BatchTriggerTaskServiceOptions;

    const result = await this.#processBatchTaskRunItems(
      batch,
      environment,
      options.range.start,
      options.range.count,
      $payload,
      $options
    );

    if (result.error) {
      logger.error("[BatchTriggerV2][processBatchTaskRun] Batch processing error", {
        batchId: batch.friendlyId,
        currentIndex: result.workingIndex,
        error: result.error,
        attemptCount: $attemptCount,
      });

      // if the strategy is sequential, we will requeue processing with a count of the PROCESSING_BATCH_SIZE
      // if the strategy is parallel, we will requeue processing with a range starting at the workingIndex and a count that is the remainder of this "slice" of the batch
      await this.#enqueueBatchTaskRun({
        batchId: batch.id,
        processingId: options.processingId,
        range: {
          start: result.workingIndex,
          count:
            options.strategy === "sequential"
              ? options.range.count
              : options.range.count - result.workingIndex - options.range.start,
        },
        attemptCount: $attemptCount,
        strategy: options.strategy,
      });

      return;
    }

    switch (options.strategy) {
      case "sequential": {
        // Done once we've walked past the last item in the batch
        if (result.workingIndex >= batch.runCount) {
          await this.runStore.updateBatchTaskRun({
            where: { id: batch.id },
            data: { sealed: true, sealedAt: new Date() },
            select: { id: true },
          });

          logger.debug("[BatchTriggerV2][processBatchTaskRun] Batch processing complete", {
            batchId: batch.friendlyId,
            runCount: batch.runCount,
            currentIndex: result.workingIndex,
            attemptCount: $attemptCount,
          });
        } else {
          await this.#enqueueBatchTaskRun({
            batchId: batch.id,
            processingId: options.processingId,
            range: {
              start: result.workingIndex,
              count: options.range.count,
            },
            attemptCount: 0,
            strategy: options.strategy,
          });
        }

        break;
      }
      case "parallel": {
        // Each processing job increments the count; the last one to arrive seals the batch
        const { processingJobsCount, processingJobsExpectedCount } =
          await this.runStore.updateBatchTaskRun({
            where: { id: batch.id },
            data: {
              processingJobsCount: {
                increment: 1,
              },
            },
            select: {
              processingJobsExpectedCount: true,
              processingJobsCount: true,
            },
          });

        if (processingJobsCount >= processingJobsExpectedCount) {
          await this.runStore.updateBatchTaskRun({
            where: { id: batch.id },
            data: { sealed: true, sealedAt: new Date() },
            select: { id: true },
          });

          logger.debug("[BatchTriggerV2][processBatchTaskRun] Batch processing complete", {
            batchId: batch.friendlyId,
            currentIndex: result.workingIndex,
            attemptCount: $attemptCount,
          });
        }
      }
    }
  }

  async #processBatchTaskRunItems(
    batch: BatchTaskRun,
    environment: AuthenticatedEnvironment,
    currentIndex: number,
    batchSize: number,
    items: BatchTriggerTaskV2RequestBody["items"],
    options?: BatchTriggerTaskServiceOptions
  ): Promise<{ workingIndex: number; error?: Error }> {
    const runIds = batch.runIds.slice(currentIndex, currentIndex + batchSize);

    logger.debug("[BatchTriggerV2][processBatchTaskRun] Processing batch items", {
      batchId: batch.friendlyId,
      currentIndex,
      runIds,
      runCount: batch.runCount,
    });

    // Pair each runId in this window with its item from the payload array
    const itemsToProcess = runIds.map((runId, index) => ({
      runId,
      item: items[index + currentIndex],
    }));

    let workingIndex = currentIndex;
    let expectedCount = 0;

    for (const item of itemsToProcess) {
      try {
        const created = await this.#processBatchTaskRunItem(
          batch,
          environment,
          item,
          workingIndex,
          options
        );

        if (created) {
          expectedCount++;
        }

        workingIndex++;
      } catch (error) {
        logger.error("[BatchTriggerV2][processBatchTaskRun] Failed to process item", {
          batchId: batch.friendlyId,
          currentIndex: workingIndex,
          error,
        });

        return {
          error: error instanceof Error ? error : new Error(String(error)),
          workingIndex,
        };
      }
    }

    if (expectedCount > 0) {
      await this.runStore.updateBatchTaskRun({
        where: { id: batch.id },
        data: {
          expectedCount: {
            increment: expectedCount,
          },
        },
        select: { id: true },
      });
    }

    return { workingIndex };
  }

  async #processBatchTaskRunItem(
    batch: BatchTaskRun,
    environment: AuthenticatedEnvironment,
    task: { runId: string; item: BatchTriggerTaskV2RequestBody["items"][number] },
    currentIndex: number,
    options?: BatchTriggerTaskServiceOptions
  ) {
    logger.debug("[BatchTriggerV2][processBatchTaskRunItem] Processing item", {
      batchId: batch.friendlyId,
      runId: task.runId,
      currentIndex,
    });

    const triggerTaskService = new TriggerTaskService();

    const result = await triggerTaskService.call(
      task.item.task,
      environment,
      {
        ...task.item,
        options: {
          ...task.item.options,
          dependentBatch: batch.dependentTaskAttemptId ? batch.friendlyId : undefined, // Only set dependentBatch if dependentAttempt is set which means batchTriggerAndWait was called
          parentBatch: batch.dependentTaskAttemptId ? undefined : batch.friendlyId, // Only set parentBatch if dependentAttempt is NOT set which means batchTrigger was called
        },
      },
      {
        triggerVersion: options?.triggerVersion,
        traceContext: options?.traceContext,
        spanParentAsLink: options?.spanParentAsLink,
        batchId: batch.id,
        skipChecks: true,
        runFriendlyId: task.runId,
        realtimeStreamsVersion: options?.realtimeStreamsVersion,
        triggerSource: options?.triggerSource ?? "api",
        triggerAction: options?.triggerAction ?? "trigger",
      }
    );

    if (!result) {
      throw new Error(`Failed to trigger run ${task.runId} for batch ${batch.friendlyId}`);
    }

    if (!result.isCached) {
      try {
        await this.runStore.createBatchTaskRunItem({
          batchTaskRunId: batch.id,
          taskRunId: result.run.id,
          status: batchTaskRunItemStatusForRunStatus(result.run.status),
        });

        return true;
      } catch (error) {
        if (isUniqueConstraintError(error, ["batchTaskRunId", "taskRunId"])) {
          // This means there is already a batchTaskRunItem for this batch and taskRun
          logger.debug(
            "[BatchTriggerV2][processBatchTaskRunItem] BatchTaskRunItem already exists",
            {
              batchId: batch.friendlyId,
              runId: task.runId,
              currentIndex,
            }
          );

          return false;
        }

        throw error;
      }
    }

    return false;
  }

  async #enqueueBatchTaskRun(options: BatchProcessingOptions) {
    await batchTriggerWorker.enqueue({
      id: `BatchTriggerV2Service.process:${options.batchId}:${options.processingId}`,
      job: "v3.processBatchTaskRun",
      payload: options,
    });
  }

  async #handlePayloadPacket(
    payload: any,
    pathPrefix: string,
    environment: AuthenticatedEnvironment
  ) {
    return await startActiveSpan("handlePayloadPacket()", async (span) => {
      const packet = { data: JSON.stringify(payload), dataType: "application/json" };

      if (!packet.data) {
        return packet;
      }

      const { needsOffloading } = packetRequiresOffloading(
        packet,
        env.TASK_PAYLOAD_OFFLOAD_THRESHOLD
      );

      if (!needsOffloading) {
        return packet;
      }

      const filename = `${pathPrefix}/payload.json`;

      const uploadedFilename = await uploadPacketToObjectStore(
        filename,
        packet.data,
        packet.dataType,
        environment
      );

      return {
        data: uploadedFilename,
        dataType: "application/store",
      };
    });
  }
}

export async function tryCompleteBatchV3(
  batchId: string,
  tx: PrismaClientOrTransaction,
  scheduleResumeOnComplete: boolean,
  // Threaded in so a run-ops id (NEW-resident) batch + its items are read/written on the owning
  // store, not the control-plane `tx`. Defaults to the singleton (single-DB = passthrough).
  runStore: RunStore = defaultRunStore
) {
  const batch = await runStore.findBatchTaskRunById(batchId);

  if (!batch) {
    logger.debug("tryCompleteBatchV3: Batch not found", { batchId });
    return;
  }

  if (batch.status === "COMPLETED") {
    logger.debug("tryCompleteBatchV3: Already completed", { batchId });
    return;
  }

  if (!batch.sealed) {
    logger.debug("tryCompleteBatchV3: Not sealed yet", { batchId });
    return;
  }

  const completedCount = await runStore.countBatchTaskRunItems({
    batchTaskRunId: batchId,
    status: "COMPLETED",
  });

  if (completedCount < batch.expectedCount) {
    logger.debug("tryCompleteBatchV3: Not all items completed", {
      batchId,
      completedCount,
      expectedCount: batch.expectedCount,
    });
    return;
  }

  // Mark batch COMPLETED (idempotent via status check)
  const updated = await runStore.updateManyBatchTaskRun({
    where: { id: batchId, status: "PENDING" },
    data: { status: "COMPLETED", completedAt: new Date(), completedCount },
  });

  if (updated.count === 0) {
    logger.debug("tryCompleteBatchV3: Already transitioned", { batchId });
    return;
  }

  logger.debug("tryCompleteBatchV3: Batch completed", { batchId, completedCount });

  if (scheduleResumeOnComplete && batch.dependentTaskAttemptId) {
    await ResumeBatchRunService.enqueue(batchId, true, tx);
  }
}

export async function completeBatchTaskRunItemV3(
  itemId: string,
  batchTaskRunId: string,
  tx: PrismaClientOrTransaction,
  scheduleResumeOnComplete = false,
  taskRunAttemptId?: string,
  retryAttempt?: number,
  // Threaded in so a run-ops id (NEW-resident) batch's item lands on the owning store; route by
  // batchTaskRunId (items co-reside with their batch). Defaults to the singleton.
  runStore: RunStore = defaultRunStore
) {
  const isRetry = retryAttempt !== undefined;

  logger.debug("completeBatchTaskRunItemV3", {
    itemId,
    batchTaskRunId,
    scheduleResumeOnComplete,
    taskRunAttemptId,
    retryAttempt,
    isRetry,
  });

  try {
    // Update item to COMPLETED (no transaction needed, no contention). Routed by
    // batchTaskRunId so the item write lands on the batch's owning DB.
    const updated = await runStore.updateManyBatchTaskRunItems({
      where: { id: itemId, batchTaskRunId, status: "PENDING" },
      data: { status: "COMPLETED", taskRunAttemptId },
    });

    if (updated.count === 0) {
      logger.debug("completeBatchTaskRunItemV3: Item already completed", {
        itemId,
        batchTaskRunId,
      });
      return;
    }

    // Schedule debounced completion check
    // enqueue with same ID overwrites, resetting the 200ms timer (debounce behavior)
    await legacyRunEngineWorker.enqueue({
      id: `tryCompleteBatchV3:${batchTaskRunId}`,
      job: "tryCompleteBatchV3",
      payload: { batchId: batchTaskRunId, scheduleResumeOnComplete },
      availableAt: new Date(Date.now() + 200),
    });
  } catch (error) {
    if (isPrismaRetriableError(error) || isPrismaRaceConditionError(error)) {
      logger.error("completeBatchTaskRunItemV3 failed, scheduling retry", {
        itemId,
        batchTaskRunId,
        error,
        retryAttempt,
        isRetry,
      });

      if (isRetry) {
        throw error;
      } else {
        await legacyRunEngineWorker.enqueue({
          id: `completeBatchTaskRunItem:${itemId}`,
          job: "completeBatchTaskRunItem",
          payload: { itemId, batchTaskRunId, scheduleResumeOnComplete, taskRunAttemptId },
          availableAt: new Date(Date.now() + 2_000),
        });
      }
    } else {
      logger.error("completeBatchTaskRunItemV3 failed with non-retriable error", {
        itemId,
        batchTaskRunId,
        error,
        retryAttempt,
        isRetry,
      });
    }
  }
}