import { OpenAIStream, StreamingTextResponse } from 'ai';
import OpenAI from 'openai';

import { auth } from '@/auth';
import { nanoid } from '@/lib/utils';
import {
	ChatCompletionMessageParam,
	ChatCompletionContentPart,
	ChatCompletionContentPartImage,
} from 'openai/resources';
import { MessageWithSelectedDataset } from '../../../lib/types';

export const runtime = 'edge';

const openai = new OpenAI({
	apiKey: process.env.OPENAI_API_KEY,
});

export async function POST(req: Request) {
	const json = await req.json();
	const { messages } = json as {
		messages: MessageWithSelectedDataset[];
	};

	const userId = (await auth())?.user.id;

	if (!userId) {
		return new Response('Unauthorized', {
			status: 401,
		});
	}

	const formattedMessage: ChatCompletionMessageParam[] = messages.map(
		message => {
			const { dataset, ...rest } = message;

			const contentWithImage: ChatCompletionContentPart[] = [
				{
					type: 'text',
					text: message.content as string,
				},
				...(dataset ?? []).map(
					entity =>
						({
							type: 'image_url',
							image_url: { url: entity.url },
						}) satisfies ChatCompletionContentPartImage,
				),
			];
			return {
				role: 'user',
				content: contentWithImage,
			};
		},
	);

	const res = await openai.chat.completions.create({
		model: 'gpt-4-vision-preview',
		messages: formattedMessage,
		temperature: 0.3,
		stream: true,
		max_tokens: 300,
	});

	const stream = OpenAIStream(res);

	return new StreamingTextResponse(stream);
}