evOo-local/packages/adapters/src/media-ingest-cluster.ts

import { readFile } from "node:fs/promises";
import type { AgentInputImage } from "@rakazo/adapter-kit";
import type { MediaClusterInput } from "@rakazo/db";
import { getLogger } from "@rakazo/logging";
import type { ClusterAsset, SessionOptions, VectorItem } from "@rakazo/media";
import {
  buildSessions,
  clusterFaces,
  clusterThemes,
  cosineSimilarity,
  decodeVectors,
  encodeVectors,
  meanVector,
  ruleLabel,
} from "@rakazo/media";
import { resolveInside } from "@rakazo/media/node";
import { z } from "zod";
import type { MediaServiceDeps } from "./media-context.js";
import { errorMessage, publishSafe } from "./media-ingest-files.js";

/*
 * Sessions, themes, and person suggestions for one space.
 *
 * Bounds: themes use average-linkage clustering (O(n²) time and memory), so
 * each run clusters at most THEME_BATCH_LIMIT of the most recently captured
 * analysed assets; older assets keep no theme until they fall into the window.
 * Face clustering is bounded the same way per run (FACE_BATCH_LIMIT unassigned
 * tracks); the rest are picked up by the next run.
 *
 * Theme titles come from the model only when the media model client returns
 * one (it checks AI consent for the space's library owner); it receives cover
 * thumbnails and nothing else. Face embeddings and person data never leave
 * this process.
 *
 * Collections the user merged or split keep their membership; re-clustering
 * never re-adds their assets elsewhere (see reconcileUserEdits).
 */

export const THEME_BATCH_LIMIT = 2000;
export const FACE_BATCH_LIMIT = 2000;
const THEME_MIN_SIZE = 3;
const THEME_SIMILARITY = 0.8;
const FACE_SIMILARITY = 0.6;
const COVER_IMAGES = 9;
const MODEL_LABELS_PER_RUN = 12;
const EMBEDDING_BATCH = 500;

const ThemeTitleSchema = z.object({ title: z.string().trim().min(1).max(80) });

const THEME_INSTRUCTIONS =
  'You name groups of photos and videos. Reply with one JSON object {"title": string}: ' +
  "a short, neutral title of at most five words for what the images have in common " +
  "(place type, activity, or subject). Never name or describe people's identities.";

interface UserCollection {
  clusterKey: string | null;
  assetIds: string[];
}

/**
 * Apply the user's edits to computed clusters of one kind:
 * - A user collection without a cluster key (split off) keeps its assets; they
 *   are withheld from every computed cluster.
 * - A user collection whose assets all still fall into one computed cluster
 *   (renamed only) follows that cluster: the cluster takes the collection's key
 *   so the collection keeps updating, also when an earlier asset changes the
 *   cluster's own key.
 * - Any other user collection (merged, or split and kept) is frozen: its key
 *   and assets are withheld, so re-clustering never re-adds them elsewhere.
 * Clusters that lose their first asset are re-keyed by the first remaining one
 * (the clustering's own rule); empty clusters and key collisions are dropped.
 */
export function reconcileUserEdits<T extends { key: string; assetIds: string[] }>(
  clusters: T[],
  userCollections: UserCollection[],
  prefix: "session" | "theme",
): T[] {
  const splitOff = new Set(
    userCollections.flatMap((collection) => (collection.clusterKey ? [] : collection.assetIds)),
  );
  const hostOf = new Map<string, T>();
  for (const cluster of clusters) {
    for (const id of cluster.assetIds) if (!splitOff.has(id)) hostOf.set(id, cluster);
  }
  const followedBy = new Map<string, string>();
  const frozenKeys = new Set<string>();
  const frozenAssets = new Set<string>();
  for (const collection of userCollections) {
    const [first] = collection.assetIds;
    if (!collection.clusterKey || first === undefined) continue;
    const host = hostOf.get(first);
    const contained = host && collection.assetIds.every((id) => hostOf.get(id) === host);
    if (host && contained && !followedBy.has(host.key)) {
      followedBy.set(host.key, collection.clusterKey);
    } else {
      frozenKeys.add(collection.clusterKey);
      for (const id of collection.assetIds) frozenAssets.add(id);
    }
  }
  const seen = new Set<string>(frozenKeys);
  const result: T[] = [];
  for (const cluster of clusters) {
    const assetIds = cluster.assetIds.filter((id) => !splitOff.has(id) && !frozenAssets.has(id));
    const first = assetIds[0];
    if (first === undefined) continue;
    const key =
      followedBy.get(cluster.key) ??
      (first === cluster.assetIds[0] ? cluster.key : `${prefix}:${first}`);
    if (seen.has(key)) continue;
    seen.add(key);
    result.push({ ...cluster, key, assetIds });
  }
  return result;
}

type ClusterRow = Awaited<ReturnType<MediaServiceDeps["repos"]["listAssetsForClustering"]>>[number];

/** Mean frame embedding per asset from the latest embeddings analysis, read in batches. */
async function meanEmbeddings(
  deps: MediaServiceDeps,
  spaceId: string,
  assetIds: string[],
): Promise<Map<string, Float32Array>> {
  const means = new Map<string, Float32Array>();
  for (let offset = 0; offset < assetIds.length; offset += EMBEDDING_BATCH) {
    const batch = assetIds.slice(offset, offset + EMBEDDING_BATCH);
    const rows = await deps.prisma.mediaAnalysis.findMany({
      where: { spaceId, kind: "embeddings", assetId: { in: batch }, vectors: { not: null } },
      select: { assetId: true, data: true, vectors: true },
      orderBy: [{ createdAt: "desc" }, { id: "desc" }],
    });
    for (const row of rows) {
      if (means.has(row.assetId) || !row.vectors) continue;
      const data = row.data as { dim?: unknown } | null;
      const dim = typeof data?.dim === "number" ? data.dim : 0;
      if (!Number.isInteger(dim) || dim <= 0 || row.vectors.byteLength % (dim * 4) !== 0) continue;
      const mean = meanVector(decodeVectors(row.vectors, dim));
      if (mean) means.set(row.assetId, mean);
    }
  }
  return means;
}

/** The vectors of the most common dimension (a model change mixes dimensions). */
function sameDimension<T extends { vector: Float32Array }>(items: T[]): T[] {
  const counts = new Map<number, number>();
  for (const item of items)
    counts.set(item.vector.length, (counts.get(item.vector.length) ?? 0) + 1);
  const dim = [...counts.entries()].sort((a, b) => b[1] - a[1] || b[0] - a[0])[0]?.[0];
  return items.filter((item) => item.vector.length === dim);
}

async function userCollections(
  deps: MediaServiceDeps,
  spaceId: string,
  kind: "session" | "theme",
): Promise<UserCollection[]> {
  const rows = await deps.prisma.mediaCollection.findMany({
    where: { spaceId, kind, labelSource: "user" },
    select: {
      clusterKey: true,
      items: { select: { assetId: true }, orderBy: { position: "asc" } },
    },
  });
  return rows.map((row) => ({
    clusterKey: row.clusterKey,
    assetIds: row.items.map((item) => item.assetId),
  }));
}

/** Whose AI consent governs model labels: the owner of the space's first media source. */
async function libraryOwner(deps: MediaServiceDeps, spaceId: string): Promise<string | null> {
  const source = await deps.prisma.mediaSource.findFirst({
    where: { spaceId },
    orderBy: [{ createdAt: "asc" }, { id: "asc" }],
    select: { ownerUserId: true },
  });
  if (source) return source.ownerUserId;
  const space = await deps.prisma.space.findUnique({
    where: { id: spaceId },
    select: { createdByUserId: true },
  });
  return space?.createdByUserId ?? null;
}

function validTimeZone(zone: string | null | undefined): string | null {
  if (!zone) return null;
  try {
    new Intl.DateTimeFormat("en", { timeZone: zone });
    return zone;
  } catch {
    return null;
  }
}

export interface ClusterOutcome {
  sessions: number;
  themes: number;
  modelLabels: number;
  assignedTracks: number;
  suggestedPeople: number;
}

export async function clusterSpace(
  deps: MediaServiceDeps,
  spaceId: string,
  options: { signal?: AbortSignal; sessionOptions?: Partial<SessionOptions> } = {},
): Promise<ClusterOutcome> {
  const userId = await libraryOwner(deps, spaceId);
  const profile = userId
    ? await deps.prisma.userProfile.findUnique({
        where: { userId },
        select: { timeZone: true, locale: true },
      })
    : null;
  const timeZone = validTimeZone(profile?.timeZone) ?? "UTC";
  const locale = profile?.locale ?? undefined;
  const label = (range: { startAt: Date | null; endAt: Date | null }, assetIds: string[]) =>
    ruleLabel({ ...range, albumNames: albums(assetIds) }, locale, timeZone);
  const rows = await deps.repos.listAssetsForClustering(spaceId);
  const extras = await deps.prisma.mediaAsset.findMany({
    where: { spaceId },
    select: { id: true, thumbPath: true, qualityScore: true },
  });
  const extra = new Map(extras.map((row) => [row.id, row]));
  const embeddings = await meanEmbeddings(
    deps,
    spaceId,
    rows.map((row) => row.id),
  );
  const byId = new Map(rows.map((row) => [row.id, row]));

  const cover = (assetIds: string[]): string | null => {
    const ranked = assetIds
      .map((id) => extra.get(id))
      .filter((row) => row?.thumbPath)
      .sort((a, b) => (b?.qualityScore ?? 0) - (a?.qualityScore ?? 0));
    return ranked[0]?.id ?? assetIds[0] ?? null;
  };
  const span = (assetIds: string[]) => {
    const dates = assetIds
      .flatMap((id) => {
        const at = byId.get(id)?.capturedAt;
        return at ? [at.getTime()] : [];
      })
      .sort((a, b) => a - b);
    return {
      startAt: dates.length > 0 ? new Date(dates[0] as number) : null,
      endAt: dates.length > 0 ? new Date(dates.at(-1) as number) : null,
    };
  };
  const albums = (assetIds: string[]) => assetIds.flatMap((id) => byId.get(id)?.albumNames ?? []);

  // Sessions -----------------------------------------------------------------
  const clusterAssets: ClusterAsset[] = rows.map((row: ClusterRow) => ({
    id: row.id,
    kind: row.kind as ClusterAsset["kind"],
    capturedAt: row.capturedAt,
    gpsLat: row.gpsLat,
    gpsLon: row.gpsLon,
    deviceModel: row.deviceModel,
    albumNames: row.albumNames,
    embedding: embeddings.get(row.id) ?? null,
  }));
  const computedSessions = buildSessions(clusterAssets, { timeZone, ...options.sessionOptions });
  const sessions: MediaClusterInput[] = reconcileUserEdits(
    computedSessions,
    await userCollections(deps, spaceId, "session"),
    "session",
  ).map((session) => {
    const range = span(session.assetIds);
    return {
      key: session.key,
      title: label(range, session.assetIds),
      labelSource: "rule",
      ...range,
      centroid: session.centroid,
      coverAssetId: cover(session.assetIds),
      assetIds: session.assetIds,
    };
  });
  options.signal?.throwIfAborted();
  await deps.repos.replaceClusters(spaceId, "session", sessions);
  await publishSafe(deps, spaceId, { type: "media.collections", kind: "session" });

  // Themes -------------------------------------------------------------------
  const recent = rows
    .flatMap((row) => {
      const vector = embeddings.get(row.id);
      return vector ? [{ id: row.id, vector, at: row.capturedAt?.getTime() ?? 0 }] : [];
    })
    .sort((a, b) => b.at - a.at || a.id.localeCompare(b.id))
    .slice(0, THEME_BATCH_LIMIT);
  const items: VectorItem[] = sameDimension(recent).map(({ id, vector }) => ({ id, vector }));
  const computedThemes = reconcileUserEdits(
    clusterThemes(items, { similarity: THEME_SIMILARITY, minSize: THEME_MIN_SIZE }).map(
      (theme) => ({ key: theme.key, assetIds: theme.ids }),
    ),
    await userCollections(deps, spaceId, "theme"),
    "theme",
  );
  const existing = await deps.prisma.mediaCollection.findMany({
    where: { spaceId, kind: "theme", clusterKey: { in: computedThemes.map((theme) => theme.key) } },
    select: { clusterKey: true, title: true, labelSource: true },
  });
  const modelTitles = new Map(
    existing.flatMap((row) =>
      row.labelSource === "model" && row.clusterKey ? [[row.clusterKey, row.title] as const] : [],
    ),
  );
  let modelCalls = 0;
  let modelLabels = 0;
  const themes: MediaClusterInput[] = [];
  for (const theme of computedThemes) {
    options.signal?.throwIfAborted();
    const range = span(theme.assetIds);
    const coverAssetId = cover(theme.assetIds);
    let title = modelTitles.get(theme.key) ?? null;
    if (!title && userId && modelCalls < MODEL_LABELS_PER_RUN) {
      modelCalls += 1;
      title = await themeTitle(deps, spaceId, userId, theme.assetIds, extra, options.signal);
    }
    if (title) modelLabels += 1;
    themes.push({
      key: theme.key,
      title: title ?? label(range, theme.assetIds),
      labelSource: title ? "model" : "rule",
      ...range,
      centroid: null,
      coverAssetId,
      assetIds: theme.assetIds,
    });
  }
  await deps.repos.replaceClusters(spaceId, "theme", themes);
  await publishSafe(deps, spaceId, { type: "media.collections", kind: "theme" });

  // People -------------------------------------------------------------------
  const people = await suggestPeople(deps, spaceId);
  if (people.assignedTracks + people.suggestedPeople > 0) {
    await publishSafe(deps, spaceId, { type: "media.people" });
  }
  return { sessions: sessions.length, themes: themes.length, modelLabels, ...people };
}

/** Ask the model for a theme title from cover thumbnails only; null falls back to rules. */
async function themeTitle(
  deps: MediaServiceDeps,
  spaceId: string,
  userId: string,
  assetIds: string[],
  extra: Map<string, { id: string; thumbPath: string | null; qualityScore: number | null }>,
  signal?: AbortSignal,
): Promise<string | null> {
  const covers = assetIds
    .flatMap((id) => {
      const row = extra.get(id);
      return row?.thumbPath ? [row] : [];
    })
    .sort((a, b) => (b.qualityScore ?? 0) - (a.qualityScore ?? 0) || a.id.localeCompare(b.id))
    .slice(0, COVER_IMAGES);
  const images: AgentInputImage[] = [];
  for (const [index, row] of covers.entries()) {
    try {
      const file = await resolveInside(deps.mediaRoot, row.thumbPath as string);
      images.push({
        name: `cover-${index + 1}.jpg`,
        mimeType: "image/jpeg",
        data: await readFile(file),
      });
    } catch {
      // A missing thumbnail only shrinks the set.
    }
  }
  if (images.length === 0) return null;
  try {
    const reply = await deps.model.generate({
      spaceId,
      userId,
      instructions: THEME_INSTRUCTIONS,
      prompt: `Name what these ${images.length} images have in common.`,
      images,
      schema: ThemeTitleSchema,
      tier: "fast",
      ...(signal ? { signal } : {}),
    });
    return reply?.title ?? null;
  } catch (error) {
    getLogger().warn("media theme label failed", { reason: errorMessage(error) });
    return null;
  }
}

/**
 * Unassigned face tracks join an unconfirmed suggestion with the same identity;
 * the rest are clustered into new suggestions. Confirmed people only ever gain
 * tracks through the user.
 */
async function suggestPeople(
  deps: MediaServiceDeps,
  spaceId: string,
): Promise<{ assignedTracks: number; suggestedPeople: number }> {
  const unassigned = (await deps.repos.faceTrackEmbeddings(spaceId, { unassignedOnly: true }))
    .slice(0, FACE_BATCH_LIMIT)
    .flatMap((track) => {
      const bytes = track.embedding;
      if (bytes.byteLength === 0 || bytes.byteLength % 4 !== 0) return [];
      const [vector] = decodeVectors(bytes, bytes.byteLength / 4);
      return vector ? [{ id: track.trackId, vector }] : [];
    });
  if (unassigned.length === 0) return { assignedTracks: 0, suggestedPeople: 0 };

  const confirmed = new Set(
    (await deps.repos.listPersonRecords(spaceId))
      .filter((person) => person.confirmed)
      .map((person) => person.id),
  );
  const suggestions = (await deps.repos.personCentroids(spaceId)).flatMap((person) => {
    if (confirmed.has(person.personId)) return [];
    const bytes = person.centroid;
    if (bytes.byteLength === 0 || bytes.byteLength % 4 !== 0) return [];
    const [vector] = decodeVectors(bytes, bytes.byteLength / 4);
    return vector ? [{ personId: person.personId, vector }] : [];
  });

  let assignedTracks = 0;
  const remaining: VectorItem[] = [];
  for (const track of unassigned) {
    let best: string | null = null;
    let bestSimilarity = FACE_SIMILARITY;
    for (const suggestion of suggestions) {
      if (suggestion.vector.length !== track.vector.length) continue;
      const similarity = cosineSimilarity(suggestion.vector, track.vector);
      if (similarity >= bestSimilarity) {
        best = suggestion.personId;
        bestSimilarity = similarity;
      }
    }
    if (best) {
      await deps.repos.assignTrack(spaceId, track.id, best);
      assignedTracks += 1;
    } else {
      remaining.push(track);
    }
  }
  const groups = clusterFaces(sameDimension(remaining), { similarity: FACE_SIMILARITY });
  const created = await deps.repos.createPersonSuggestions(
    spaceId,
    groups.map((group) => ({ centroid: encodeVectors([group.centroid]), trackIds: group.ids })),
  );
  return { assignedTracks, suggestedPeople: created.length };
}