atelier/sidecar/db/collections/queries.ts

import { and, desc, eq, inArray, sql } from 'drizzle-orm';
import { db } from '../db';
import { generations } from '../generations/schema';
import { collectionItems, collections } from './schema';

export type Collection = typeof collections.$inferSelect;
export type CollectionWithCount = Collection & { itemCount: number; coverImagePath: string | null };

/**
 * Every collection, with how many images it holds and one to show for it.
 *
 * Counted in the query rather than by loading the memberships and calling `.length`: the list screen
 * asks for this on every render, and the alternative reads every row of every collection to produce
 * a number.
 */
export async function listCollections(): Promise<CollectionWithCount[]> {
  const rows = await db
    .select({
      id: collections.id,
      name: collections.name,
      description: collections.description,
      createdAt: collections.createdAt,
      updatedAt: collections.updatedAt,
      itemCount: sql<number>`count(${collectionItems.generationId})::int`,
      // The newest image that actually has a file — a failed generation makes a poor cover.
      coverImagePath: sql<string | null>`
        (select g.image_path from ${collectionItems} ci
         join ${generations} g on g.id = ci.generation_id
         where ci.collection_id = ${collections.id} and g.image_path is not null
         order by g.created_at desc limit 1)
      `,
    })
    .from(collections)
    .leftJoin(collectionItems, eq(collectionItems.collectionId, collections.id))
    .groupBy(collections.id)
    .orderBy(desc(collections.updatedAt));

  return rows;
}

export async function getCollection(id: string): Promise<Collection | undefined> {
  const [row] = await db.select().from(collections).where(eq(collections.id, id)).limit(1);
  return row;
}

export async function getCollectionByName(name: string): Promise<Collection | undefined> {
  const [row] = await db.select().from(collections).where(eq(collections.name, name)).limit(1);
  return row;
}

export async function createCollection(name: string, description?: string | null): Promise<Collection> {
  const [row] = await db.insert(collections).values({ name, description }).returning();
  return row!;
}

/**
 * Find by name, or make it.
 *
 * `onConflictDoUpdate` rather than select-then-insert: two requests naming the same new collection
 * arrive together often — that is exactly what "put this afternoon's work in one place" produces —
 * and the check-then-act version loses that race and fails one of them on the unique constraint.
 */
export async function findOrCreateCollection(name: string): Promise<Collection> {
  const [row] = await db
    .insert(collections)
    .values({ name })
    .onConflictDoUpdate({ target: collections.name, set: { updatedAt: new Date() } })
    .returning();
  return row!;
}

export async function renameCollection(id: string, values: { name?: string; description?: string | null }) {
  const [row] = await db
    .update(collections)
    .set({ ...values, updatedAt: new Date() })
    .where(eq(collections.id, id))
    .returning();
  return row;
}

/** The memberships go with it; the images do not. */
export async function deleteCollection(id: string): Promise<void> {
  await db.delete(collections).where(eq(collections.id, id));
}

/** Adding something already in the collection is a no-op, not an error. */
export async function addToCollection(collectionId: string, generationIds: string[]): Promise<number> {
  if (generationIds.length === 0) return 0;

  const rows = await db
    .insert(collectionItems)
    .values(generationIds.map((generationId) => ({ collectionId, generationId })))
    .onConflictDoNothing()
    .returning({ generationId: collectionItems.generationId });

  await db.update(collections).set({ updatedAt: new Date() }).where(eq(collections.id, collectionId));
  return rows.length;
}

export async function removeFromCollection(collectionId: string, generationIds: string[]): Promise<void> {
  if (generationIds.length === 0) return;

  await db
    .delete(collectionItems)
    .where(and(eq(collectionItems.collectionId, collectionId), inArray(collectionItems.generationId, generationIds)));

  await db.update(collections).set({ updatedAt: new Date() }).where(eq(collections.id, collectionId));
}

/** The images in a collection, newest first. */
export async function collectionImages(collectionId: string) {
  return db
    .select({ generation: generations, addedAt: collectionItems.addedAt })
    .from(collectionItems)
    .innerJoin(generations, eq(generations.id, collectionItems.generationId))
    .where(eq(collectionItems.collectionId, collectionId))
    .orderBy(desc(generations.createdAt))
    .then((rows) => rows.map((row) => row.generation));
}

/** Which collections a generation is in — what the "add to collection" control needs to render state. */
export async function collectionsForGenerations(generationIds: string[]): Promise<Record<string, string[]>> {
  if (generationIds.length === 0) return {};

  const rows = await db
    .select({ generationId: collectionItems.generationId, collectionId: collectionItems.collectionId })
    .from(collectionItems)
    .where(inArray(collectionItems.generationId, generationIds));

  const map: Record<string, string[]> = {};
  for (const row of rows) (map[row.generationId] ??= []).push(row.collectionId);
  return map;
}