import { v } from 'convex/values';
import { mutation, query } from './_generated/server';

export const list = query({
	args: {},
	returns: v.array(
		v.object({
			_id: v.id('pendingGenerations'),
			_creationTime: v.number(),
			userId: v.id('users'),
			chatId: v.id('chats'),
			userMessage: v.string(),
			imagesBase64: v.optional(v.array(v.string())),
			imagesMediaTypes: v.optional(v.array(v.string())),
			createdAt: v.number()
		})
	),
	handler: async (ctx) => {
		const pending = await ctx.db.query('pendingGenerations').collect();

		const result = [];
		for (const p of pending) {
			const images = await ctx.db
				.query('pendingGenerationImages')
				.withIndex('by_pending_generation_id', (q) => q.eq('pendingGenerationId', p._id))
				.collect();

			const sortedImages = images.sort((a, b) => a.order - b.order);

			result.push({
				...p,
				imagesBase64:
					sortedImages.length > 0 ? sortedImages.map((img) => img.base64) : p.imagesBase64,
				imagesMediaTypes:
					sortedImages.length > 0 ? sortedImages.map((img) => img.mediaType) : p.imagesMediaTypes
			});
		}

		return result;
	}
});

export const create = mutation({
	args: {
		userId: v.id('users'),
		chatId: v.id('chats'),
		userMessage: v.string()
	},
	returns: v.id('pendingGenerations'),
	handler: async (ctx, args) => {
		return await ctx.db.insert('pendingGenerations', {
			userId: args.userId,
			chatId: args.chatId,
			userMessage: args.userMessage,
			createdAt: Date.now()
		});
	}
});

export const remove = mutation({
	args: { id: v.id('pendingGenerations') },
	returns: v.null(),
	handler: async (ctx, args) => {
		const images = await ctx.db
			.query('pendingGenerationImages')
			.withIndex('by_pending_generation_id', (q) => q.eq('pendingGenerationId', args.id))
			.collect();
		for (const img of images) {
			await ctx.db.delete(img._id);
		}
		await ctx.db.delete(args.id);
		return null;
	}
});

export const getImages = query({
	args: { pendingGenerationId: v.id('pendingGenerations') },
	returns: v.array(
		v.object({
			base64: v.string(),
			mediaType: v.string()
		})
	),
	handler: async (ctx, args) => {
		const images = await ctx.db
			.query('pendingGenerationImages')
			.withIndex('by_pending_generation_id', (q) =>
				q.eq('pendingGenerationId', args.pendingGenerationId)
			)
			.collect();
		return images
			.sort((a, b) => a.order - b.order)
			.map((img) => ({ base64: img.base64, mediaType: img.mediaType }));
	}
});