import { VoyageAIClient } from "voyageai";
import { pipeline } from "node:stream/promises";
import { Readable } from "node:stream";
import { batch } from "../lib/streams.js";
import type { EmbeddingMap } from "../types/domain.js";

const client = new VoyageAIClient({
    apiKey: process.env.VOYAGEAI_API_KEY,
});

const EMBED_BATCH_SIZE = 128;

type EmbeddableNote = { id: string; text: string };

/**
 * @returns noteId to embedding vector
 */
export const embedNotes = async (notes: EmbeddableNote[]): Promise<EmbeddingMap> => {
    const embeddingMap: EmbeddingMap = new Map();

    if (!notes || notes.length === 0) {
        return embeddingMap;
    }

    await pipeline(
        Readable.from(notes, { objectMode: true }),
        batch<EmbeddableNote>(EMBED_BATCH_SIZE),
        async (batches: AsyncIterable<EmbeddableNote[]>) => {
            for await (const chunk of batches) {
                const response = await client.embed({
                    input: chunk.map((n) => n.text),
                    model: "voyage-3.5",
                });

                response.data?.forEach((item, i) => {
                    if (item.embedding) {
                        embeddingMap.set(chunk[i].id, item.embedding);
                    }
                });
            }
        }
    );

    return embeddingMap;
};