/**
 * Gather (Embedding Lookup) Kernels
 *
 * Provides token embedding lookups from embedding tables.
 */

import type { CommandRecorder } from '../command-recorder.js';
import type { Tensor, TensorDtype } from '../tensor.js';
import type { OutputBufferOptions } from './types.js';
import type { SplitWeightBuffer, WeightDtype, WeightStorageEncoding } from '../weight-buffer.js';

/** Gather kernel options */
export interface GatherOptions extends OutputBufferOptions {
  useVec4?: boolean;
  embeddingDtype?: Extract<WeightDtype, 'f16' | 'f32' | 'litert_int4'>;
  /**
   * Output dtype. When 'f16', converts F32 embeddings to F16 output.
   * Default: 'f32'
   */
  outputDtype?: 'f16' | 'f32';
  /**
   * True if embeddings are stored as [hidden_size, vocab_size] (GGUF layout).
   * False if embeddings are stored as [vocab_size, hidden_size] (PyTorch layout).
   * Default: false (RDRR format uses PyTorch layout from SafeTensors).
   */
  transpose?: boolean;
  /** Optional index offset into the token indices buffer. */
  indexOffset?: number;
  /** Total hidden size stored in the source embedding table (defaults to hiddenSize). */
  inputHiddenSize?: number;
  /** Hidden-dimension offset inside the source embedding row (defaults to 0). */
  hiddenOffset?: number;
  /** Optional indirect dispatch buffer for GPU-driven workgroup counts. */
  indirectBuffer?: GPUBuffer | null;
  /** Byte offset into indirect dispatch buffer (default: 0). */
  indirectOffset?: number;
  /** Required when embeddingDtype is 'litert_int4'. */
  storageEncoding?: WeightStorageEncoding | null;
  /** Optional authored split gather section count, usually stamped from the manifest-selected embedding kernel. */
  splitGatherSectionCount?: number | null;
}

/**
 * Run gather/embedding lookup
 */
export declare function runGather(
  indices: GPUBuffer,
  embeddings: GPUBuffer,
  numTokens: number,
  hiddenSize: number,
  vocabSize: number,
  options?: GatherOptions
): Promise<Tensor>;

/**
 * Record gather (batched, no submit)
 */
export declare function recordGather(
  recorder: CommandRecorder,
  indices: GPUBuffer,
  embeddings: GPUBuffer,
  numTokens: number,
  hiddenSize: number,
  vocabSize: number,
  options?: GatherOptions
): Promise<Tensor>;

/**
 * Run gather/embedding lookup against a row-split F16 embedding table.
 */
export declare function runGatherSplit4(
  indices: GPUBuffer,
  splitEmbedding: SplitWeightBuffer,
  numTokens: number,
  hiddenSize: number,
  vocabSize: number,
  options?: GatherOptions
): Promise<Tensor>;

/**
 * Record gather against a row-split F16 embedding table.
 */
export declare function recordGatherSplit4(
  recorder: CommandRecorder,
  indices: GPUBuffer,
  splitEmbedding: SplitWeightBuffer,
  numTokens: number,
  hiddenSize: number,
  vocabSize: number,
  options?: GatherOptions
): Promise<Tensor>;

/**
 * Run gather/embedding lookup against a row-split F16 embedding table with up
 * to eight sections.
 */
export declare function runGatherSplit8(
  indices: GPUBuffer,
  splitEmbedding: SplitWeightBuffer,
  numTokens: number,
  hiddenSize: number,
  vocabSize: number,
  options?: GatherOptions
): Promise<Tensor>;

/**
 * Record gather against a row-split F16 embedding table with up to eight sections.
 */
export declare function recordGatherSplit8(
  recorder: CommandRecorder,
  indices: GPUBuffer,
  splitEmbedding: SplitWeightBuffer,
  numTokens: number,
  hiddenSize: number,
  vocabSize: number,
  options?: GatherOptions
): Promise<Tensor>;

/**
 * Run gather/embedding lookup against a row-split F16 embedding table, honoring
 * the manifest-stamped split variant when the split buffer provides one.
 */
export declare function runGatherSplit(
  indices: GPUBuffer,
  splitEmbedding: SplitWeightBuffer,
  numTokens: number,
  hiddenSize: number,
  vocabSize: number,
  options?: GatherOptions
): Promise<Tensor>;

/**
 * Record gather against a row-split F16 embedding table, honoring the
 * manifest-stamped split variant when the split buffer provides one.
 */
export declare function recordGatherSplit(
  recorder: CommandRecorder,
  indices: GPUBuffer,
  splitEmbedding: SplitWeightBuffer,
  numTokens: number,
  hiddenSize: number,
  vocabSize: number,
  options?: GatherOptions
): Promise<Tensor>;