/* eslint-disable @typescript-eslint/no-explicit-any */
import type { TrackedEnvelope } from "@trpc/server";
import { isTrackedEnvelope, tracked } from "@trpc/server";
import { z } from "zod";
function isAsyncIterable<TValue, TReturn = unknown>(
  value: unknown,
): value is AsyncIterable<TValue, TReturn> {
  return !!value && typeof value === "object" && Symbol.asyncIterator in value;
}
const trackedEnvelopeSchema =
  z.custom<TrackedEnvelope<unknown>>(isTrackedEnvelope);
/**
 * A Zod schema helper designed specifically for validating async iterables. This schema ensures that:
 * 1. The value being validated is an async iterable.
 * 2. Each item yielded by the async iterable conforms to a specified type.
 * 3. The return value of the async iterable, if any, also conforms to a specified type.
 */
export function zAsyncIterable<
  TYield,
  TReturn = void,
  Tracked extends boolean = false,
>(opts: {
  yield: z.ZodType<TYield, any, any>;
  return?: z.ZodType<TReturn, any, any>;
  tracked?: Tracked;
}) {
  const baseSchema = z.custom<
    AsyncIterable<
      Tracked extends true
        ? TrackedEnvelope<z.input<typeof opts.yield>>
        : z.input<typeof opts.yield>,
      z.input<typeof opts.return>
    >
  >((val: unknown): val is AsyncIterable<any, any> => isAsyncIterable(val), {
    message: "Expected AsyncIterable",
  });

  return baseSchema.transform(async function* (iter) {
    const iterator = iter[Symbol.asyncIterator]();
    try {
      let next;
      while ((next = await iterator.next()) && !next.done) {
        if (opts.tracked) {
          const [id, data] = await trackedEnvelopeSchema.parseAsync(next.value);
          yield tracked(id, await opts.yield.parseAsync(data));
          continue;
        }
        yield await opts.yield.parseAsync(next.value);
      }
      if (opts.return) {
        return (await opts.return.parseAsync(next.value)) as TReturn;
      }
      return;
    } finally {
      await iterator.return?.();
    }
  });
}