import {
  defineBatchStrategyMap,
  type AnyBatchOperationStrategy,
  type BatchOperationStrategy,
} from './batching-types';

type TestRateLimitPayload = {
  key: string;
  stage?: string;
  lead_id?: string;
  row_number?: number;
  simulated_delay_ms?: number;
};

type TestRateLimitBatchPayload = {
  key: string;
  stage?: string;
  simulated_delay_ms?: number;
  items: Array<{
    itemKey: string;
    payload: TestRateLimitPayload;
  }>;
};

type TestRateLimitBatchResult = {
  status: 'completed';
  key: string;
  stage?: string;
  batch: true;
  batch_size: number;
  items: Array<{
    itemKey: string;
    result: {
      status: 'completed';
      key: string;
      batch: false;
      lead_id: string | null;
      row_number: number | null;
    };
  }>;
};

function testRateLimitBatchItems(
  value: unknown,
): Array<{ itemKey?: string; result?: unknown }> {
  if (!value || typeof value !== 'object' || Array.isArray(value)) {
    return [];
  }
  const record = value as Record<string, unknown>;
  if (Array.isArray(record.items)) {
    return record.items as Array<{ itemKey?: string; result?: unknown }>;
  }
  const candidates = [
    record.data,
    record.result,
    record.output,
    record.toolResponse &&
    typeof record.toolResponse === 'object' &&
    !Array.isArray(record.toolResponse)
      ? (record.toolResponse as Record<string, unknown>).raw
      : undefined,
    record.toolOutput &&
    typeof record.toolOutput === 'object' &&
    !Array.isArray(record.toolOutput)
      ? (record.toolOutput as Record<string, unknown>).raw
      : undefined,
  ];
  for (const candidate of candidates) {
    const items = testRateLimitBatchItems(candidate);
    if (items.length > 0) return items;
  }
  return [];
}

const testRateLimitBatchStrategy: BatchOperationStrategy<
  TestRateLimitPayload,
  TestRateLimitBatchPayload,
  TestRateLimitBatchResult
> = {
  sourceOperation: 'test_rate_limit',
  batchOperation: 'test_batch_rate_limit',
  kind: 'identifier_batch',
  maxBatchSize: 200,
  bucketKeyPayloadFields: ['stage'],
  canBatchWith(left, right) {
    return String(left.stage || '') === String(right.stage || '');
  },
  toBucketKey(payload) {
    return `test_batch_rate_limit:${String(payload.stage || '')}`;
  },
  toItemKey(payload) {
    return String(payload.lead_id || payload.row_number || payload.key || '');
  },
  compile(payloads) {
    const stage = String(payloads[0]?.stage || '');
    const simulatedDelayMs =
      typeof payloads[0]?.simulated_delay_ms === 'number'
        ? payloads[0].simulated_delay_ms
        : undefined;
    const items = payloads.map((payload, index) => ({
      itemKey: String(
        payload.lead_id || payload.row_number || payload.key || `row_${index}`,
      ),
      payload,
    }));
    return {
      batchOperation: 'test_batch_rate_limit',
      batchPayload: {
        key: 'batch',
        ...(stage ? { stage } : {}),
        ...(simulatedDelayMs !== undefined
          ? { simulated_delay_ms: simulatedDelayMs }
          : {}),
        items,
      },
      items,
    };
  },
  splitResult(fullResult, compiled) {
    const resultItems = testRateLimitBatchItems(fullResult);
    return compiled.items.map((item, index) => ({
      itemKey: item.itemKey,
      result:
        index < resultItems.length
          ? (resultItems[index]?.result ?? null)
          : null,
      rawResult:
        index < resultItems.length
          ? (resultItems[index]?.result ?? null)
          : null,
    }));
  },
};

export const DEFAULT_PLAY_RUNTIME_BATCH_STRATEGIES = defineBatchStrategyMap({
  test_rate_limit: testRateLimitBatchStrategy,
});

export function getDefaultPlayRuntimeBatchStrategy(
  operation: string | null | undefined,
): AnyBatchOperationStrategy | null {
  if (!operation) {
    return null;
  }
  return DEFAULT_PLAY_RUNTIME_BATCH_STRATEGIES[operation] ?? null;
}