path: root/packages/benchmarks/src/seed.ts

import pLimit from "p-limit";

import type { ZBookmarkList } from "@karakeep/shared/types/lists";
import type { ZTagBasic } from "@karakeep/shared/types/tags";
import { BookmarkTypes } from "@karakeep/shared/types/bookmarks";

import { logInfo, logStep, logSuccess } from "./log";
import { getTrpcClient, TrpcClient } from "./trpc";
import { waitUntil } from "./utils";

export interface SeedConfig {
  bookmarkCount: number;
  tagCount: number;
  listCount: number;
  concurrency: number;
  userCount: number;
}

export interface SeededBookmark {
  id: string;
  tags: ZTagBasic[];
  listId?: string;
  title: string | null | undefined;
}

export interface UserSeedData {
  apiKey: string;
  trpc: TrpcClient;
  email: string;
  tags: ZTagBasic[];
  lists: ZBookmarkList[];
  bookmarks: SeededBookmark[];
}

export interface SeedResult {
  users: UserSeedData[];
  searchTerm: string;
  // For backwards compatibility, expose the first user's data
  apiKey: string;
  trpc: TrpcClient;
  tags: ZTagBasic[];
  lists: ZBookmarkList[];
  bookmarks: SeededBookmark[];
}

const TOPICS = [
  "performance",
  "search",
  "reading",
  "workflow",
  "api",
  "workers",
  "backend",
  "frontend",
  "productivity",
  "cli",
];

async function seedUserData(
  authlessClient: TrpcClient,
  userIndex: number,
  config: SeedConfig,
  timestamp: number,
): Promise<UserSeedData> {
  const email = `benchmarks+${timestamp}+user${userIndex}@example.com`;
  const password = "benchmarks1234";

  logStep(`Creating user ${userIndex + 1}/${config.userCount}`);
  await authlessClient.users.create.mutate({
    name: `Benchmark User ${userIndex + 1}`,
    email,
    password,
    confirmPassword: password,
  });
  const { key } = await authlessClient.apiKeys.exchange.mutate({
    email,
    password,
    keyName: `benchmark-key-${userIndex}`,
  });

  const trpc = getTrpcClient(key);
  logSuccess(`User ${userIndex + 1} ready`);

  logStep(`Creating ${config.tagCount} tags for user ${userIndex + 1}`);
  const tags: ZTagBasic[] = [];
  for (let i = 0; i < config.tagCount; i++) {
    const tag = await trpc.tags.create.mutate({
      name: `user${userIndex}-topic-${i + 1}`,
    });
    tags.push(tag);
  }
  logSuccess(`Tags created for user ${userIndex + 1}`);

  logStep(`Creating ${config.listCount} lists for user ${userIndex + 1}`);
  const lists: ZBookmarkList[] = [];
  for (let i = 0; i < config.listCount; i++) {
    const list = await trpc.lists.create.mutate({
      name: `User ${userIndex + 1} List ${i + 1}`,
      description: `Auto-generated benchmark list #${i + 1} for user ${userIndex + 1}`,
      icon: "bookmark",
    });
    lists.push(list);
  }
  logSuccess(`Lists created for user ${userIndex + 1}`);

  logStep(
    `Creating ${config.bookmarkCount} bookmarks for user ${userIndex + 1}`,
  );
  const limit = pLimit(config.concurrency);
  const bookmarks: SeededBookmark[] = [];

  await Promise.all(
    Array.from({ length: config.bookmarkCount }).map((_, index) =>
      limit(async () => {
        const topic = TOPICS[index % TOPICS.length];
        const createdAt = new Date(Date.now() - index * 3000);
        const bookmark = await trpc.bookmarks.createBookmark.mutate({
          type: BookmarkTypes.LINK,
          url: `https://example.com/user${userIndex}/${topic}/${index}`,
          title: `User ${userIndex + 1} ${topic} article ${index}`,
          source: "api",
          summary: `Benchmark dataset entry about ${topic} for user ${userIndex + 1}.`,
          favourited: index % 7 === 0,
          archived: index % 11 === 0,
          createdAt,
        });

        const primaryTag = tags[index % tags.length];
        const secondaryTag = tags[(index + 5) % tags.length];
        const attachedTags = [primaryTag, secondaryTag];
        await trpc.bookmarks.updateTags.mutate({
          bookmarkId: bookmark.id,
          attach: attachedTags.map((tag) => ({
            tagId: tag.id,
            tagName: tag.name,
          })),
          detach: [],
        });

        let listId: string | undefined;
        if (lists.length > 0) {
          const list = lists[index % lists.length];
          await trpc.lists.addToList.mutate({
            listId: list.id,
            bookmarkId: bookmark.id,
          });
          listId = list.id;
        }

        bookmarks.push({
          id: bookmark.id,
          tags: attachedTags,
          listId,
          title: bookmark.title,
        });
      }),
    ),
  );
  logSuccess(`Bookmarks created for user ${userIndex + 1}`);

  return {
    apiKey: key,
    trpc,
    email,
    tags,
    lists,
    bookmarks,
  };
}

export async function seedData(config: SeedConfig): Promise<SeedResult> {
  const authlessClient = getTrpcClient();
  const timestamp = Date.now();

  logInfo(`Seeding data for ${config.userCount} users`);
  const users: UserSeedData[] = [];

  // Create all users sequentially to avoid race conditions
  for (let i = 0; i < config.userCount; i++) {
    const userData = await seedUserData(authlessClient, i, config, timestamp);
    users.push(userData);
  }

  const searchTerm = "benchmark";
  logStep("Waiting for search index to be ready");
  // Use the first user's client to check search readiness
  await waitUntil(
    async () => {
      const results = await users[0].trpc.bookmarks.searchBookmarks.query({
        text: searchTerm,
        limit: 1,
      });
      return results.bookmarks.length > 0;
    },
    "search data to be indexed",
    120_000,
    2_000,
  );
  logSuccess("Search index warmed up");

  const totalBookmarks = users.reduce(
    (sum, user) => sum + user.bookmarks.length,
    0,
  );
  const totalTags = users.reduce((sum, user) => sum + user.tags.length, 0);
  const totalLists = users.reduce((sum, user) => sum + user.lists.length, 0);

  logInfo(
    `Seeded ${totalBookmarks} bookmarks across ${totalTags} tags and ${totalLists} lists for ${config.userCount} users`,
  );

  // Return first user's data for backwards compatibility
  const firstUser = users[0];
  return {
    users,
    searchTerm,
    apiKey: firstUser.apiKey,
    trpc: firstUser.trpc,
    tags: firstUser.tags,
    lists: firstUser.lists,
    bookmarks: firstUser.bookmarks,
  };
}