import type { Attributes, Tracer } from "@opentelemetry/api";
import type {
  ExceptionEventProperties,
  SpanEvents,
  TaskEventEnvironment,
  TaskEventStyle,
  TaskRunError,
} from "@trigger.dev/core/v3";
import type {
  Prisma,
  TaskEvent,
  TaskEventKind,
  TaskEventLevel,
  TaskEventStatus,
  TaskRun,
} from "@trigger.dev/database";
import type { MetricsV1Input } from "@internal/clickhouse";
import type { DetailedTraceEvent, TaskEventStoreTable } from "../taskEventStore.server";
export type { ExceptionEventProperties };

// ============================================================================
// Event Creation Types
// ============================================================================

export type LlmMetricsData = {
  genAiSystem: string;
  requestModel: string;
  responseModel: string;
  baseResponseModel: string;
  matchedModelId: string;
  operationId: string;
  finishReason: string;
  costSource: string;
  pricingTierId: string;
  pricingTierName: string;
  inputTokens: number;
  outputTokens: number;
  totalTokens: number;
  usageDetails: Record<string, number>;
  inputCost: number;
  outputCost: number;
  totalCost: number;
  costDetails: Record<string, number>;
  providerCost: number;
  msToFirstChunk: number;
  tokensPerSecond: number;
  metadata: Record<string, string>;
  promptSlug: string;
  promptVersion: number;
};

export type CreateEventInput = Omit<
  Prisma.TaskEventCreateInput,
  | "id"
  | "createdAt"
  | "properties"
  | "metadata"
  | "style"
  | "output"
  | "payload"
  | "serviceName"
  | "serviceNamespace"
  | "tracestate"
  | "projectRef"
  | "runIsTest"
  | "workerId"
  | "queueId"
  | "queueName"
  | "batchId"
  | "taskPath"
  | "taskExportName"
  | "workerVersion"
  | "idempotencyKey"
  | "attemptId"
  | "usageDurationMs"
  | "usageCostInCents"
  | "machinePreset"
  | "machinePresetCpu"
  | "machinePresetMemory"
  | "machinePresetCentsPerMs"
  | "links"
> & {
  properties: Attributes;
  resourceProperties?: Attributes;
  metadata: Attributes | undefined;
  style: Attributes | undefined;
  machineId?: string;
  runTags?: string[];
  /** Side-channel data for LLM cost tracking, populated by enrichCreatableEvents */
  _llmMetrics?: LlmMetricsData;
};

export type CreatableEventKind = TaskEventKind;
export type CreatableEventStatus = TaskEventStatus;

// ============================================================================
// Task Run Types
// ============================================================================

export type CompleteableTaskRun = Pick<
  TaskRun,
  | "friendlyId"
  | "traceId"
  | "spanId"
  | "parentSpanId"
  | "createdAt"
  | "completedAt"
  | "taskIdentifier"
  | "projectId"
  | "runtimeEnvironmentId"
  | "organizationId"
  | "isTest"
>;

// ============================================================================
// Trace and Event Types
// ============================================================================

export type TraceAttributes = Partial<
  Pick<
    CreateEventInput,
    "isError" | "isCancelled" | "isDebug" | "runId" | "metadata" | "properties" | "style"
  >
>;

export type SetAttribute<T extends TraceAttributes> = (key: keyof T, value: T[keyof T]) => void;

export type TraceEventOptions = {
  kind?: CreatableEventKind;
  context?: Record<string, unknown>;
  spanParentAsLink?: boolean;
  spanIdSeed?: string;
  attributes: TraceAttributes;
  environment: TaskEventEnvironment;
  taskSlug: string;
  startTime?: bigint;
  endTime?: Date;
  immediate?: boolean;
};

export type EventBuilder = {
  traceId: string;
  spanId: string;
  setAttribute: SetAttribute<TraceAttributes>;
  stop: () => void;
  failWithError: (error: TaskRunError) => void;
};

export type UpdateEventOptions = {
  attributes: TraceAttributes;
  endTime?: Date;
  immediate?: boolean;
  events?: SpanEvents;
};

// ============================================================================
// Configuration Types
// ============================================================================

export type EventRepoConfig = {
  batchSize: number;
  batchInterval: number;
  retentionInDays: number;
  partitioningEnabled: boolean;
  tracer?: Tracer;
  minConcurrency?: number;
  maxConcurrency?: number;
  maxBatchSize?: number;
  memoryPressureThreshold?: number;
  loadSheddingThreshold?: number;
  loadSheddingEnabled?: boolean;
};

// ============================================================================
// Query Types
// ============================================================================

export type QueryOptions = Prisma.TaskEventWhereInput;

export type TaskEventRecord = TaskEvent;

export type QueriedEvent = Prisma.TaskEventGetPayload<{
  select: {
    spanId: true;
    parentId: true;
    runId: true;
    message: true;
    style: true;
    startTime: true;
    duration: true;
    isError: true;
    isPartial: true;
    isCancelled: true;
    level: true;
    events: true;
    kind: true;
    attemptNumber: true;
  };
}>;

export type PreparedEvent = Omit<QueriedEvent, "events" | "style" | "duration"> & {
  duration: number;
  events: SpanEvents;
  style: TaskEventStyle;
};

export type PreparedDetailedEvent = Omit<DetailedTraceEvent, "events" | "style" | "duration"> & {
  duration: number;
  events: SpanEvents;
  style: TaskEventStyle;
};

export type RunPreparedEvent = PreparedEvent & {
  taskSlug?: string;
};

export type SpanDetail = {
  // ============================================================================
  // Core Identity & Structure
  // ============================================================================
  spanId: string; // Tree structure, span identification
  parentId: string | null; // Tree hierarchy
  message: string; // Displayed as span title

  // ============================================================================
  // Status & State
  // ============================================================================
  isError: boolean; // Error status display, filtering, status icons
  isPartial: boolean; // In-progress status display, timeline calculations
  isCancelled: boolean; // Cancelled status display, status determination
  level: TaskEventLevel; // Text styling, timeline rendering decisions

  // ============================================================================
  // Timing
  // ============================================================================
  startTime: Date; // Timeline calculations, display
  duration: number; // Timeline width, duration display, calculations

  // ============================================================================
  // Content & Display
  // ============================================================================
  events: SpanEvents; // Timeline events, SpanEvents component
  style: TaskEventStyle; // Icons, variants, accessories (RunIcon, SpanTitle)
  properties: Record<string, unknown> | string | number | boolean | null | undefined; // Displayed as JSON in span properties (CodeBlock)
  resourceProperties?: Record<string, unknown> | string | number | boolean | null | undefined; // Displayed as JSON in span resource properties (CodeBlock)

  // ============================================================================
  // Entity & Relationships
  // ============================================================================
  entity: {
    // Used for entity type switching in SpanEntity
    type: string | undefined;
    id: string | undefined;
    metadata: string | undefined;
  };

  metadata: any; // Used by SpanPresenter for entity processing
};

// ============================================================================
// Span and Link Types
// ============================================================================

export type SpanSummaryCommon = {
  id: string;
  parentId: string | undefined;
  runId: string;
  data: {
    message: string;
    events: SpanEvents;
    startTime: Date;
    duration: number;
    isError: boolean;
    isPartial: boolean;
    isCancelled: boolean;
    level: NonNullable<CreateEventInput["level"]>;
    attemptNumber?: number;
  };
};

export type SpanSummary = {
  id: string;
  parentId: string | undefined;
  runId: string;
  data: {
    message: string;
    style: TaskEventStyle;
    events: SpanEvents;
    startTime: Date;
    duration: number;
    isError: boolean;
    isPartial: boolean;
    isCancelled: boolean;
    isDebug: boolean;
    level: NonNullable<CreateEventInput["level"]>;
    attemptNumber?: number;
  };
};

export type SpanOverride = {
  isCancelled?: boolean;
  isError?: boolean;
  duration?: number;
  events?: SpanEvents;
};

export type TraceSummary = {
  rootSpan: SpanSummary;
  spans: Array<SpanSummary>;
  overridesBySpanId?: Record<string, SpanOverride>;
  /** Set when a subtree fetch hit the row cap before collecting all descendants. */
  isTruncated?: boolean;
};

export type SpanDetailedSummary = {
  id: string;
  parentId: string | undefined;
  runId: string;
  data: {
    message: string;
    taskSlug?: string;
    events: SpanEvents;
    startTime: Date;
    duration: number;
    isError: boolean;
    isPartial: boolean;
    isCancelled: boolean;
    level: NonNullable<CreateEventInput["level"]>;
    attemptNumber?: number;
    properties?: Attributes;
  };
  children: Array<SpanDetailedSummary>;
};

// A single trace event for the streaming export path (the "Download trace"
// feature). Deliberately flat and self-contained: it carries its own parent ref
// so hierarchy is reconstructable downstream without ever building a tree. Used
// by `streamTraceEvents`, which yields these one at a time so an arbitrarily
// large trace is never fully resident in memory.
export type StreamedTraceEvent = {
  spanId: string;
  parentSpanId: string;
  startTime: Date;
  durationNs: number;
  level: string;
  message: string;
  isError: boolean;
  // Span attributes/properties as a raw JSON string, emitted verbatim (the
  // ClickHouse store already materialises it as text — no per-row parse).
  propertiesText: string;
};

export type TraceDetailedSummary = {
  traceId: string;
  rootSpan: SpanDetailedSummary;
  /** Set when a fetch hit the row cap before collecting all spans. */
  isTruncated?: boolean;
};

// ============================================================================
// Event Repository Interface
// ============================================================================

/**
 * Interface for the EventRepository class.
 * Defines the public API for managing task events, traces, and spans.
 */
export interface IEventRepository {
  maximumLiveReloadingSetting: number;
  // Event insertion methods
  insertMany(events: CreateEventInput[]): void;
  insertManyImmediate(events: CreateEventInput[]): Promise<void>;
  insertManyMetrics(rows: MetricsV1Input[]): void;

  // Run event completion methods
  completeSuccessfulRunEvent(params: { run: CompleteableTaskRun; endTime?: Date }): Promise<void>;

  completeCachedRunEvent(params: {
    run: CompleteableTaskRun;
    blockedRun: CompleteableTaskRun;
    spanId: string;
    parentSpanId: string;
    spanCreatedAt: Date;
    isError: boolean;
    endTime?: Date;
  }): Promise<void>;

  completeFailedRunEvent(params: {
    run: CompleteableTaskRun;
    endTime?: Date;
    exception: { message?: string; type?: string; stacktrace?: string };
  }): Promise<void>;

  completeExpiredRunEvent(params: {
    run: CompleteableTaskRun;
    endTime?: Date;
    ttl: string;
  }): Promise<void>;

  createAttemptFailedRunEvent(params: {
    run: CompleteableTaskRun;
    endTime?: Date;
    attemptNumber: number;
    exception: { message?: string; type?: string; stacktrace?: string };
  }): Promise<void>;

  cancelRunEvent(params: {
    reason: string;
    run: CompleteableTaskRun;
    cancelledAt: Date;
  }): Promise<void>;

  // Query methods
  getTraceSummary(
    storeTable: TaskEventStoreTable,
    environmentId: string,
    traceId: string,
    startCreatedAt: Date,
    endCreatedAt?: Date,
    options?: { includeDebugLogs?: boolean }
  ): Promise<TraceSummary | undefined>;

  /** Fetch the anchor span, its ancestors (for override propagation), and all descendants. */
  getTraceSubtreeSummary(
    storeTable: TaskEventStoreTable,
    environmentId: string,
    traceId: string,
    anchorSpanId: string,
    startCreatedAt: Date,
    endCreatedAt?: Date,
    options?: { includeDebugLogs?: boolean }
  ): Promise<TraceSummary | undefined>;

  getTraceDetailedSummary(
    storeTable: TaskEventStoreTable,
    environmentId: string,
    traceId: string,
    startCreatedAt: Date,
    endCreatedAt?: Date,
    options?: { includeDebugLogs?: boolean }
  ): Promise<TraceDetailedSummary | undefined>;

  /** Fetch the anchor span subtree as a detailed hierarchical trace rooted at anchorSpanId. */
  getTraceDetailedSubtreeSummary(
    storeTable: TaskEventStoreTable,
    environmentId: string,
    traceId: string,
    anchorSpanId: string,
    startCreatedAt: Date,
    endCreatedAt?: Date,
    options?: { includeDebugLogs?: boolean }
  ): Promise<TraceDetailedSummary | undefined>;

  // Streams a trace's events in start_time order, one at a time, without ever
  // materialising the full result set or a tree. Powers the streaming trace
  // export so arbitrarily large traces download with bounded memory.
  streamTraceEvents(
    storeTable: TaskEventStoreTable,
    environmentId: string,
    traceId: string,
    startCreatedAt: Date,
    endCreatedAt?: Date,
    options?: { includeDebugLogs?: boolean }
  ): AsyncIterable<StreamedTraceEvent>;

  getRunEvents(
    storeTable: TaskEventStoreTable,
    environmentId: string,
    traceId: string,
    runId: string,
    startCreatedAt: Date,
    endCreatedAt?: Date
  ): Promise<RunPreparedEvent[]>;

  getSpan(
    storeTable: TaskEventStoreTable,
    environmentId: string,
    spanId: string,
    traceId: string,
    startCreatedAt: Date,
    endCreatedAt?: Date,
    options?: { includeDebugLogs?: boolean }
  ): Promise<SpanDetail | undefined>;

  // Event recording methods
  recordEvent(
    message: string,
    options: TraceEventOptions & { duration?: number; parentId?: string }
  ): Promise<void>;

  traceEvent<TResult>(
    message: string,
    options: TraceEventOptions & { incomplete?: boolean; isError?: boolean },
    callback: (
      e: EventBuilder,
      traceContext: Record<string, string | undefined>,
      traceparent?: { traceId: string; spanId: string }
    ) => Promise<TResult>
  ): Promise<TResult>;
}
