import { PdfPageSelection } from "../../../public/types";
import { IronPdfServiceClient } from "../../generated_proto/ironpdfengineproto/IronPdfService";
import { Access } from "../../access";
import { ImagesResultStreamP__Output } from "../../generated_proto/ironpdfengineproto/ImagesResultStreamP";
import { ClientReadableStream } from "@grpc/grpc-js";
import {
	AsyncPdfPageSelectionToIndexes, isNullOrUndefined
} from "../util";

export async function rasterizeToImageBuffers(
	id: string,
	pdfPageSelection: PdfPageSelection
): Promise<Buffer[]> {
	const client: IronPdfServiceClient = await Access.ensureConnection();

	const pi = await AsyncPdfPageSelectionToIndexes(id, pdfPageSelection);

	return new Promise(
		(
			resolve: (_: Buffer[]) => void,
			reject: (errorMsg: string) => void
		) => {
			const stream: ClientReadableStream<ImagesResultStreamP__Output> =
				client.pdfiumImagePdfToImages({
					document: { documentId: id },
					dpi: 96,
					pageIndexes: pi,
				});
			handleImageResultStream(stream, resolve, reject);
		}
	);
}

export async function extractRawImages(
	id: string,
	pdfPageSelection: PdfPageSelection
): Promise<Buffer[]> {
	const client: IronPdfServiceClient = await Access.ensureConnection();
	const pi = await AsyncPdfPageSelectionToIndexes(id, pdfPageSelection);
	return new Promise(
		(
			resolve: (_: Buffer[]) => void,
			reject: (errorMsg: string) => void
		) => {
			const stream: ClientReadableStream<ImagesResultStreamP__Output> =
				client.pdfiumImageExtractAllRawImages({
					document: { documentId: id },
					pageIndexes: pi,
				});
			handleImageResultStream(stream, resolve, reject);
		}
	);
}

function handleImageResultStream(
	stream: ClientReadableStream<ImagesResultStreamP__Output>,
	resolve: (_: Buffer[]) => void,
	reject: (errorMsg: string) => void
) {
	const temp: { imageIndex: number; buffer: Buffer }[] = [];

	stream.on("data", (data: ImagesResultStreamP__Output) => {
		if (data.exception) {
			reject(
				`${data.exception.message}/n${data.exception.remoteStackTrace}`
			);
		} else if (data.rawImagesChunk) {
			if (data.rawImagesChunk) {
				temp.push({
					imageIndex: data.rawImagesChunk.imageIndex!,
					buffer: data.rawImagesChunk.rawImageChunk!,
				});
			}
		}
	});

	stream.on("end", () => {
		const groupedData = temp.reduce(
			(
				result: { [imageIndex: number]: Buffer },
				item: { imageIndex: number; buffer: Buffer }
			) => {
				if (isNullOrUndefined(result[item.imageIndex])) {
					result[item.imageIndex] = Buffer.from([]);
				}
				const oldItem = result[item.imageIndex];
				if(!isNullOrUndefined(oldItem)){
					result[item.imageIndex] = Buffer.concat([
						oldItem,
						item.buffer,
					]);
				}
				return result;
			},
			{}
		);

		return resolve(Object.values(groupedData));
	});
}