stealth-ai-relay/frontend/src/lib/convex/messages.ts

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

export const listByChat = query({
	args: { chatId: v.id('chats') },
	returns: v.array(
		v.object({
			_id: v.id('messages'),
			_creationTime: v.number(),
			chatId: v.id('chats'),
			role: v.union(v.literal('user'), v.literal('assistant')),
			content: v.string(),
			followUpOptions: v.optional(v.array(v.string())),
			source: v.union(v.literal('telegram'), v.literal('web')),
			createdAt: v.number(),
			isStreaming: v.optional(v.boolean())
		})
	),
	handler: async (ctx, args) => {
		const messages = await ctx.db
			.query('messages')
			.withIndex('by_chat_id_and_created_at', (q) => q.eq('chatId', args.chatId))
			.order('asc')
			.collect();

		return messages.map((m) => ({
			_id: m._id,
			_creationTime: m._creationTime,
			chatId: m.chatId,
			role: m.role,
			content: m.content,
			followUpOptions: m.followUpOptions,
			source: m.source,
			createdAt: m.createdAt,
			isStreaming: m.isStreaming
		}));
	}
});

export const create = mutation({
	args: {
		chatId: v.id('chats'),
		role: v.union(v.literal('user'), v.literal('assistant')),
		content: v.string(),
		source: v.union(v.literal('telegram'), v.literal('web')),
		imageBase64: v.optional(v.string()),
		imageMediaType: v.optional(v.string()),
		imagesBase64: v.optional(v.array(v.string())),
		imagesMediaTypes: v.optional(v.array(v.string())),
		photoDraftIds: v.optional(v.array(v.id('photoDrafts'))),
		followUpOptions: v.optional(v.array(v.string())),
		isStreaming: v.optional(v.boolean())
	},
	returns: v.id('messages'),
	handler: async (ctx, args) => {
		const messageId = await ctx.db.insert('messages', {
			chatId: args.chatId,
			role: args.role,
			content: args.content,
			source: args.source,
			imageBase64: args.imageBase64,
			imageMediaType: args.imageMediaType,
			imagesBase64: args.imagesBase64,
			imagesMediaTypes: args.imagesMediaTypes,
			followUpOptions: args.followUpOptions,
			createdAt: Date.now(),
			isStreaming: args.isStreaming
		});

		const drafts: Array<{ base64: string; mediaType: string; id: Id<'photoDrafts'> }> = [];
		if (args.photoDraftIds && args.photoDraftIds.length > 0) {
			for (const draftId of args.photoDraftIds) {
				const draft = await ctx.db.get(draftId);
				if (draft) {
					drafts.push({ base64: draft.base64, mediaType: draft.mediaType, id: draft._id });
				}
			}
		}

		for (let i = 0; i < drafts.length; i++) {
			await ctx.db.insert('messageImages', {
				messageId,
				base64: drafts[i].base64,
				mediaType: drafts[i].mediaType,
				order: i
			});
		}

		if (args.source === 'web' && args.role === 'user') {
			const chat = await ctx.db.get(args.chatId);
			if (chat) {
				const pendingGenId = await ctx.db.insert('pendingGenerations', {
					userId: chat.userId,
					chatId: args.chatId,
					userMessage: args.content,
					createdAt: Date.now()
				});

				for (let i = 0; i < drafts.length; i++) {
					await ctx.db.insert('pendingGenerationImages', {
						pendingGenerationId: pendingGenId,
						base64: drafts[i].base64,
						mediaType: drafts[i].mediaType,
						order: i
					});
				}
			}
		}

		for (const draft of drafts) {
			await ctx.db.delete(draft.id);
		}

		return messageId;
	}
});

export const update = mutation({
	args: {
		messageId: v.id('messages'),
		content: v.optional(v.string()),
		followUpOptions: v.optional(v.array(v.string())),
		isStreaming: v.optional(v.boolean())
	},
	returns: v.null(),
	handler: async (ctx, args) => {
		const updates: {
			content?: string;
			followUpOptions?: string[];
			isStreaming?: boolean;
		} = {};

		if (args.content !== undefined) {
			updates.content = args.content;
		}
		if (args.followUpOptions !== undefined) {
			updates.followUpOptions = args.followUpOptions;
		}
		if (args.isStreaming !== undefined) {
			updates.isStreaming = args.isStreaming;
		}

		await ctx.db.patch(args.messageId, updates);
		return null;
	}
});

export const setTelegramMessageId = mutation({
	args: { messageId: v.id('messages'), telegramMessageId: v.number() },
	returns: v.null(),
	handler: async (ctx, args) => {
		await ctx.db.patch(args.messageId, { telegramMessageId: args.telegramMessageId });
		return null;
	}
});

export const getTelegramMessageId = query({
	args: { messageId: v.id('messages') },
	returns: v.union(v.number(), v.null()),
	handler: async (ctx, args) => {
		const msg = await ctx.db.get(args.messageId);
		return msg?.telegramMessageId ?? null;
	}
});

export const getHistoryForAI = query({
	args: { chatId: v.id('chats'), limit: v.optional(v.number()) },
	returns: v.array(
		v.object({
			role: v.union(v.literal('user'), v.literal('assistant')),
			content: v.string()
		})
	),
	handler: async (ctx, args) => {
		const messages = await ctx.db
			.query('messages')
			.withIndex('by_chat_id_and_created_at', (q) => q.eq('chatId', args.chatId))
			.order('asc')
			.collect();

		const limit = args.limit ?? 50;
		const limited = messages.slice(-limit);

		return limited.map((m) => ({
			role: m.role,
			content: m.content
		}));
	}
});

export const getLastAssistantMessage = query({
	args: { chatId: v.id('chats') },
	returns: v.union(
		v.object({
			_id: v.id('messages'),
			_creationTime: v.number(),
			chatId: v.id('chats'),
			role: v.union(v.literal('user'), v.literal('assistant')),
			content: v.string(),
			imageBase64: v.optional(v.string()),
			imageMediaType: v.optional(v.string()),
			imagesBase64: v.optional(v.array(v.string())),
			imagesMediaTypes: v.optional(v.array(v.string())),
			followUpOptions: v.optional(v.array(v.string())),
			source: v.union(v.literal('telegram'), v.literal('web')),
			createdAt: v.number(),
			isStreaming: v.optional(v.boolean())
		}),
		v.null()
	),
	handler: async (ctx, args) => {
		const messages = await ctx.db
			.query('messages')
			.withIndex('by_chat_id_and_created_at', (q) => q.eq('chatId', args.chatId))
			.order('desc')
			.collect();

		return messages.find((m) => m.role === 'assistant') ?? null;
	}
});

export const getChatImages = query({
	args: { chatId: v.id('chats') },
	returns: v.array(
		v.object({
			base64: v.string(),
			mediaType: v.string()
		})
	),
	handler: async (ctx, args) => {
		const messages = await ctx.db
			.query('messages')
			.withIndex('by_chat_id', (q) => q.eq('chatId', args.chatId))
			.collect();

		const images: Array<{ base64: string; mediaType: string }> = [];

		for (const m of messages) {
			const msgImages = await ctx.db
				.query('messageImages')
				.withIndex('by_message_id', (q) => q.eq('messageId', m._id))
				.collect();

			for (const img of msgImages.sort((a, b) => a.order - b.order)) {
				images.push({ base64: img.base64, mediaType: img.mediaType });
			}

			if (m.imagesBase64 && m.imagesMediaTypes) {
				for (let i = 0; i < m.imagesBase64.length; i++) {
					images.push({
						base64: m.imagesBase64[i],
						mediaType: m.imagesMediaTypes[i]
					});
				}
			} else if (m.imageBase64 && m.imageMediaType) {
				images.push({
					base64: m.imageBase64,
					mediaType: m.imageMediaType
				});
			}
		}

		return images;
	}
});