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

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

// Well under Voyage's per-request input and token ceilings.
const EMBED_BATCH_SIZE = 128;

/**
 * @param {Array<{id: string, text: string}>} notes
 * @returns {Promise<Map<string, number[]>>} noteId → embedding vector
 */
export const embedNotes = async (notes) => {
  const embeddingMap = new Map();

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

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

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

  return embeddingMap;
};