import { readFileSync, writeFileSync } from "node:fs";
import path from "node:path";
import { EMBEDDING_DIM, embedTexts, MODEL_ID } from "../lib/embedder";

type CatalogImage = {
  id: string;
  image: string;
  captions: string[];
};

type IndexEntry = {
  id: string;
  image: string;
  caption: string;
};

const BATCH_SIZE = 8;

async function main() {
  const root = process.cwd();
  const catalog = JSON.parse(
    readFileSync(path.join(root, "data/catalog.json"), "utf8"),
  ) as CatalogImage[];

  const entries: IndexEntry[] = catalog.flatMap((item) =>
    item.captions.map((caption) => ({
      id: item.id,
      image: item.image,
      caption,
    })),
  );

  const embeddings = new Float32Array(entries.length * EMBEDDING_DIM);
  console.log(
    `Embedding ${entries.length} captions from ${catalog.length} images with ${MODEL_ID}`,
  );

  for (let start = 0; start < entries.length; start += BATCH_SIZE) {
    const batch = entries
      .slice(start, start + BATCH_SIZE)
      .map((entry) => entry.caption);
    const vectors = await embedTexts(batch);
    if (vectors.some((vector) => vector.length !== EMBEDDING_DIM)) {
      throw new Error(
        `Unexpected embedding length ${vectors[0]?.length ?? 0}; expected ${EMBEDDING_DIM}`,
      );
    }
    vectors.forEach((vector, index) => {
      embeddings.set(vector, (start + index) * EMBEDDING_DIM);
    });
    console.log(
      `embedded ${Math.min(start + BATCH_SIZE, entries.length)}/${entries.length}`,
    );
  }

  const dataDir = path.join(root, "data");
  writeFileSync(
    path.join(dataDir, "index.json"),
    `${JSON.stringify(entries)}\n`,
  );
  writeFileSync(
    path.join(dataDir, "embeddings.f32"),
    Buffer.from(
      embeddings.buffer,
      embeddings.byteOffset,
      embeddings.byteLength,
    ),
  );
  console.log(`wrote ${entries.length} embeddings`);
}

main().catch((error) => {
  console.error(error);
  process.exit(1);
});