import type { XRPCQueries } from "@atcute/lexicons/ambient";
import type {
	InferInput,
	InferOutput,
	ObjectSchema,
	XRPCLexBodyParam
} from "@atcute/lexicons/validations";
import { ok, type BobbinContext, type XrpcRequestInit } from "./client";

export type QueryOutput<TName extends keyof XRPCQueries> =
	XRPCQueries[TName]["output"] extends XRPCLexBodyParam
		? InferOutput<XRPCQueries[TName]["output"]["schema"]>
		: never;

export type QueryParams<TName extends keyof XRPCQueries> =
	XRPCQueries[TName]["params"] extends ObjectSchema
		? InferInput<XRPCQueries[TName]["params"]>
		: Record<string, never>;

type CursorPage = { items: readonly unknown[]; cursor?: string };

// registered queries whose output is an { items, cursor? } page.
export type PaginatedQuery = {
	[K in keyof XRPCQueries]: QueryOutput<K> extends CursorPage ? K : never;
}[keyof XRPCQueries];

export type PageItem<TName extends PaginatedQuery> =
	QueryOutput<TName> extends { items: readonly (infer I)[] } ? I : never;

export type PageParams<TName extends PaginatedQuery> = Omit<QueryParams<TName>, "cursor">;

export interface PaginateOptions extends XrpcRequestInit {
	maxPages?: number;
}

const requestPage = <TName extends PaginatedQuery>(
	ctx: BobbinContext,
	name: TName,
	params: PageParams<TName> & { cursor?: string },
	init?: XrpcRequestInit
): Promise<QueryOutput<TName>> =>
	ok(
		ctx.xrpc.get(name, {
			params,
			signal: init?.signal,
			headers: init?.headers
		} as never)
	) as Promise<QueryOutput<TName>>;
export const fetchPage = <TName extends PaginatedQuery>(
	ctx: BobbinContext,
	name: TName,
	params: PageParams<TName> & { cursor?: string },
	init?: XrpcRequestInit
): Promise<QueryOutput<TName>> => requestPage(ctx, name, params, init);

export async function* pages<TName extends PaginatedQuery>(
	ctx: BobbinContext,
	name: TName,
	params: PageParams<TName>,
	options: PaginateOptions = {}
): AsyncGenerator<QueryOutput<TName>, void, unknown> {
	let cursor: string | undefined;
	let seen = 0;
	do {
		const data = await requestPage(
			ctx,
			name,
			{ ...params, cursor } as PageParams<TName> & { cursor?: string },
			options
		);
		yield data;
		const next = (data as CursorPage).cursor;
		if (next !== undefined && next === cursor) break;
		cursor = next;
		if (options.maxPages && ++seen >= options.maxPages) break;
	} while (cursor);
}

export async function* items<TName extends PaginatedQuery>(
	ctx: BobbinContext,
	name: TName,
	params: PageParams<TName>,
	options: PaginateOptions = {}
): AsyncGenerator<PageItem<TName>, void, unknown> {
	for await (const data of pages(ctx, name, params, options)) {
		for (const item of (data as CursorPage).items) yield item as PageItem<TName>;
	}
}

export interface CollectOptions extends PaginateOptions {
	max?: number;
}

export const collect = async <TName extends PaginatedQuery>(
	ctx: BobbinContext,
	name: TName,
	params: PageParams<TName>,
	options: CollectOptions = {}
): Promise<PageItem<TName>[]> => {
	const out: PageItem<TName>[] = [];
	for await (const item of items(ctx, name, params, options)) {
		out.push(item);
		if (options.max && out.length >= options.max) break;
	}
	return out;
};

export async function* paginateBy<T>(
	load: (cursor: string | undefined) => Promise<{ items: readonly T[]; cursor?: string | null }>,
	options: { maxPages?: number } = {}
): AsyncGenerator<T, void, unknown> {
	let cursor: string | undefined;
	let seen = 0;
	do {
		const page = await load(cursor);
		for (const item of page.items) yield item;
		const next = page.cursor ?? undefined;
		if (next !== undefined && next === cursor) break;
		cursor = next;
		if (options.maxPages && ++seen >= options.maxPages) break;
	} while (cursor);
}