import {
  type MachinePreset,
  prettyPrintPacket,
  RunAnnotations,
  SemanticInternalAttributes,
  type TaskRunContext,
  TaskRunError,
  TriggerTraceContext,
  type V3TaskRunContext,
} from "@trigger.dev/core/v3";

import { AttemptId, getMaxDuration, parseTraceparent } from "@trigger.dev/core/v3/isomorphic";
import { $replica, runOpsNewReplica, runOpsSplitReadEnabled } from "~/db.server";
import {
  extractIdempotencyKeyScope,
  getUserProvidedIdempotencyKey,
} from "@trigger.dev/core/v3/serverOnly";
import {
  extractAIEmbedData,
  extractAISpanData,
  extractAISummarySpanData,
  extractAIToolCallData,
} from "~/components/runs/v3/ai";
import { RUNNING_STATUSES } from "~/components/runs/v3/TaskRunStatus";
import { baseWorkerQueue } from "~/runEngine/concerns/workerQueueSplit.server";
import { logger } from "~/services/logger.server";
import { safeJsonParse } from "~/utils/json";
import { rehydrateAttribute } from "~/v3/eventRepository/eventRepository.server";
import type { IEventRepository } from "~/v3/eventRepository/eventRepository.types";
import { getEventRepositoryForStore } from "~/v3/eventRepository/index.server";
import { machinePresetFromRun } from "~/v3/machinePresets.server";
import { findRunByIdWithMollifierFallback } from "~/v3/mollifier/readFallback.server";
import { buildSyntheticSpanRun } from "~/v3/mollifier/syntheticSpanRun.server";
import { engine } from "~/v3/runEngine.server";
import { runStore } from "~/v3/runStore.server";
import { getTaskEventStoreTableForRun, type TaskEventStoreTable } from "~/v3/taskEventStore.server";
import { isFailedRunStatus, isFinalRunStatus } from "~/v3/taskStatus";
import { BasePresenter } from "./basePresenter.server";
import { WaitpointPresenter } from "./WaitpointPresenter.server";
import {
  controlPlaneResolver,
  type ResolvedRunLockedWorker,
} from "~/v3/runOpsMigration/controlPlaneResolver.server";
import type { AuthenticatedEnvironment } from "@trigger.dev/core/v3/auth/environment";

export type PromptSpanData = {
  slug: string;
  version: number;
  labels: string;
  model?: string;
  template?: string;
  text?: string;
  input?: string;
  config?: string;
};

function extractPromptSpanData(properties: Record<string, unknown>): PromptSpanData | undefined {
  // Properties come as an unflattened nested object from ClickHouse,
  // e.g. { prompt: { slug: "...", version: 3, ... } }
  const prompt = properties.prompt;
  if (!prompt || typeof prompt !== "object") {
    return undefined;
  }

  const p = prompt as Record<string, unknown>;
  const slug = p.slug;
  const version = p.version;

  if (typeof slug !== "string" || typeof version !== "number") {
    return undefined;
  }

  return {
    slug,
    version,
    labels: typeof p.labels === "string" ? p.labels : "",
    model: typeof p.model === "string" ? p.model : undefined,
    template: typeof p.template === "string" ? p.template : undefined,
    text: typeof p.text === "string" ? p.text : undefined,
    input: typeof p.input === "string" ? p.input : undefined,
    config: typeof p.config === "string" ? p.config : undefined,
  };
}

// Grounded in `getRun` (the canonical shape source), not inferred from `call`:
// `call`'s buffered branch returns `buildSyntheticSpanRun` which is annotated
// `Promise<SpanRun>`, so deriving SpanRun from `call` would be a circular type
// reference TS rejects. `getRun` doesn't recurse, breaking the cycle.
export type SpanRun = NonNullable<
  Awaited<ReturnType<InstanceType<typeof SpanPresenter>["getRun"]>>
>;
type Result = Awaited<ReturnType<SpanPresenter["call"]>>;
export type Span = NonNullable<NonNullable<Result>["span"]>;
type FindRunResult = NonNullable<
  Awaited<ReturnType<InstanceType<typeof SpanPresenter>["findRun"]>>
>;

// Run-ops TaskRun reads (parent run in `call`, hydrate in `findRun`, children in
// `#getSpan`) go through the `runStore` seam; split routing is the RoutingRunStore's
// job below it. Control-plane reads stay on `this._replica`/`this._prisma`.
export class SpanPresenter extends BasePresenter {
  public async call({
    userId,
    projectSlug,
    envSlug,
    spanId,
    runFriendlyId,
    linkedRunId,
  }: {
    userId: string;
    projectSlug: string;
    // Optional for backwards compatibility, required for the mollifier
    // buffer fallback when the parent run isn't yet in PG — we need to
    // resolve the env id to satisfy `findRunByIdWithMollifierFallback`'s
    // auth check.
    envSlug?: string;
    spanId: string;
    runFriendlyId: string;
    linkedRunId?: string;
  }) {
    const project = await this._replica.project.findFirst({
      where: {
        slug: projectSlug,
        organization: {
          members: {
            some: {
              userId,
            },
          },
        },
      },
    });

    if (!project) {
      throw new Error("Project not found");
    }

    const parentRun = await runStore.findRun(
      {
        friendlyId: runFriendlyId,
        projectId: project.id,
      },
      {
        select: {
          traceId: true,
          runtimeEnvironmentId: true,
          projectId: true,
          taskEventStore: true,
          createdAt: true,
          completedAt: true,
        },
      },
      this._prisma
    );

    if (!parentRun) {
      // PG miss → fall back to the mollifier buffer. Without this the
      // right-side span detail panel on the run-detail page never
      // resolves for buffered runs: `call()` returns undefined, the
      // resource route redirects with an "Event not found" toast, the
      // run-detail page reloads, the toast fires again — a perpetual
      // spin until the drainer materialises the row. Synthesise a
      // SpanRun straight from the buffer snapshot, reusing
      // `buildSyntheticSpanRun` (the same helper the run-detail
      // loader's header fallback already uses).
      if (!envSlug) return;
      const envRow = await this._replica.runtimeEnvironment.findFirst({
        where: { project: { id: project.id }, slug: envSlug },
        select: { id: true, slug: true, type: true, organizationId: true },
      });
      if (!envRow) return;
      const buffered = await findRunByIdWithMollifierFallback({
        runId: runFriendlyId,
        environmentId: envRow.id,
        organizationId: envRow.organizationId,
      });
      if (!buffered) return;
      const synth = await buildSyntheticSpanRun({
        run: buffered,
        environment: { id: envRow.id, slug: envRow.slug, type: envRow.type },
      });
      return { type: "run" as const, run: synth };
    }

    const { traceId } = parentRun;

    const repository = await getEventRepositoryForStore(
      parentRun.taskEventStore,
      project.organizationId
    );

    const eventStore = getTaskEventStoreTableForRun(parentRun);

    const run = await this.getRun({
      eventStore,
      traceId,
      eventRepository: repository,
      spanId,
      linkedRunId,
      createdAt: parentRun.createdAt,
      completedAt: parentRun.completedAt,
      environmentId: parentRun.runtimeEnvironmentId,
    });
    if (run) {
      return {
        type: "run" as const,
        run,
      };
    }

    const span = await this.#getSpan({
      eventStore,
      spanId,
      traceId,
      environmentId: parentRun.runtimeEnvironmentId,
      projectId: parentRun.projectId,
      createdAt: parentRun.createdAt,
      completedAt: parentRun.completedAt,
      eventRepository: repository,
    });

    if (!span) {
      throw new Error("Span not found");
    }

    return {
      type: "span" as const,
      span,
    };
  }

  async getRun({
    eventStore,
    environmentId,
    traceId,
    eventRepository,
    spanId,
    linkedRunId,
    createdAt,
    completedAt,
  }: {
    eventStore: TaskEventStoreTable;
    environmentId: string;
    traceId: string;
    eventRepository: IEventRepository;
    spanId: string;
    linkedRunId?: string;
    createdAt: Date;
    completedAt: Date | null;
  }) {
    // Use linkedRunId if provided (for cached spans), otherwise look up by spanId
    const run = await this.findRun({ originalRunId: linkedRunId, spanId, environmentId });

    if (!run) {
      return;
    }

    const environment = await controlPlaneResolver.resolveAuthenticatedEnv(
      run.runtimeEnvironmentId
    );

    if (!environment) {
      return undefined;
    }

    const lockedWorker = await controlPlaneResolver.resolveRunLockedWorker({
      lockedById: run.lockedById,
      lockedToVersionId: run.lockedToVersionId,
    });

    const isFinished = isFinalRunStatus(run.status);
    const output = !isFinished
      ? undefined
      : run.outputType === "application/store"
        ? `/resources/packets/${environment.id}/${run.output}`
        : typeof run.output !== "undefined" && run.output !== null
          ? await prettyPrintPacket(run.output, run.outputType ?? undefined)
          : undefined;

    const payload =
      run.payloadType === "application/store"
        ? `/resources/packets/${environment.id}/${run.payload}`
        : typeof run.payload !== "undefined" && run.payload !== null
          ? await prettyPrintPacket(run.payload, run.payloadType ?? undefined)
          : undefined;

    let error: TaskRunError | undefined = undefined;

    if (run?.error) {
      const result = TaskRunError.safeParse(run.error);
      if (result.success) {
        error = result.data;
      } else {
        error = {
          type: "CUSTOM_ERROR",
          raw: JSON.stringify(run.error),
        };
      }
    }

    const metadata = run.metadata
      ? await prettyPrintPacket(run.metadata, run.metadataType, {
          filteredKeys: ["$$streams", "$$streamsVersion", "$$streamsBaseUrl"],
        })
      : undefined;

    const machine = run.machinePreset ? machinePresetFromRun(run) : undefined;

    const context = await this.#getTaskRunContext({
      run,
      machine: machine ?? undefined,
      environment,
      lockedWorker,
    });

    const externalTraceId = this.#getExternalTraceId(run.traceContext);

    const taskKind = RunAnnotations.safeParse(run.annotations).data?.taskKind;
    const isAgentRun = taskKind === "AGENT";
    const isScheduled = taskKind === "SCHEDULED";

    let region: { name: string; location: string | null } | null = null;

    if (environment.type !== "DEVELOPMENT" && run.engine !== "V1") {
      const workerGroup = await this._replica.workerInstanceGroup.findFirst({
        select: {
          name: true,
          location: true,
        },
        where: {
          // masterQueue is unique and IS the run's backing queue, so this finds
          // the group the run actually ran on.
          masterQueue: baseWorkerQueue(run.workerQueue),
        },
      });

      // Show the stamped geo region as the name so a migrated run never reveals
      // its compute backing; fall back to the group name for unstamped runs.
      region = workerGroup
        ? { name: run.region ?? workerGroup.name, location: workerGroup.location }
        : null;
    }

    // Only AGENT-tagged runs can be session-bound, so skip the SessionRun lookup
    // for the much larger set of standard runs — the cheapest query is the one we
    // don't run.
    const sessionRun = isAgentRun
      ? await this._replica.sessionRun.findFirst({
          where: { runId: run.id },
          select: {
            reason: true,
            triggeredAt: true,
            session: {
              select: {
                friendlyId: true,
                externalId: true,
                type: true,
                taskIdentifier: true,
                closedAt: true,
                expiresAt: true,
              },
            },
          },
        })
      : null;

    const session = sessionRun
      ? {
          friendlyId: sessionRun.session.friendlyId,
          externalId: sessionRun.session.externalId,
          type: sessionRun.session.type,
          taskIdentifier: sessionRun.session.taskIdentifier,
          status:
            sessionRun.session.closedAt != null
              ? ("CLOSED" as const)
              : sessionRun.session.expiresAt != null &&
                  sessionRun.session.expiresAt.getTime() < Date.now()
                ? ("EXPIRED" as const)
                : ("ACTIVE" as const),
          reason: sessionRun.reason,
          triggeredAt: sessionRun.triggeredAt,
        }
      : undefined;

    return {
      id: run.id,
      friendlyId: run.friendlyId,
      status: run.status,
      statusReason: run.statusReason ?? undefined,
      createdAt: run.createdAt,
      startedAt: run.startedAt,
      executedAt: run.executedAt,
      updatedAt: run.updatedAt,
      delayUntil: run.delayUntil,
      expiredAt: run.expiredAt,
      completedAt: run.completedAt,
      logsDeletedAt: run.logsDeletedAt,
      ttl: run.ttl,
      taskIdentifier: run.taskIdentifier,
      version: lockedWorker?.lockedToVersion?.version,
      sdkVersion: lockedWorker?.lockedToVersion?.sdkVersion,
      runtime: lockedWorker?.lockedToVersion?.runtime,
      runtimeVersion: lockedWorker?.lockedToVersion?.runtimeVersion,
      isTest: run.isTest,
      replayedFromTaskRunFriendlyId: run.replayedFromTaskRunFriendlyId,
      environmentId: environment.id,
      idempotencyKey: getUserProvidedIdempotencyKey(run),
      idempotencyKeyExpiresAt: run.idempotencyKeyExpiresAt,
      idempotencyKeyScope: extractIdempotencyKeyScope(run),
      idempotencyKeyStatus: this.getIdempotencyKeyStatus(run),
      debounce: run.debounce as { key: string; delay: string; createdAt: Date } | null,
      schedule: await this.resolveSchedule(run.scheduleId ?? undefined),
      queue: {
        name: run.queue,
        isCustomQueue: !run.queue.startsWith("task/"),
        concurrencyKey: run.concurrencyKey,
      },
      tags: run.runTags,
      baseCostInCents: run.baseCostInCents,
      costInCents: run.costInCents,
      totalCostInCents: run.costInCents + run.baseCostInCents,
      usageDurationMs: run.usageDurationMs,
      isFinished,
      isRunning: RUNNING_STATUSES.includes(run.status),
      isError: isFailedRunStatus(run.status),
      isAgentRun,
      isScheduled,
      payload,
      payloadType: run.payloadType,
      output,
      outputType: run?.outputType ?? "application/json",
      error,
      relationships: {
        root: run.rootTaskRun
          ? {
              ...run.rootTaskRun,
              isParent: run.parentTaskRun?.friendlyId === run.rootTaskRun.friendlyId,
            }
          : undefined,
        parent: run.parentTaskRun ?? undefined,
      },
      context: JSON.stringify(context, null, 2),
      metadata,
      maxDurationInSeconds: getMaxDuration(run.maxDurationInSeconds),
      batch: run.batch ? { friendlyId: run.batch.friendlyId } : undefined,
      session,
      engine: run.engine,
      region,
      workerQueue: run.workerQueue,
      traceId: run.traceId,
      spanId: run.spanId,
      isCached: !!linkedRunId,
      isBuffered: false,
      machinePreset: machine?.name,
      taskEventStore: run.taskEventStore,
      externalTraceId,
    };
  }

  private getIdempotencyKeyStatus(run: {
    idempotencyKey: string | null;
    idempotencyKeyExpiresAt: Date | null;
    idempotencyKeyOptions: unknown;
  }): "active" | "inactive" | "expired" | undefined {
    // No idempotency configured if no scope exists
    const scope = extractIdempotencyKeyScope(run);
    if (!scope) {
      return undefined;
    }

    // Check if expired first (takes precedence)
    if (run.idempotencyKeyExpiresAt && run.idempotencyKeyExpiresAt < new Date()) {
      return "expired";
    }

    // Check if reset (hash is null but options exist)
    if (run.idempotencyKey === null) {
      return "inactive";
    }

    return "active";
  }

  async resolveSchedule(scheduleId?: string) {
    if (!scheduleId) {
      return;
    }

    const schedule = await this._replica.taskSchedule.findFirst({
      where: {
        id: scheduleId,
      },
      select: {
        friendlyId: true,
        generatorExpression: true,
        timezone: true,
        generatorDescription: true,
      },
    });

    if (!schedule) {
      return;
    }

    return {
      friendlyId: schedule.friendlyId,
      generatorExpression: schedule.generatorExpression,
      description: schedule.generatorDescription,
      timezone: schedule.timezone,
    };
  }

  async findRun({
    originalRunId,
    spanId,
    environmentId,
  }: {
    originalRunId?: string;
    spanId: string;
    environmentId: string;
  }) {
    const run = await runStore.findRun(
      originalRunId
        ? {
            friendlyId: originalRunId,
            runtimeEnvironmentId: environmentId,
          }
        : {
            spanId,
            runtimeEnvironmentId: environmentId,
          },
      {
        select: {
          id: true,
          runtimeEnvironmentId: true,
          lockedById: true,
          lockedToVersionId: true,
          spanId: true,
          traceId: true,
          traceContext: true,
          //metadata
          number: true,
          taskIdentifier: true,
          friendlyId: true,
          isTest: true,
          maxDurationInSeconds: true,
          taskEventStore: true,
          runTags: true,
          machinePreset: true,
          engine: true,
          workerQueue: true,
          region: true,
          error: true,
          output: true,
          outputType: true,
          //status + duration
          status: true,
          statusReason: true,
          startedAt: true,
          executedAt: true,
          createdAt: true,
          updatedAt: true,
          queuedAt: true,
          completedAt: true,
          logsDeletedAt: true,
          //idempotency
          idempotencyKey: true,
          idempotencyKeyExpiresAt: true,
          idempotencyKeyOptions: true,
          //debounce
          debounce: true,
          //delayed
          delayUntil: true,
          //ttl
          ttl: true,
          expiredAt: true,
          //queue
          queue: true,
          concurrencyKey: true,
          //schedule
          scheduleId: true,
          //usage
          baseCostInCents: true,
          costInCents: true,
          usageDurationMs: true,
          payload: true,
          payloadType: true,
          metadata: true,
          metadataType: true,
          annotations: true,
          maxAttempts: true,
          //relationships
          rootTaskRun: {
            select: {
              taskIdentifier: true,
              friendlyId: true,
              spanId: true,
              createdAt: true,
            },
          },
          parentTaskRun: {
            select: {
              taskIdentifier: true,
              friendlyId: true,
              spanId: true,
            },
          },
          batch: {
            select: {
              friendlyId: true,
            },
          },
          replayedFromTaskRunFriendlyId: true,
          attempts: {
            take: 1,
            orderBy: {
              createdAt: "desc",
            },
            select: {
              number: true,
              status: true,
              createdAt: true,
              friendlyId: true,
            },
          },
        },
      },
      this._replica
    );

    return run;
  }

  async #getSpan({
    eventStore,
    eventRepository,
    traceId,
    spanId,
    environmentId,
    projectId,
    createdAt,
    completedAt,
  }: {
    eventRepository: IEventRepository;
    traceId: string;
    spanId: string;
    environmentId: string;
    projectId: string;
    eventStore: TaskEventStoreTable;
    createdAt: Date;
    completedAt: Date | null;
  }) {
    const span = await eventRepository.getSpan(
      eventStore,
      environmentId,
      spanId,
      traceId,
      createdAt,
      completedAt ?? undefined,
      { includeDebugLogs: true }
    );

    if (!span) {
      return;
    }

    const triggeredRuns = await runStore.findRuns(
      {
        where: {
          parentSpanId: spanId,
        },
        select: {
          friendlyId: true,
          taskIdentifier: true,
          spanId: true,
          createdAt: true,
          status: true,
        },
      },
      this._replica
    );

    const data = {
      spanId: span.spanId,
      parentId: span.parentId,
      message: span.message,
      isError: span.isError,
      isPartial: span.isPartial,
      isCancelled: span.isCancelled,
      level: span.level,
      startTime: span.startTime,
      duration: span.duration,
      events: span.events,
      style: span.style,
      properties:
        span.properties &&
        typeof span.properties === "object" &&
        Object.keys(span.properties).length > 0
          ? JSON.stringify(span.properties, null, 2)
          : undefined,
      resourceProperties:
        span.resourceProperties &&
        typeof span.resourceProperties === "object" &&
        Object.keys(span.resourceProperties).length > 0
          ? JSON.stringify(span.resourceProperties, null, 2)
          : undefined,
      entity: span.entity,
      metadata: span.metadata,
      triggeredRuns,
      aiData:
        span.properties && typeof span.properties === "object"
          ? extractAISpanData(span.properties as Record<string, unknown>, span.duration / 1_000_000)
          : undefined,
    };

    switch (span.entity.type) {
      case "waitpoint": {
        if (!span.entity.id) {
          logger.error(`SpanPresenter: No waitpoint id`, {
            spanId,
            waitpointFriendlyId: span.entity.id,
          });
          return { ...data, entity: null };
        }

        const presenter = new WaitpointPresenter(undefined, undefined, {
          newClient: runOpsNewReplica,
          legacyReplica: $replica,
          splitEnabled: runOpsSplitReadEnabled,
        });
        const waitpoint = await presenter.call({
          friendlyId: span.entity.id,
          environmentId,
          projectId,
        });

        if (!waitpoint) {
          logger.error(`SpanPresenter: Waitpoint not found`, {
            spanId,
            waitpointFriendlyId: span.entity.id,
          });
          return { ...data, entity: null };
        }

        return {
          ...data,
          entity: {
            type: "waitpoint" as const,
            object: waitpoint,
          },
        };
      }
      case "attempt": {
        const isWarmStart = rehydrateAttribute<boolean>(
          span.metadata,
          SemanticInternalAttributes.WARM_START
        );

        return {
          ...data,
          entity: {
            type: "attempt" as const,
            object: {
              isWarmStart,
            },
          },
        };
      }
      case "realtime-stream": {
        if (!span.entity.id) {
          logger.error(`SpanPresenter: No realtime stream id`, {
            spanId,
            realtimeStreamId: span.entity.id,
          });
          return { ...data, entity: null };
        }

        const [runId, streamKey] = span.entity.id.split(":");

        if (!runId || !streamKey) {
          logger.error(`SpanPresenter: Invalid realtime stream id`, {
            spanId,
            realtimeStreamId: span.entity.id,
          });
          return { ...data, entity: null };
        }

        const metadata = span.entity.metadata
          ? (safeJsonParse(span.entity.metadata) as Record<string, unknown> | undefined)
          : undefined;

        return {
          ...data,
          entity: {
            type: "realtime-stream" as const,
            object: {
              runId,
              streamKey,
              metadata,
            },
          },
        };
      }
      case "input-stream": {
        if (!span.entity.id) {
          logger.error(`SpanPresenter: No input stream id`, {
            spanId,
            inputStreamId: span.entity.id,
          });
          return { ...data, entity: null };
        }

        const [runId, streamId] = span.entity.id.split(":");

        if (!runId || !streamId) {
          logger.error(`SpanPresenter: Invalid input stream id`, {
            spanId,
            inputStreamId: span.entity.id,
          });
          return { ...data, entity: null };
        }

        // Translate user-facing stream ID to internal S2 stream name
        const s2StreamKey = `$trigger.input:${streamId}`;

        return {
          ...data,
          entity: {
            type: "realtime-stream" as const,
            object: {
              runId,
              streamKey: s2StreamKey,
              displayName: streamId,
              metadata: undefined,
            },
          },
        };
      }
      case "prompt": {
        const promptData = extractPromptSpanData(span.properties as Record<string, unknown>);

        if (promptData) {
          return {
            ...data,
            entity: { type: "prompt" as const, object: promptData },
          };
        }

        return { ...data, entity: null };
      }
      default: {
        // Check for top-level AI SDK parent spans by message name
        const AI_SUMMARY_MESSAGES = [
          "ai.generateText",
          "ai.streamText",
          "ai.generateObject",
          "ai.streamObject",
        ];

        if (typeof span.message === "string" && AI_SUMMARY_MESSAGES.includes(span.message)) {
          const aiSummaryData = extractAISummarySpanData(
            span.properties as Record<string, unknown>,
            span.duration / 1_000_000
          );
          if (aiSummaryData) {
            const promptVersionData = await this.#lookupPromptVersion(
              aiSummaryData.promptSlug,
              aiSummaryData.promptVersion,
              aiSummaryData.promptLabels,
              aiSummaryData.promptInput,
              projectId,
              environmentId
            );
            return {
              ...data,
              entity: {
                type: "ai-summary" as const,
                object: aiSummaryData,
                promptVersionData,
              },
            };
          }
        }

        if (span.message === "ai.toolCall") {
          const toolCallData = extractAIToolCallData(
            span.properties as Record<string, unknown>,
            span.duration / 1_000_000
          );
          if (toolCallData) {
            return {
              ...data,
              entity: { type: "ai-tool-call" as const, object: toolCallData },
            };
          }
        }

        if (span.message === "ai.embed") {
          const embedData = extractAIEmbedData(
            span.properties as Record<string, unknown>,
            span.duration / 1_000_000
          );
          if (embedData) {
            return {
              ...data,
              entity: { type: "ai-embed" as const, object: embedData },
            };
          }
        }

        // Child generation spans (doGenerate/doStream) with gen_ai.* attributes
        if (data.aiData) {
          const promptVersionData = await this.#lookupPromptVersion(
            data.aiData.promptSlug,
            data.aiData.promptVersion,
            data.aiData.promptLabels,
            data.aiData.promptInput,
            projectId,
            environmentId
          );
          return {
            ...data,
            entity: {
              type: "ai-generation" as const,
              object: data.aiData,
              promptVersionData,
            },
          };
        }
        return { ...data, entity: null };
      }
    }
  }

  async #lookupPromptVersion(
    promptSlug: string | undefined,
    promptVersion: string | undefined,
    promptLabels: string | undefined,
    promptInput: string | undefined,
    projectId: string,
    environmentId: string
  ): Promise<PromptSpanData | undefined> {
    if (!promptSlug || !promptVersion) return undefined;

    const prompt = await this._replica.prompt.findUnique({
      where: {
        projectId_runtimeEnvironmentId_slug: {
          projectId,
          runtimeEnvironmentId: environmentId,
          slug: promptSlug,
        },
      },
    });
    if (!prompt) return undefined;

    const version = await this._replica.promptVersion.findUnique({
      where: {
        promptId_version: {
          promptId: prompt.id,
          version: Number(promptVersion),
        },
      },
    });
    if (!version) return undefined;

    return {
      slug: promptSlug,
      version: version.version,
      labels: promptLabels ?? "",
      model: version.model ?? undefined,
      template: version.textContent ?? undefined,
      text: undefined,
      input: promptInput,
    };
  }

  async #getTaskRunContext({
    run,
    machine,
    environment,
    lockedWorker,
  }: {
    run: FindRunResult;
    machine?: MachinePreset;
    environment: AuthenticatedEnvironment;
    lockedWorker: ResolvedRunLockedWorker | null;
  }) {
    if (run.engine === "V1") {
      return this.#getV3TaskRunContext({ run, machine, environment, lockedWorker });
    } else {
      return this.#getV4TaskRunContext({ run });
    }
  }

  async #getV3TaskRunContext({
    run,
    machine,
    environment,
    lockedWorker,
  }: {
    run: FindRunResult;
    machine?: MachinePreset;
    environment: AuthenticatedEnvironment;
    lockedWorker: ResolvedRunLockedWorker | null;
  }): Promise<V3TaskRunContext> {
    const attempt = run.attempts[0];

    const context = {
      attempt: attempt
        ? {
            id: attempt.friendlyId,
            number: attempt.number,
            status: attempt.status,
            startedAt: attempt.createdAt,
          }
        : {
            id: AttemptId.generate().friendlyId,
            number: 1,
            status: "PENDING" as const,
            startedAt: run.updatedAt,
          },
      task: {
        id: run.taskIdentifier,
        filePath: lockedWorker?.lockedBy?.filePath ?? "",
      },
      run: {
        id: run.friendlyId,
        createdAt: run.createdAt,
        tags: run.runTags,
        isTest: run.isTest,
        isReplay: !!run.replayedFromTaskRunFriendlyId,
        idempotencyKey: getUserProvidedIdempotencyKey(run) ?? undefined,
        startedAt: run.startedAt ?? run.createdAt,
        durationMs: run.usageDurationMs,
        costInCents: run.costInCents,
        baseCostInCents: run.baseCostInCents,
        maxAttempts: run.maxAttempts ?? undefined,
        version: lockedWorker?.lockedToVersion?.version,
        maxDuration: run.maxDurationInSeconds ?? undefined,
      },
      queue: {
        name: run.queue,
        id: run.queue,
      },
      environment: {
        id: environment.id,
        slug: environment.slug,
        type: environment.type,
      },
      organization: {
        id: environment.organization.id,
        slug: environment.organization.slug,
        name: environment.organization.title,
      },
      project: {
        id: environment.project.id,
        ref: environment.project.externalRef,
        slug: environment.project.slug,
        name: environment.project.name,
      },
      machine,
    } satisfies V3TaskRunContext;

    return context;
  }

  async #getV4TaskRunContext({ run }: { run: FindRunResult }): Promise<TaskRunContext> {
    return engine.resolveTaskRunContext(run.id);
  }

  #getExternalTraceId(traceContext: unknown) {
    if (!traceContext) {
      return;
    }

    const parsedTraceContext = TriggerTraceContext.safeParse(traceContext);

    if (!parsedTraceContext.success) {
      return;
    }

    const externalTraceparent = parsedTraceContext.data.external?.traceparent;

    if (!externalTraceparent) {
      return;
    }

    const parsedTraceparent = parseTraceparent(externalTraceparent);

    return parsedTraceparent?.traceId;
  }
}