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

export interface ChunkExecutionResult<TRequest, TResult> {
  request: TRequest;
  result: TResult | null;
  /**
   * Present when this request's execution rejected. The request failed but
   * its siblings in the chunk kept their results — one provider hiccup must
   * stay a per-request failure, not a run-level abort that discards billed
   * work (rows already persisted by completed calls stay recoverable).
   */
  error?: unknown;
}

export function formatChunkExecutionError(error: unknown): string {
  if (error instanceof Error) {
    return error.message;
  }
  return String(error);
}

export interface CompiledRequestBatch<TRequest> {
  batchOperation: string;
  memberRequests: TRequest[];
  batchPayload: Record<string, unknown>;
  splitResults: (value: unknown) => Array<unknown | null>;
}

/**
 * Plan chunk boundaries. A chunk closes when it reaches `batchSize` units OR,
 * when a `maxChunkWeight` + `weightOf` exposure cap is supplied, when its
 * cumulative weight reaches that cap — whichever comes first. A single unit that
 * alone exceeds either limit still forms its own chunk (batches are atomic and
 * never split). The weight cap is the persistence-latch dispatch-wave bound:
 * without it a large `batchSize` (provider-pacing hint) would let a whole
 * chunk's provider calls bill before the latch can gate the rest.
 */
function planExecutionChunks<TRequest>(
  requests: TRequest[],
  batchSize: number,
  maxChunkWeight: number | undefined,
  weightOf: ((request: TRequest) => number) | undefined,
): TRequest[][] {
  const maxUnits = Math.max(1, Math.floor(batchSize));
  const weightCap =
    maxChunkWeight != null && weightOf != null
      ? Math.max(1, Math.floor(maxChunkWeight))
      : null;
  const chunks: TRequest[][] = [];
  let current: TRequest[] = [];
  let weight = 0;
  for (const request of requests) {
    if (
      current.length > 0 &&
      (current.length >= maxUnits || (weightCap != null && weight >= weightCap))
    ) {
      chunks.push(current);
      current = [];
      weight = 0;
    }
    current.push(request);
    if (weightCap != null && weightOf != null) {
      weight += Math.max(1, Math.floor(weightOf(request)));
    }
  }
  if (current.length > 0) chunks.push(current);
  return chunks;
}

export async function executeChunkedRequests<TRequest, TResult>(input: {
  requests: TRequest[];
  batchSize: number;
  /**
   * Optional exposure cap: the maximum cumulative `weightOf` a single chunk may
   * dispatch before the caller's per-unit latch check runs again. Set to
   * PROVIDER_DISPATCH_WAVE_SIZE by the batch dispatch paths so a persistence
   * failure prevents the rest of a large batch group instead of billing it all.
   */
  maxChunkWeight?: number;
  /** Per-request weight toward `maxChunkWeight` (e.g. batch member count). */
  weightOf?: (request: TRequest) => number;
  execute: (request: TRequest) => Promise<TResult>;
  /**
   * Loud per-request failure hook. A rejected request is recorded as a
   * `result: null` entry with the raw `error` set so typed callers can keep
   * domain errors intact; it must never abort the chunk, the sibling requests,
   * or the run. Callers use this to log and persist the failure.
   */
  onRequestError?: (request: TRequest, error: unknown) => void;
  onChunkComplete?: (
    results: Array<ChunkExecutionResult<TRequest, TResult>>,
  ) => void | Promise<void>;
  /** Keep settled entries for the returned aggregate. Defaults to true. */
  retainResults?: boolean;
}): Promise<Array<ChunkExecutionResult<TRequest, TResult>>> {
  const results: Array<ChunkExecutionResult<TRequest, TResult>> = [];

  const chunks = planExecutionChunks(
    input.requests,
    input.batchSize,
    input.maxChunkWeight,
    input.weightOf,
  );
  for (const chunk of chunks) {
    // notifyChain serializes onChunkComplete calls (the caller's callback can
    // touch shared, non-concurrency-safe state) without letting one entry's
    // callback failure skip a sibling's callback. A naive `.then()` chain
    // would propagate a rejection forward and silently drop every queued
    // notify() behind it -- for BetterContact/native-batch callers that
    // means a correlation-mismatch throw on one row would leak the tool
    // slots of every other row in the same chunk, since their
    // onChunkComplete (and its `finally { releaseToolSlot() }`) would never
    // run. Each callback failure is caught and re-thrown only after every
    // queued entry in the chunk has had its callback invoked.
    let notifyChain: Promise<void> = Promise.resolve();
    let firstCallbackFailure: unknown;
    const notify = async (
      entry: ChunkExecutionResult<TRequest, TResult>,
    ): Promise<void> => {
      if (input.retainResults !== false) {
        results.push(entry);
      }
      notifyChain = notifyChain.then(async () => {
        try {
          await input.onChunkComplete?.([entry]);
        } catch (error) {
          firstCallbackFailure ??= error;
        }
      });
      await notifyChain;
    };

    await Promise.all(
      chunk.map(async (request) => {
        let entry: ChunkExecutionResult<TRequest, TResult>;
        try {
          entry = {
            request,
            result: await input.execute(request),
          };
        } catch (error) {
          input.onRequestError?.(request, error);
          entry = {
            request,
            result: null,
            error,
          };
        }
        await notify(entry);
      }),
    );
    await notifyChain;
    if (firstCallbackFailure !== undefined) throw firstCallbackFailure;
  }

  return results;
}

export function compileRequestsWithStrategy<TRequest>(input: {
  requests: TRequest[];
  strategy: AnyBatchOperationStrategy;
  getPayload: (request: TRequest) => Record<string, unknown>;
}): Array<CompiledRequestBatch<TRequest>> {
  const compiledBatches: Array<CompiledRequestBatch<TRequest>> = [];
  const bucketedRequests = new Map<string, TRequest[]>();

  for (const request of input.requests) {
    const payload = input.getPayload(request);
    const bucketKey = String(input.strategy.toBucketKey(payload));
    if (!bucketedRequests.has(bucketKey)) {
      bucketedRequests.set(bucketKey, []);
    }
    bucketedRequests.get(bucketKey)!.push(request);
  }

  for (const bucketRequests of bucketedRequests.values()) {
    let currentBatch: TRequest[] = [];

    const flushBatch = () => {
      if (currentBatch.length === 0) {
        return;
      }

      const memberRequests = [...currentBatch];
      const compiled = input.strategy.compile(
        memberRequests.map((request) => input.getPayload(request)),
      );
      compiledBatches.push({
        batchOperation: compiled.batchOperation,
        memberRequests,
        batchPayload: compiled.batchPayload as Record<string, unknown>,
        splitResults: (value: unknown) => {
          const splitResults = input.strategy.splitResult(
            value as never,
            compiled as never,
          );
          return memberRequests.map(
            (_, index) => splitResults[index]?.result ?? null,
          );
        },
      });
      currentBatch = [];
    };

    for (const request of bucketRequests) {
      const payload = input.getPayload(request);
      const canAppend =
        currentBatch.length > 0 &&
        currentBatch.length < input.strategy.maxBatchSize &&
        currentBatch.every((existing) =>
          input.strategy.canBatchWith(input.getPayload(existing), payload),
        );

      if (!canAppend && currentBatch.length > 0) {
        flushBatch();
      }

      currentBatch.push(request);
    }

    flushBatch();
  }

  return compiledBatches;
}