export type BatchStrategyKind =
  | 'identifier_batch'
  | 'query_share'
  | 'async_dataset_job';

export type BatchCompileItem<TSingle extends object> = {
  // Preserved for diagnostics and job bookkeeping only.
  // Batch fanout is positional: compiled.items[index] maps to splitResult[index].
  itemKey: string;
  payload: TSingle;
};

export type BatchCompileResult<
  TBatch extends object,
  TSingle extends object,
> = {
  batchOperation: string;
  batchPayload: TBatch;
  items: Array<BatchCompileItem<TSingle>>;
};

export type BatchExecutionDataEnvelope<TBatchResult> = {
  data: TBatchResult;
};

export type BatchExecutionResultPayload<TBatchResult> =
  | TBatchResult
  | BatchExecutionDataEnvelope<TBatchResult>;

export type BatchSplitItemResult<TResult = unknown, TRawResult = unknown> = {
  // Preserved as metadata only. The runtime does not use itemKey to match rows.
  itemKey: string;
  result: TResult;
  rawResult?: TRawResult;
};

export type BatchOperationStrategy<
  TSingle extends Record<string, unknown> = Record<string, unknown>,
  TBatch extends Record<string, unknown> = Record<string, unknown>,
  TBatchResult extends object = Record<string, unknown>,
  TResult = unknown,
  TRawResult = unknown,
> = {
  sourceOperation: string;
  batchOperation: string;
  kind: BatchStrategyKind;
  maxBatchSize: number;
  /**
   * Payload field names that can change the batch bucket. Static map planning
   * may treat a tool call as one bucket only when these fields are absent from
   * the static params/source text; otherwise it must assume worst-case bucket
   * fanout because the field may vary per row.
   */
  bucketKeyPayloadFields?: string[];
  canBatchWith(left: TSingle, right: TSingle): boolean;
  toBucketKey(payload: TSingle): string;
  toItemKey(payload: TSingle): string;
  compile(payloads: TSingle[]): BatchCompileResult<TBatch, TSingle>;
  splitResult(
    fullResult: TBatchResult,
    compiled: BatchCompileResult<TBatch, TSingle>,
    // Implementations must return results in the same order as compiled.items.
  ): BatchSplitItemResult<TResult, TRawResult>[];
};

export type AnyBatchOperationStrategy = {
  sourceOperation: string;
  batchOperation: string;
  kind: BatchStrategyKind;
  maxBatchSize: number;
  bucketKeyPayloadFields?: string[];
  canBatchWith(
    left: Record<string, unknown>,
    right: Record<string, unknown>,
  ): boolean;
  toBucketKey(payload: Record<string, unknown>): string;
  toItemKey(payload: Record<string, unknown>): string;
  compile(payloads: Record<string, unknown>[]): {
    batchOperation: string;
    batchPayload: Record<string, unknown>;
    items: Array<{
      itemKey: string;
      payload: Record<string, unknown>;
    }>;
  };
  splitResult(
    fullResult: unknown,
    compiled: {
      batchOperation: string;
      batchPayload: Record<string, unknown>;
      items: Array<{ itemKey: string; payload: Record<string, unknown> }>;
    },
  ): BatchSplitItemResult<unknown, unknown>[];
};

export type BatchStrategyMap = Record<string, AnyBatchOperationStrategy>;

type StrictBatchOperationStrategy = BatchOperationStrategy<
  Record<string, unknown>,
  Record<string, unknown>,
  object,
  unknown,
  unknown
>;

export function defineBatchStrategyMap<
  TStrategies extends Record<string, StrictBatchOperationStrategy>,
>(strategies: TStrategies): BatchStrategyMap {
  return strategies as BatchStrategyMap;
}